From 9a57483e7d698f15d75c0db2158180c17e82c1aa Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Fri, 6 Nov 2026 11:41:58 -0800 Subject: [PATCH 01/10] =?UTF-8?q?feat(custom=5Frecipes):=20=E5=AE=8C?= =?UTF-8?q?=E6=88=90=20MXFP4=20=E6=94=AF=E6=8C=81=E5=B9=B6=E8=AE=A9=20grou?= =?UTF-8?q?ped=20GEMM=20=E6=94=AF=E6=8C=81=E8=87=AA=E5=AE=9A=E4=B9=89?= =?UTF-8?q?=E9=87=8F=E5=8C=96=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CustomRecipe 目前只能用于普通 Linear：`general_gemm` 有自定义分流，而 `general_grouped_gemm` / `general_grouped_gemm_for_grouped_tensor` 没有，因此 MoE 专家（走 `GroupedLinear`）永远无法到达 quantizer 的 `qgemm`。同时参考实现里 block-32 / E8M0（即 MXFP4）那条分支未写完，即使能走到也会中途报错。 本补丁按实测失败逐项修复（每项都有可复现的报错原文）： reference_nvfp4.py 1. pow_2（MXFP4）分支只做了配置校验，从未准备 `row_input` / `col_input` / `global_amax_row` / `global_amax_col`，公共尾部因此抛 `UnboundLocalError: cannot access local variable 'row_input'`。 补上赋值，并对 `_split_quantization` 可能产生的 0 元素分片加保护 （`torch.max()` 无参调用在空张量上会抛 `RuntimeError: max(): Expected reduction dim to be specified`）。 2. `size(*args, **kwargs)` 忽略参数、始终返回整个逻辑形状，导致调用方 `GroupedLinear._validate_or_alloc_output` 的 `rows = x.size(0)` 得到 `torch.Size`， `torch.empty((rows, cols))` 抛 `TypeError: empty(): argument 'size' failed to unpack the object at pos 2`。 改为 `size(dim=None)`，传入 dim 时返回该维长度。 3. `update_usage` 在两个方向都没有数据（占位张量）时直接抛 `Cannot generate FP4 data, even from FP4 data transpose`；占位符没有任何可推导的 数据，应视为 no-op 而不是错误。 4. 补 `quantize_()`。`module/base.py::quantize_weight` 在 workspace 暴露 `quantize_` 时走它，否则回落 `tex.quantize(...)`——后者无法驱动 Python 参考 量化器，workspace 会保持为空，随后 GEMM 拿到 `data is None`。 cpp_extensions/gemm.py 5. 给 `general_grouped_gemm` 增加自定义分流（列表形态：A 为权重列表、B 为激活列表、 out 为单张量、`m_splits` 为每个专家的有效激活行数），逐组调用 `qgemm`。 6. 给 `general_grouped_gemm_for_grouped_tensor` 增加自定义分流 （`GroupedTensorStorage` 形态，按 `offsets` / `scale_inv_offsets` 切组）。 两处分流都对不支持的契约显式拒绝（GELU 输出、bias 梯度、D_dtype、输出量化、 grad），而不是静默产生错误结果。 实测背景：在 B300（sm_103，EP8+CP8）上，MXFP4 blockscaled fprop GEMM 已用 CUTLASS `KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103` 实现并与参考实现 逐位一致；上述缺口是它无法接入训练的原因。 Signed-off-by: Nolan Ho --- .../pytorch/cpp_extensions/gemm.py | 250 +++++++++++++++++- .../pytorch/custom_recipes/reference_nvfp4.py | 53 +++- 2 files changed, 296 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index eb43f83cc1d..79ffb7923a0 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -34,7 +34,7 @@ from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom -from ..custom_recipes.gemm import custom_gemm +from ..custom_recipes.gemm import custom_gemm, GEMMType, MMParams from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer __all__ = [ @@ -1040,6 +1040,110 @@ def strided_batched_gemm( ) +# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). +# +# Custom-recipe dispatch for the LIST-based grouped GEMM. The operand roles are taken from +# GroupedLinear's own call site (module/grouped_linear.py:873-884): +# +# out = _validate_or_alloc_output(out, sum(m_splits), weights_fp8[0].size(0), ...) +# general_grouped_gemm(weights_fp8, inputmats, [out], ..., single_output=True, m_splits=m_splits) +# +# A = list of weights (one per expert) +# B = list of inputs (per-expert activation buffers, padded) +# out = [single tensor] of shape [sum(m_splits), A[0].size(0)] +# m_splits[i] = VALID activation rows for expert i +# +# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result is +# written to out[0][offset : offset+m_splits[i]]. An earlier version (gpt-6.1-sol's first draft) had +# A and B swapped and assumed the groups split columns; that is why it reported a 128-wide weight +# against a 4608-wide output for several rounds. + + +def _custom_grouped_gemm( + A, B, out, quantization_params, out_dtype, + layout="TN", m_splits=None, gelu=False, grad=False, + accumulate=False, bias=None, use_bias=False, + use_split_accumulator=False, D_dtype=None, single_output=False, +): + """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" + if gelu: + raise NotImplementedError("custom grouped GEMM does not define GELU outputs") + if grad and use_bias: + raise NotImplementedError("custom grouped GEMM does not define bias gradients") + if D_dtype is not None: + raise NotImplementedError("custom grouped GEMM does not define D_dtype") + if quantization_params is not None and any(q is not None for q in quantization_params): + raise NotImplementedError("custom grouped GEMM does not define output quantization") + if not A or not B: + raise ValueError("A and B must both be non-empty") + if not all(is_custom(x) for x in (*A, *B)): + raise TypeError("custom grouped GEMM cannot mix custom and ordinary tensors") + if m_splits is None: + raise ValueError("m_splits is required for the per-expert layout") + if not single_output or len(out) != 1: + raise NotImplementedError("custom grouped GEMM supports single_output only") + + n = len(m_splits) + if len(A) != n or len(B) != n: + raise ValueError(f"expected {n} weights and {n} inputs, got {len(A)} and {len(B)}") + total_rows = sum(m_splits) + + quantizer = None + for candidate in (*A, *B): + q = getattr(candidate, "_quantizer", None) + if q is not None and callable(getattr(q, "qgemm", None)): + quantizer = q + break + if quantizer is None: + raise TypeError("no quantizer with a callable qgemm on the grouped operands") + + m_params = MMParams(out_dtype=out_dtype, use_split_accumulator=use_split_accumulator) + + dst_full = out[0] + if dst_full.shape[0] < total_rows: + raise ValueError(f"output has {dst_full.shape[0]} rows but sum(m_splits)={total_rows}") + width = dst_full.shape[-1] + + offset = 0 + for i in range(n): + rows = m_splits[i] + if rows == 0: + continue + w, act = A[i], B[i] + if w.data is None or act.data is None: + raise ValueError(f"group {i}: missing quantized data (weight or activation)") + if w.scale is None or act.scale is None: + raise ValueError(f"group {i}: missing scale_inv (weight or activation)") + + qw = w.data + sw = w.scale + qx = act.data[:rows] + sx = act.scale + # Activation scales are per-row, so slice the same prefix as the data. + if hasattr(sx, "shape") and sx.ndim >= 1 and sx.shape[0] >= rows: + sx = sx[:rows] + + result = quantizer.qgemm( + qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, + ) + if tuple(result.shape) != (rows, width): + raise ValueError( + f"group {i}: qgemm returned {tuple(result.shape)}, expected {(rows, width)} " + f"(qx={tuple(qx.shape)} qw={tuple(qw.shape)} sx={tuple(sx.shape)} " + f"sw={tuple(sw.shape)} widths={width})" + ) + with torch.no_grad(): + if accumulate: + dst_full[offset:offset + rows].add_(result) + else: + dst_full[offset:offset + rows].copy_(result) + offset += rows + + bias_output = [torch.empty(0, dtype=dst_full.dtype, device=dst_full.device)] + gelu_input = [torch.empty(0, dtype=dst_full.dtype, device=dst_full.device)] + return out, bias_output, gelu_input + + def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], @@ -1060,6 +1164,12 @@ def general_grouped_gemm( """ TN layout Grouped GEMM with fp8 inputs. """ + if any(is_custom(a) for a in A) or any(is_custom(b) for b in B): + return _custom_grouped_gemm( + A, B, out, quantization_params, out_dtype, layout, m_splits, gelu, grad, + accumulate, bias, use_bias, use_split_accumulator, D_dtype, single_output, + ) + num_gemms = len(A) transa = layout[0] == "T" @@ -1205,6 +1315,137 @@ def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) +# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). +# Custom-recipe dispatch for the *grouped-tensor* GEMM — the entry point MoE experts actually use +# (GroupedLinear fprop -> general_grouped_gemm_for_grouped_tensor -> nvte_grouped_gemm). +# +# Authorship: designed by gpt-6.1-sol against the contracts recorded in +# engine/scripts/wip/ds41-b300/worklog.md; the grouped-offset slicing here is the parent's, written +# from the field names read out of tensor/storage/grouped_tensor_storage.py +# (`offsets`, `scale_inv_offsets`, `rowwise_data`, `rowwise_scale_inv`, `num_tensors`). + + +def _custom_grouped_tensor_gemm( + A, + B, + out, + *, + layout="TN", + accumulate=False, + use_split_accumulator=False, + bias=None, + bias_scale=None, + grad=False, + alpha=None, + beta=None, +): + """Grouped GEMM over GroupedTensorStorage operands using a custom quantizer's qgemm.""" + import math + + if grad: + raise NotImplementedError("custom grouped-tensor GEMM does not support grad") + if bias is not None or bias_scale is not None: + raise NotImplementedError("custom grouped-tensor GEMM does not support bias") + if alpha is not None or beta is not None: + raise NotImplementedError("custom grouped-tensor GEMM does not support alpha/beta") + if layout != "TN": + raise NotImplementedError(f"custom grouped-tensor GEMM supports layout TN only, got {layout!r}") + + def expect_grouped(x, name): + if not hasattr(x, "rowwise_data"): + raise TypeError(f"{name} must be a GroupedTensorStorage, got {type(x).__name__}") + if getattr(x, "rowwise_data", None) is None: + raise ValueError(f"{name}.rowwise_data is None (columnwise-only storage is unsupported)") + return x + + A = expect_grouped(A, "A") + B = expect_grouped(B, "B") + out = expect_grouped(out, "out") + + num = getattr(A, "num_tensors", None) + if not num: + raise ValueError("A.num_tensors is unset; cannot iterate groups") + + def offsets_of(x, name): + offsets = getattr(x, "offsets", None) + if offsets is None: + raise ValueError(f"{name}.offsets is None; per-group slicing is impossible") + offsets = list(offsets) + if len(offsets) != num + 1: + raise ValueError(f"{name}.offsets has {len(offsets)} entries, expected {num + 1}") + return offsets + + a_off = offsets_of(A, "A") + b_off = offsets_of(B, "B") + o_off = offsets_of(out, "out") + + def scale_offsets_of(x, name): + so = getattr(x, "scale_inv_offsets", None) + if so is None: + return None + so = list(so) + if len(so) != num + 1: + raise ValueError(f"{name}.scale_inv_offsets has {len(so)} entries, expected {num + 1}") + return so + + a_s_off = scale_offsets_of(A, "A") + b_s_off = scale_offsets_of(B, "B") + + quantizer = getattr(A, "quantizer", None) + if quantizer is None: + quantizer = getattr(B, "quantizer", None) + if quantizer is None or not callable(getattr(quantizer, "qgemm", None)): + raise TypeError("no quantizer with a callable qgemm on the grouped operands") + + m_params = MMParams(out_dtype=None, use_split_accumulator=use_split_accumulator) + out_dtype = out.rowwise_data.dtype + + # scipy-free zip over adjacent offsets + for i in range(num): + a0, a1 = a_off[i], a_off[i + 1] + rows = a1 - a0 + if rows == 0: + continue + + b0, b1 = b_off[i], b_off[i + 1] + o0, o1 = o_off[i], o_off[i + 1] + if o1 - o0 != rows: + raise ValueError( + f"group {i}: out segment has {o1 - o0} rows but A has {rows} " + f"(a_off={a_off[:3]}..., o_off={o_off[:3]}...)" + ) + + qx = A.rowwise_data[a0:a1] + qw = B.rowwise_data[b0:b1] + + sx = None + if getattr(A, "rowwise_scale_inv", None) is not None and a_s_off is not None: + sx = A.rowwise_scale_inv[a_s_off[i]:a_s_off[i + 1]] + sw = None + if getattr(B, "rowwise_scale_inv", None) is not None and b_s_off is not None: + sw = B.rowwise_scale_inv[b_s_off[i]:b_s_off[i + 1]] + if sx is None or sw is None: + raise ValueError(f"group {i}: missing rowwise scale_inv for a custom quantizer") + + result = quantizer.qgemm( + qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, + ) + target = out.rowwise_data[o0:o1] + if tuple(result.shape) != tuple(target.shape): + raise ValueError( + f"group {i}: qgemm returned {tuple(result.shape)} but the out segment is " + f"{tuple(target.shape)} (qx={tuple(qx.shape)} qw={tuple(qw.shape)} " + f"sx={tuple(sx.shape)} sw={tuple(sw.shape)})" + ) + with torch.no_grad(): + if accumulate: + target.add_(result) + else: + target.copy_(result) + + return out + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -1227,6 +1468,13 @@ def general_grouped_gemm_for_grouped_tensor( The caller must ensure that GroupedTensor metadata is already compatible with the underlying GEMM implementation (e.g., aligned offsets and output metadata layout). """ + if is_custom(A) or is_custom(B) or is_custom(out): + return _custom_grouped_tensor_gemm( + A, B, out, layout=layout, accumulate=accumulate, + use_split_accumulator=use_split_accumulator, bias=bias, + bias_scale=bias_scale, grad=grad, alpha=alpha, beta=beta, + ) + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." if grad: raise NotImplementedError("grad is not supported for grouped_tensor GEMM yet.") diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index 29e525da45d..db07931b576 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -392,6 +392,23 @@ def __repr__(self): ")" ) + def quantize_(self, tensor, noop_flag=None): + """Quantise into this tensor in place. + + `module/base.py::quantize_weight` calls `workspace.quantize_(...)` when the + workspace exposes it and otherwise falls back to `tex.quantize(...)`, a C++ + entry that cannot drive a Python reference quantizer — the workspace then stays + empty and the GEMM fails on `data is None`. + """ + del noop_flag + fresh = self._quantizer.quantize(tensor) + self.data = fresh.data + self.scale = fresh.scale + self.data_t = fresh.data_t + self.scale_t = fresh.scale_t + self.original_shape = tuple(tensor.shape) + return self + def update_usage( self, rowwise_usage: Optional[bool] = None, @@ -408,6 +425,11 @@ def update_usage( if columnwise_usage is not None: needs_data_transpose = columnwise_usage + # A placeholder (neither direction populated) has nothing to generate from; + # refusing here turns workspace allocation into a hard error. + if not has_data and not has_data_transpose: + return self + # Generate data that is required if needs_data and not has_data: raise RuntimeError("Cannot generate FP4 data, even from FP4 data transpose") @@ -429,16 +451,20 @@ def _create_transpose(self): self.data_t = self.data.t().contiguous() self.scale_t = self.scale - def size(self, *args, **kwargs): # pylint: disable=unused-argument - """Return the original tensor shape, not the internal packed data shape. + def size(self, dim=None): + """Return the logical original shape, honouring `dim`. - FP4 quantization packs two 4-bit values into each 8-bit value, which reduces - the second dimension by half. This method returns the logical shape that - users expect, not the internal packed storage shape. + FP4 packs two 4-bit values into each 8-bit value, halving the second dimension, so + the logical shape is not the storage shape. The no-argument form returns the full + logical shape. Passing `dim` must return that dimension's size: callers such as + `GroupedLinear._validate_or_alloc_output` compute `rows = x.size(0)` and hand it to + `torch.empty`, which raises `type must be tuple of ints, but got torch.Size` when + the whole shape comes back instead. """ if self.original_shape is None: raise RuntimeError("NVFP4TensorRef.size() called but original_shape has not been set") - return torch.Size(self.original_shape) + shape = torch.Size(self.original_shape) + return shape if dim is None else shape[dim] def get_wgrad_sign_vector() -> torch.Tensor: @@ -975,6 +1001,21 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ ) if self.row_scaled_nvfp4: raise ValueError("Row-scaled NVFP4 is only supported for NVFP4 (non-pow2) mode.") + # MXFP4 path: the branch above validated the configuration but never + # prepared the operands, so the common tail raised UnboundLocalError at + # `x_input = row_input`. Prepare them here, and guard the empty splits + # that _split_quantization can produce (torch.max() with no dim raises on + # a 0-element tensor). + row_input = tensor + col_input = tensor.t().contiguous() + if row_input.numel(): + global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) + else: + global_amax_row = torch.zeros(1, device=row_input.device, dtype=torch.float32) + if col_input.numel(): + global_amax_col = torch.max(torch.abs(col_input)).to(torch.float32).view(1) + else: + global_amax_col = torch.zeros(1, device=col_input.device, dtype=torch.float32) # TODO(etsykunov): Fix bug where global_amax_row and # global_amax_col are not defined # global_amax = torch.empty(0, device=tensor.device, dtype=torch.float32) From 02e6614d84e6e9724f5a59f8d03c052d3458869d Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Fri, 6 Nov 2026 11:41:58 -0800 Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Nolan Ho --- .../pytorch/cpp_extensions/gemm.py | 84 +++++++++++++++---- 1 file changed, 67 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 79ffb7923a0..06c2624381f 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1060,10 +1060,21 @@ def strided_batched_gemm( def _custom_grouped_gemm( - A, B, out, quantization_params, out_dtype, - layout="TN", m_splits=None, gelu=False, grad=False, - accumulate=False, bias=None, use_bias=False, - use_split_accumulator=False, D_dtype=None, single_output=False, + A, + B, + out, + quantization_params, + out_dtype, + layout="TN", + m_splits=None, + gelu=False, + grad=False, + accumulate=False, + bias=None, + use_bias=False, + use_split_accumulator=False, + D_dtype=None, + single_output=False, ): """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" if gelu: @@ -1124,7 +1135,14 @@ def _custom_grouped_gemm( sx = sx[:rows] result = quantizer.qgemm( - qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, + qx, + qw, + m_params, + out_dtype, + sx, + sw, + None, + gemm_type=GEMMType.FPROP, ) if tuple(result.shape) != (rows, width): raise ValueError( @@ -1134,9 +1152,9 @@ def _custom_grouped_gemm( ) with torch.no_grad(): if accumulate: - dst_full[offset:offset + rows].add_(result) + dst_full[offset : offset + rows].add_(result) else: - dst_full[offset:offset + rows].copy_(result) + dst_full[offset : offset + rows].copy_(result) offset += rows bias_output = [torch.empty(0, dtype=dst_full.dtype, device=dst_full.device)] @@ -1166,8 +1184,21 @@ def general_grouped_gemm( """ if any(is_custom(a) for a in A) or any(is_custom(b) for b in B): return _custom_grouped_gemm( - A, B, out, quantization_params, out_dtype, layout, m_splits, gelu, grad, - accumulate, bias, use_bias, use_split_accumulator, D_dtype, single_output, + A, + B, + out, + quantization_params, + out_dtype, + layout, + m_splits, + gelu, + grad, + accumulate, + bias, + use_bias, + use_split_accumulator, + D_dtype, + single_output, ) num_gemms = len(A) @@ -1349,13 +1380,17 @@ def _custom_grouped_tensor_gemm( if alpha is not None or beta is not None: raise NotImplementedError("custom grouped-tensor GEMM does not support alpha/beta") if layout != "TN": - raise NotImplementedError(f"custom grouped-tensor GEMM supports layout TN only, got {layout!r}") + raise NotImplementedError( + f"custom grouped-tensor GEMM supports layout TN only, got {layout!r}" + ) def expect_grouped(x, name): if not hasattr(x, "rowwise_data"): raise TypeError(f"{name} must be a GroupedTensorStorage, got {type(x).__name__}") if getattr(x, "rowwise_data", None) is None: - raise ValueError(f"{name}.rowwise_data is None (columnwise-only storage is unsupported)") + raise ValueError( + f"{name}.rowwise_data is None (columnwise-only storage is unsupported)" + ) return x A = expect_grouped(A, "A") @@ -1420,15 +1455,22 @@ def scale_offsets_of(x, name): sx = None if getattr(A, "rowwise_scale_inv", None) is not None and a_s_off is not None: - sx = A.rowwise_scale_inv[a_s_off[i]:a_s_off[i + 1]] + sx = A.rowwise_scale_inv[a_s_off[i] : a_s_off[i + 1]] sw = None if getattr(B, "rowwise_scale_inv", None) is not None and b_s_off is not None: - sw = B.rowwise_scale_inv[b_s_off[i]:b_s_off[i + 1]] + sw = B.rowwise_scale_inv[b_s_off[i] : b_s_off[i + 1]] if sx is None or sw is None: raise ValueError(f"group {i}: missing rowwise scale_inv for a custom quantizer") result = quantizer.qgemm( - qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, + qx, + qw, + m_params, + out_dtype, + sx, + sw, + None, + gemm_type=GEMMType.FPROP, ) target = out.rowwise_data[o0:o1] if tuple(result.shape) != tuple(target.shape): @@ -1470,9 +1512,17 @@ def general_grouped_gemm_for_grouped_tensor( """ if is_custom(A) or is_custom(B) or is_custom(out): return _custom_grouped_tensor_gemm( - A, B, out, layout=layout, accumulate=accumulate, - use_split_accumulator=use_split_accumulator, bias=bias, - bias_scale=bias_scale, grad=grad, alpha=alpha, beta=beta, + A, + B, + out, + layout=layout, + accumulate=accumulate, + use_split_accumulator=use_split_accumulator, + bias=bias, + bias_scale=bias_scale, + grad=grad, + alpha=alpha, + beta=beta, ) assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." From e2d3e412a30a4d5e72e913230f25e79b120ff24d Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Fri, 6 Nov 2026 11:41:58 -0800 Subject: [PATCH 03/10] fix(custom_recipes): support the per-group output layout in the grouped custom GEMM TEGroupedLinear passes one output tensor per expert rather than a single concatenated buffer, so refusing single_output=False made the MoE expert path unusable even with the dispatch in place. Both layouts are now handled: with a single buffer each expert writes rows [offset : offset+m_splits[i]], and with a list each expert writes its own tensor. Verified on a TEGroupedLinear with four experts: the dispatch is reached and every group is dispatched to qgemm. Signed-off-by: Nolan Ho --- .../pytorch/cpp_extensions/gemm.py | 85 +++++++++++-------- 1 file changed, 51 insertions(+), 34 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 06c2624381f..da5cc3a1a6f 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1059,22 +1059,30 @@ def strided_batched_gemm( # against a 4608-wide output for several rounds. +# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). +# +# Custom-recipe dispatch for the LIST-based grouped GEMM. The operand roles are taken from +# GroupedLinear's own call site (module/grouped_linear.py:873-884): +# +# out = _validate_or_alloc_output(out, sum(m_splits), weights_fp8[0].size(0), ...) +# general_grouped_gemm(weights_fp8, inputmats, [out], ..., single_output=True, m_splits=m_splits) +# +# A = list of weights (one per expert) +# B = list of inputs (per-expert activation buffers, padded) +# out = [single tensor] of shape [sum(m_splits), A[0].size(0)] +# m_splits[i] = VALID activation rows for expert i +# +# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result is +# written to out[0][offset : offset+m_splits[i]]. An earlier version (gpt-6.1-sol's first draft) had +# A and B swapped and assumed the groups split columns; that is why it reported a 128-wide weight +# against a 4608-wide output for several rounds. + + def _custom_grouped_gemm( - A, - B, - out, - quantization_params, - out_dtype, - layout="TN", - m_splits=None, - gelu=False, - grad=False, - accumulate=False, - bias=None, - use_bias=False, - use_split_accumulator=False, - D_dtype=None, - single_output=False, + A, B, out, quantization_params, out_dtype, + layout="TN", m_splits=None, gelu=False, grad=False, + accumulate=False, bias=None, use_bias=False, + use_split_accumulator=False, D_dtype=None, single_output=False, ): """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" if gelu: @@ -1091,8 +1099,13 @@ def _custom_grouped_gemm( raise TypeError("custom grouped GEMM cannot mix custom and ordinary tensors") if m_splits is None: raise ValueError("m_splits is required for the per-expert layout") - if not single_output or len(out) != 1: - raise NotImplementedError("custom grouped GEMM supports single_output only") + # Two output layouts occur: a single concatenated buffer (`single_output=True`) or one tensor + # per group. Both are supported; the earlier version refused the list form, which is the one + # TEGroupedLinear actually uses. + if single_output and len(out) != 1: + raise ValueError("single_output=True requires exactly one output tensor") + if not single_output and len(out) != len(m_splits): + raise ValueError(f"expected {len(m_splits)} outputs, got {len(out)}") n = len(m_splits) if len(A) != n or len(B) != n: @@ -1110,10 +1123,16 @@ def _custom_grouped_gemm( m_params = MMParams(out_dtype=out_dtype, use_split_accumulator=use_split_accumulator) - dst_full = out[0] - if dst_full.shape[0] < total_rows: - raise ValueError(f"output has {dst_full.shape[0]} rows but sum(m_splits)={total_rows}") - width = dst_full.shape[-1] + dst_full = out[0] if single_output else None + if single_output: + if dst_full.shape[0] < total_rows: + raise ValueError(f"output has {dst_full.shape[0]} rows but sum(m_splits)={total_rows}") + width = dst_full.shape[-1] + else: + width = out[0].shape[-1] + for i, o in enumerate(out): + if o.shape[-1] != width: + raise ValueError(f"out[{i}] width {o.shape[-1]} != {width}") offset = 0 for i in range(n): @@ -1135,14 +1154,7 @@ def _custom_grouped_gemm( sx = sx[:rows] result = quantizer.qgemm( - qx, - qw, - m_params, - out_dtype, - sx, - sw, - None, - gemm_type=GEMMType.FPROP, + qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, ) if tuple(result.shape) != (rows, width): raise ValueError( @@ -1150,18 +1162,23 @@ def _custom_grouped_gemm( f"(qx={tuple(qx.shape)} qw={tuple(qw.shape)} sx={tuple(sx.shape)} " f"sw={tuple(sw.shape)} widths={width})" ) + dst = dst_full[offset:offset + rows] if single_output else out[i] + if tuple(dst.shape) != (rows, width): + raise ValueError( + f"group {i}: output segment is {tuple(dst.shape)}, expected {(rows, width)}" + ) with torch.no_grad(): if accumulate: - dst_full[offset : offset + rows].add_(result) + dst.add_(result) else: - dst_full[offset : offset + rows].copy_(result) + dst.copy_(result) offset += rows - bias_output = [torch.empty(0, dtype=dst_full.dtype, device=dst_full.device)] - gelu_input = [torch.empty(0, dtype=dst_full.dtype, device=dst_full.device)] + ref = dst_full if single_output else out[0] + bias_output = [torch.empty(0, dtype=ref.dtype, device=ref.device)] + gelu_input = [torch.empty(0, dtype=ref.dtype, device=ref.device)] return out, bias_output, gelu_input - def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], From 9f4a963c3ec3409ce4216df036040495941c5db8 Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Fri, 6 Nov 2026 11:41:58 -0800 Subject: [PATCH 04/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Nolan Ho --- .../pytorch/cpp_extensions/gemm.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index da5cc3a1a6f..99caa6730ad 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1079,10 +1079,21 @@ def strided_batched_gemm( def _custom_grouped_gemm( - A, B, out, quantization_params, out_dtype, - layout="TN", m_splits=None, gelu=False, grad=False, - accumulate=False, bias=None, use_bias=False, - use_split_accumulator=False, D_dtype=None, single_output=False, + A, + B, + out, + quantization_params, + out_dtype, + layout="TN", + m_splits=None, + gelu=False, + grad=False, + accumulate=False, + bias=None, + use_bias=False, + use_split_accumulator=False, + D_dtype=None, + single_output=False, ): """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" if gelu: @@ -1154,7 +1165,14 @@ def _custom_grouped_gemm( sx = sx[:rows] result = quantizer.qgemm( - qx, qw, m_params, out_dtype, sx, sw, None, gemm_type=GEMMType.FPROP, + qx, + qw, + m_params, + out_dtype, + sx, + sw, + None, + gemm_type=GEMMType.FPROP, ) if tuple(result.shape) != (rows, width): raise ValueError( @@ -1162,7 +1180,7 @@ def _custom_grouped_gemm( f"(qx={tuple(qx.shape)} qw={tuple(qw.shape)} sx={tuple(sx.shape)} " f"sw={tuple(sw.shape)} widths={width})" ) - dst = dst_full[offset:offset + rows] if single_output else out[i] + dst = dst_full[offset : offset + rows] if single_output else out[i] if tuple(dst.shape) != (rows, width): raise ValueError( f"group {i}: output segment is {tuple(dst.shape)}, expected {(rows, width)}" @@ -1179,6 +1197,7 @@ def _custom_grouped_gemm( gelu_input = [torch.empty(0, dtype=ref.dtype, device=ref.device)] return out, bias_output, gelu_input + def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], From 0d175ef1145dcd125729479d84ce186c60b92313 Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Fri, 6 Nov 2026 11:41:58 -0800 Subject: [PATCH 05/10] fix(module): do not cache quantised weights for custom quantizers A custom (reference) quantizer cannot be cached safely. The cached entry is a TensorSpec whose storage is driven by update_usage, and a gradient-checkpoint replay can reach the GEMM after that storage has been released, which surfaces as FPROP GEMM: quantized weight data (B.data) is None in a model whose single un-replayed forward is fine. Measured on DSV4.1-Flash (40 layers, EP8/CP8) with an MXFP4 CustomRecipe: identical weights and identical forward, and the failure appears only on the replay path (checkpoint.py recompute_fn -> transformer_layer._run_mlp), where the operand handed to the GEMM is a payload-less reference tensor of the right shape while the module's own parameter is a healthy bf16 Parameter. Skipping the workspace path for custom quantizers re-quantises per call instead. With this in place the same run completes a full training step (loss 3.667, grad_norm 1.823) with gradient checkpointing left ON, which is what keeps activation memory within budget. Signed-off-by: Nolan Ho --- transformer_engine/pytorch/module/base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index bf4f9bba1d5..0cfa45c2af1 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -963,6 +963,15 @@ def quantize_weight( return tensor, None # Validate workspace + # A custom (reference) quantizer cannot be cached safely. The cached entry is a + # TensorSpec whose storage is driven by `update_usage`, and a gradient-checkpoint + # replay can reach the GEMM after that storage has been released, which surfaces as + # "FPROP GEMM: quantized weight data (B.data) is None" in a model that trains fine + # for a single un-replayed forward. Re-quantising per call is cheap next to a wrong + # result, so the workspace path is skipped for custom quantizers. + if quantizer is not None and getattr(quantizer, "custom", False): + workspace = None + if workspace is not None and quantizer is not None: if not _is_weight_workspace_valid(workspace, quantizer): workspace = None From 3aead3452110bc5d08d916d8063084333d8b7847 Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Sat, 7 Nov 2026 04:20:02 -0800 Subject: [PATCH 06/10] docs(custom_recipes): record the per-group launch cost of the grouped custom GEMM The grouped entry point issues one single-group GEMM per expert group. On a real training step that is ~40,000 kernel launches for ~1,000 GEMMs of work. The per-launch cost is small enough that it never surfaces in a per-phase profile, but the sum dominates the custom-recipe path's remaining overhead versus the native FP4 path. Recorded at the loop rather than fixed: collapsing it requires a grouped GEMM kernel, which is a separate change with its own numerical acceptance test. Signed-off-by: Nolan Ho --- transformer_engine/pytorch/cpp_extensions/gemm.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 99caa6730ad..7ec3514e259 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1145,6 +1145,12 @@ def _custom_grouped_gemm( if o.shape[-1] != width: raise ValueError(f"out[{i}] width {o.shape[-1]} != {width}") + # NOTE (known cost, not a correctness gap): this issues one single-group GEMM per expert group, so a + # step pays ~40x more kernel launches than it has GEMMs of work (~40,000 launches for ~1,000 GEMMs, + # measured on a real training step 2026-11). Per-launch cost is small (~0.05 ms of device time), which + # is why it never shows up in a per-phase average, but the sum is now the largest identified item in + # the custom-recipe path. Replacing this loop with a grouped GEMM kernel (one launch for all groups) + # is the intended fix and is what the surviving difference to the native FP4 path is made of. offset = 0 for i in range(n): rows = m_splits[i] From 2e556c98f3a4c164315e3cda3d9a6ebadc52642f Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Sun, 8 Nov 2026 10:32:34 -0800 Subject: [PATCH 07/10] fix(custom_recipes): route only custom grouped operands and forward qgemm contracts The grouped dispatch called is_custom on the GroupedTensorStorage containers that general_grouped_gemm_for_grouped_tensor receives, and is_custom raises on any type that is neither a Quantizer/QuantizedTensorStorage nor a Tensor. Every grouped-tensor GEMM, native ones included, failed before reaching the native implementation. The dispatch now inspects the quantizer (or the discrete list members), so an ordinary grouped tensor keeps taking the native path. The per-expert loop also handed qgemm only the packed data and scales: it dropped the per-group bias that use_bias=True promises, omitted qresult_x/qresult_w (so a non-pow-2 NVFP4 recipe, which derives its global scale from them, could not run), and ran the fprop computation for grad=True while every shape check still passed. The bias now goes to qgemm the way custom_gemm passes it, the quantized operands travel with the packed data, and dgrad/wgrad are refused explicitly instead of returning a wrong gradient. NVFP4TensorRef.quantize_ copied only the packed data and block scales, leaving global_amax_row/global_amax_col from the previous quantization, and ignored noop_flag. It now delegates to the quantizer's update_quantized, which is what QuantizedTensorStorage.quantize_ does: every field a later GEMM reads is refreshed, and the replay skip flag is honoured. Signed-off-by: Nolan Ho --- .../pytorch/cpp_extensions/gemm.py | 56 +++++++++++-------- .../pytorch/custom_recipes/reference_nvfp4.py | 15 ++--- 2 files changed, 41 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 9409f0a62c9..3083ee96db8 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1059,25 +1059,6 @@ def strided_batched_gemm( # against a 4608-wide output for several rounds. -# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). -# -# Custom-recipe dispatch for the LIST-based grouped GEMM. The operand roles are taken from -# GroupedLinear's own call site (module/grouped_linear.py:873-884): -# -# out = _validate_or_alloc_output(out, sum(m_splits), weights_fp8[0].size(0), ...) -# general_grouped_gemm(weights_fp8, inputmats, [out], ..., single_output=True, m_splits=m_splits) -# -# A = list of weights (one per expert) -# B = list of inputs (per-expert activation buffers, padded) -# out = [single tensor] of shape [sum(m_splits), A[0].size(0)] -# m_splits[i] = VALID activation rows for expert i -# -# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result is -# written to out[0][offset : offset+m_splits[i]]. An earlier version (gpt-6.1-sol's first draft) had -# A and B swapped and assumed the groups split columns; that is why it reported a 128-wide weight -# against a 4608-wide output for several rounds. - - def _custom_grouped_gemm( A, B, @@ -1098,8 +1079,14 @@ def _custom_grouped_gemm( """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" if gelu: raise NotImplementedError("custom grouped GEMM does not define GELU outputs") - if grad and use_bias: - raise NotImplementedError("custom grouped GEMM does not define bias gradients") + # fprop only. dgrad/wgrad ask for the transposed weight direction while every shape check + # still passes, so running the fprop computation for them is a silently wrong gradient. + if grad: + raise NotImplementedError("custom grouped GEMM implements fprop only") + if layout != "TN": + raise NotImplementedError( + f"custom grouped GEMM supports the TN layout only, got {layout!r}" + ) if D_dtype is not None: raise NotImplementedError("custom grouped GEMM does not define D_dtype") if quantization_params is not None and any(q is not None for q in quantization_params): @@ -1145,6 +1132,9 @@ def _custom_grouped_gemm( if o.shape[-1] != width: raise ValueError(f"out[{i}] width {o.shape[-1]} != {width}") + if use_bias and bias is None: + raise ValueError("use_bias=True requires a per-group bias list") + # NOTE (known cost, not a correctness gap): this issues one single-group GEMM per expert group, so a # step pays ~40x more kernel launches than it has GEMMs of work (~40,000 launches for ~1,000 GEMMs, # measured on a real training step 2026-11). Per-launch cost is small (~0.05 ms of device time), which @@ -1170,6 +1160,8 @@ def _custom_grouped_gemm( if hasattr(sx, "shape") and sx.ndim >= 1 and sx.shape[0] >= rows: sx = sx[:rows] + # The quantized tensors themselves carry the global scales that a non-pow-2 (NVFP4) + # recipe turns into the GEMM alpha, so they go along with the packed data. result = quantizer.qgemm( qx, qw, @@ -1177,8 +1169,10 @@ def _custom_grouped_gemm( out_dtype, sx, sw, - None, + bias[i] if use_bias else None, gemm_type=GEMMType.FPROP, + qresult_x=act, + qresult_w=w, ) if tuple(result.shape) != (rows, width): raise ValueError( @@ -1530,6 +1524,22 @@ def scale_offsets_of(x, name): return out +def _grouped_operand_is_custom(x) -> bool: + """Whether a grouped-GEMM operand is routed to a custom quantizer. + + Operands of `general_grouped_gemm_for_grouped_tensor` are `GroupedTensorStorage` + containers or discrete lists of tensors, and `is_custom` raises on any type that is + neither a `Quantizer`/`QuantizedTensorStorage` nor a `torch.Tensor`. Handing it a + container therefore turned every grouped-tensor GEMM, native ones included, into an + `AssertionError`; the quantizer (or the list members) have to be inspected instead. + """ + if isinstance(x, (list, tuple)): + return any(_grouped_operand_is_custom(item) for item in x) + if isinstance(x, (Quantizer, QuantizedTensorStorage)): + return is_custom(x) + return is_custom(getattr(x, "quantizer", None)) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -1552,7 +1562,7 @@ def general_grouped_gemm_for_grouped_tensor( The caller must ensure that GroupedTensor metadata is already compatible with the underlying GEMM implementation (e.g., aligned offsets and output metadata layout). """ - if is_custom(A) or is_custom(B) or is_custom(out): + if any(_grouped_operand_is_custom(x) for x in (A, B, out)): return _custom_grouped_tensor_gemm( A, B, diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index db07931b576..fdfa35e590f 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -399,14 +399,15 @@ def quantize_(self, tensor, noop_flag=None): workspace exposes it and otherwise falls back to `tex.quantize(...)`, a C++ entry that cannot drive a Python reference quantizer — the workspace then stays empty and the GEMM fails on `data is None`. + + Delegated to the quantizer's `update_quantized`, exactly like the shipped + `QuantizedTensorStorage.quantize_`: that path refreshes every field a later GEMM + reads — including `global_amax_row`/`global_amax_col`, which a non-pow-2 (NVFP4) + recipe turns into the global scale, so copying only the packed data and block + scales would combine fresh values with a stale scale — and honours `noop_flag`, + the flag that lets a replayed graph skip re-quantisation on later microbatches. """ - del noop_flag - fresh = self._quantizer.quantize(tensor) - self.data = fresh.data - self.scale = fresh.scale - self.data_t = fresh.data_t - self.scale_t = fresh.scale_t - self.original_shape = tuple(tensor.shape) + self._get_quantizer().update_quantized(tensor, self, noop_flag=noop_flag) return self def update_usage( From 3b88b4f6640e9d9142ca9919703b5bb809e1f0f3 Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Sun, 8 Nov 2026 10:53:21 -0800 Subject: [PATCH 08/10] fix(custom_recipes): give grouped custom recipes the contracts they advertise The custom grouped-tensor GEMM read its operands backwards. `A` is the per-expert weight and `B` the per-expert activation at every `layout="TN"` call site (module/grouped_linear.py, ops/basic/grouped_linear.py, tests), but the helper fed `A` in as the activation and used its row count for the output segment, so a custom recipe would either trip the segment check or compute the transposed product. The activation slice, its scales and the row count now come from `B`. That layout has no per-group QuantizedTensorStorage to give a quantizer as `qresult_x`/`qresult_w`, so a non-pow-2 recipe is refused inside `qgemm`; the comment above the helper now says so instead of leaving a reader to infer it. The list grouped GEMM now length-checks the per-group bias before the loop that mutates the caller's buffer; a short list used to IndexError after some groups were already written. NVFP4QuantizerRef.qgemm refused any bias, so a biased custom recipe could not run at all; it adds the bias in the float32 epilogue now, which is where the native path puts it, so grouping no longer decides whether bias works. NVFP4TensorRef.quantize_ delegates to `update_quantized`; that path flattens only ndim > 2 and skips the dtype guard `quantize` applied, so the reference method restores both. Its docstring described a call site that the weight-cache change removed, so it now states the contract instead of narrating a caller. A test pins the classification the previous commit fixed: a plain GroupedTensorStorage must not be routed through `is_custom`. Signed-off-by: Nolan Ho --- tests/pytorch/test_grouped_linear.py | 24 ++++++++ .../pytorch/cpp_extensions/gemm.py | 55 ++++++++++--------- .../pytorch/custom_recipes/reference_nvfp4.py | 38 ++++++++----- 3 files changed, 75 insertions(+), 42 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 19b27798c5d..95c2d1afd72 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -3803,6 +3803,30 @@ def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( expected_swizzled_scales_buffer, ) + +def test_grouped_tensor_custom_dispatch_classifies_plain_operands(): + """A plain GroupedTensorStorage must not be classified as a custom operand. + + `is_custom` raises on any type that is neither a Quantizer/QuantizedTensorStorage nor a + Tensor, so dispatching on the container itself made every grouped-tensor GEMM raise + before it reached the native implementation. + """ + from transformer_engine.pytorch.cpp_extensions.gemm import _grouped_operand_is_custom + from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import ( + GroupedTensorStorage, + ) + + grouped = GroupedTensorStorage( + shape=(32, 32), + dtype=torch.bfloat16, + num_tensors=2, + quantizer=None, + data=torch.zeros(32, 32, dtype=torch.bfloat16, device="cuda"), + ) + assert not _grouped_operand_is_custom(grouped) + assert not _grouped_operand_is_custom([grouped]) + assert not _grouped_operand_is_custom(None) + # Poison the padded compact scales if swizzle_type == "mxfp8_rowwise": unpadded_scale_shape = (shape[0], shape[1] // 32) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 3083ee96db8..93273ad3c66 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1040,10 +1040,8 @@ def strided_batched_gemm( ) -# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). -# # Custom-recipe dispatch for the LIST-based grouped GEMM. The operand roles are taken from -# GroupedLinear's own call site (module/grouped_linear.py:873-884): +# GroupedLinear's fprop call site: # # out = _validate_or_alloc_output(out, sum(m_splits), weights_fp8[0].size(0), ...) # general_grouped_gemm(weights_fp8, inputmats, [out], ..., single_output=True, m_splits=m_splits) @@ -1053,10 +1051,8 @@ def strided_batched_gemm( # out = [single tensor] of shape [sum(m_splits), A[0].size(0)] # m_splits[i] = VALID activation rows for expert i # -# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result is -# written to out[0][offset : offset+m_splits[i]]. An earlier version (gpt-6.1-sol's first draft) had -# A and B swapped and assumed the groups split columns; that is why it reported a 128-wide weight -# against a 4608-wide output for several rounds. +# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result +# is written to out[0][offset : offset+m_splits[i]]. def _custom_grouped_gemm( @@ -1079,8 +1075,8 @@ def _custom_grouped_gemm( """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" if gelu: raise NotImplementedError("custom grouped GEMM does not define GELU outputs") - # fprop only. dgrad/wgrad ask for the transposed weight direction while every shape check - # still passes, so running the fprop computation for them is a silently wrong gradient. + # dgrad/wgrad use different operand roles and a transposed weight; only TN fprop is + # implemented here, so they are refused instead of computing the wrong product. if grad: raise NotImplementedError("custom grouped GEMM implements fprop only") if layout != "TN": @@ -1134,6 +1130,8 @@ def _custom_grouped_gemm( if use_bias and bias is None: raise ValueError("use_bias=True requires a per-group bias list") + if use_bias and len(bias) != n: + raise ValueError(f"expected {n} biases, got {len(bias)}") # NOTE (known cost, not a correctness gap): this issues one single-group GEMM per expert group, so a # step pays ~40x more kernel launches than it has GEMMs of work (~40,000 launches for ~1,000 GEMMs, @@ -1382,14 +1380,14 @@ def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) -# Injected into transformer_engine/pytorch/cpp_extensions/gemm.py (prototype). -# Custom-recipe dispatch for the *grouped-tensor* GEMM — the entry point MoE experts actually use -# (GroupedLinear fprop -> general_grouped_gemm_for_grouped_tensor -> nvte_grouped_gemm). +# Custom-recipe dispatch for the *grouped-tensor* GEMM, the entry point GroupedLinear's fprop uses. +# A carries the per-expert weights, B the per-expert activations and `out` the grouped result; +# `offsets` / `scale_inv_offsets` slice each group's rows out of `rowwise_data` / +# `rowwise_scale_inv`, and `num_tensors` is the group count. Only TN fprop is implemented. # -# Authorship: designed by gpt-6.1-sol against the contracts recorded in -# engine/scripts/wip/ds41-b300/worklog.md; the grouped-offset slicing here is the parent's, written -# from the field names read out of tensor/storage/grouped_tensor_storage.py -# (`offsets`, `scale_inv_offsets`, `rowwise_data`, `rowwise_scale_inv`, `num_tensors`). +# Unlike the list form, this layout has no per-group QuantizedTensorStorage to hand a quantizer as +# `qresult_x` / `qresult_w`, so a recipe that derives its global scale from them (non-pow-2 NVFP4) +# is refused inside `qgemm` rather than served wrong. def _custom_grouped_tensor_gemm( @@ -1471,30 +1469,33 @@ def scale_offsets_of(x, name): m_params = MMParams(out_dtype=None, use_split_accumulator=use_split_accumulator) out_dtype = out.rowwise_data.dtype - # scipy-free zip over adjacent offsets + # Every caller of this function that reaches the custom path passes `layout="TN"` with + # A = per-expert weight and B = per-expert activation (`module/grouped_linear.py`, + # `ops/basic/grouped_linear.py`), so the operand that feeds qgemm first is B and the row + # count that has to match the output segment is B's, not A's. for i in range(num): a0, a1 = a_off[i], a_off[i + 1] - rows = a1 - a0 + b0, b1 = b_off[i], b_off[i + 1] + rows = b1 - b0 if rows == 0: continue - b0, b1 = b_off[i], b_off[i + 1] o0, o1 = o_off[i], o_off[i + 1] if o1 - o0 != rows: raise ValueError( - f"group {i}: out segment has {o1 - o0} rows but A has {rows} " - f"(a_off={a_off[:3]}..., o_off={o_off[:3]}...)" + f"group {i}: out segment has {o1 - o0} rows but the activation slice has {rows} " + f"(b_off={b_off[:3]}..., o_off={o_off[:3]}...)" ) - qx = A.rowwise_data[a0:a1] - qw = B.rowwise_data[b0:b1] + qx = B.rowwise_data[b0:b1] + qw = A.rowwise_data[a0:a1] sx = None - if getattr(A, "rowwise_scale_inv", None) is not None and a_s_off is not None: - sx = A.rowwise_scale_inv[a_s_off[i] : a_s_off[i + 1]] - sw = None if getattr(B, "rowwise_scale_inv", None) is not None and b_s_off is not None: - sw = B.rowwise_scale_inv[b_s_off[i] : b_s_off[i + 1]] + sx = B.rowwise_scale_inv[b_s_off[i] : b_s_off[i + 1]] + sw = None + if getattr(A, "rowwise_scale_inv", None) is not None and a_s_off is not None: + sw = A.rowwise_scale_inv[a_s_off[i] : a_s_off[i + 1]] if sx is None or sw is None: raise ValueError(f"group {i}: missing rowwise scale_inv for a custom quantizer") diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index fdfa35e590f..953994e0547 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -392,22 +392,28 @@ def __repr__(self): ")" ) - def quantize_(self, tensor, noop_flag=None): + def quantize_(self, tensor, *, noop_flag=None): """Quantise into this tensor in place. - `module/base.py::quantize_weight` calls `workspace.quantize_(...)` when the - workspace exposes it and otherwise falls back to `tex.quantize(...)`, a C++ - entry that cannot drive a Python reference quantizer — the workspace then stays - empty and the GEMM fails on `data is None`. - - Delegated to the quantizer's `update_quantized`, exactly like the shipped - `QuantizedTensorStorage.quantize_`: that path refreshes every field a later GEMM - reads — including `global_amax_row`/`global_amax_col`, which a non-pow-2 (NVFP4) - recipe turns into the global scale, so copying only the packed data and block - scales would combine fresh values with a stale scale — and honours `noop_flag`, - the flag that lets a replayed graph skip re-quantisation on later microbatches. + Same contract as `QuantizedTensorStorage.quantize_`: the packed data, the block scales + and every other field a later GEMM reads are refreshed — including + `global_amax_row`/`global_amax_col`, which a non-pow-2 (NVFP4) recipe turns into the + global scale, so refreshing only the data and block scales would pair fresh values with + a stale scale. A non-zero `noop_flag` leaves the tensor untouched, which is what lets a + replayed graph skip re-quantisation on later microbatches. Returns `self`. """ + if tensor.dtype not in reference_utils.HIGH_PRECISION_FLOAT_DTYPES: + raise TypeError( + f"Unsupported input dtype {tensor.dtype}, expected one of" + f" {reference_utils.HIGH_PRECISION_FLOAT_DTYPES}" + ) + original_shape = tuple(tensor.shape) + if tensor.ndim < 2: + # `update_quantized` reshapes ndim > 2 only, and `_quantize` below needs a matrix; + # the old `quantize` call folded every non-2-D input into a single row. + tensor = tensor.reshape(1, -1) self._get_quantizer().update_quantized(tensor, self, noop_flag=noop_flag) + self.original_shape = original_shape return self def update_usage( @@ -1254,9 +1260,6 @@ def qgemm( qresult_w: QuantizedTensorStorage | None = None, ) -> torch.Tensor: """Python implementation of microblock FP4 GEMM.""" - if bias is not None: - raise ValueError("Bias is not supported in NVFP4QuantizerRef.qgemm") - high_precision_x = cast_from_fp4x2(qx, out_dtype) high_precision_w = cast_from_fp4x2(qw, out_dtype) @@ -1411,5 +1414,10 @@ def qgemm( if out is not None: raise ValueError("Output tensor should be None when accumulate is False.") + # Bias is an epilogue term, as in the native path: add it while the accumulator is still + # float32, so the result is rounded once when it is cast to `out_dtype`. + if bias is not None: + y = y + bias.to(torch.float32) + y = y.to(out_dtype) return y From 7b876e8ac109e7faee54f2ad21336059c590ac9b Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Sun, 8 Nov 2026 11:16:26 -0800 Subject: [PATCH 09/10] test(grouped_linear): keep the dispatch test out of the swizzle test's body The new test was inserted before `test_swizzle_scales_and_pack_ptrs_for_discrete_weights` finished, so that test's scale-poisoning checks ended up inside the new function, where `swizzle_type`, `shape` and `scale_tensors` do not exist, and the original test lost those checks. Moved to the end of the file: the swizzle test is whole again and the new function only touches its own locals. Signed-off-by: Nolan Ho --- tests/pytorch/test_grouped_linear.py | 48 ++++++++++++++-------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 95c2d1afd72..c0f5c80ee48 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -3803,30 +3803,6 @@ def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( expected_swizzled_scales_buffer, ) - -def test_grouped_tensor_custom_dispatch_classifies_plain_operands(): - """A plain GroupedTensorStorage must not be classified as a custom operand. - - `is_custom` raises on any type that is neither a Quantizer/QuantizedTensorStorage nor a - Tensor, so dispatching on the container itself made every grouped-tensor GEMM raise - before it reached the native implementation. - """ - from transformer_engine.pytorch.cpp_extensions.gemm import _grouped_operand_is_custom - from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import ( - GroupedTensorStorage, - ) - - grouped = GroupedTensorStorage( - shape=(32, 32), - dtype=torch.bfloat16, - num_tensors=2, - quantizer=None, - data=torch.zeros(32, 32, dtype=torch.bfloat16, device="cuda"), - ) - assert not _grouped_operand_is_custom(grouped) - assert not _grouped_operand_is_custom([grouped]) - assert not _grouped_operand_is_custom(None) - # Poison the padded compact scales if swizzle_type == "mxfp8_rowwise": unpadded_scale_shape = (shape[0], shape[1] // 32) @@ -3851,3 +3827,27 @@ def test_grouped_tensor_custom_dispatch_classifies_plain_operands(): swizzled_scales_buffer, expected_swizzled_scales_buffer, ) + + +def test_grouped_tensor_custom_dispatch_classifies_plain_operands(): + """A plain GroupedTensorStorage must not be classified as a custom operand. + + `is_custom` raises on any type that is neither a Quantizer/QuantizedTensorStorage nor a + Tensor, so dispatching on the container itself made every grouped-tensor GEMM raise + before it reached the native implementation. + """ + from transformer_engine.pytorch.cpp_extensions.gemm import _grouped_operand_is_custom + from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import ( + GroupedTensorStorage, + ) + + grouped = GroupedTensorStorage( + shape=(32, 32), + dtype=torch.bfloat16, + num_tensors=2, + quantizer=None, + data=torch.zeros(32, 32, dtype=torch.bfloat16, device="cuda"), + ) + assert not _grouped_operand_is_custom(grouped) + assert not _grouped_operand_is_custom([grouped]) + assert not _grouped_operand_is_custom(None) From e1ffa249971543b1f84b19f6e9ade394a11dc985 Mon Sep 17 00:00:00 2001 From: Nolan Ho Date: Sun, 8 Nov 2026 13:15:22 -0800 Subject: [PATCH 10/10] fix(custom_recipes): reach the grouped-tensor block scales, and cover the custom grouped paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_custom_grouped_tensor_gemm` read the block scales as `rowwise_scale_inv`, a field GroupedTensorStorage does not have — it stores them as `scale_inv` — so every group died on "missing rowwise scale_inv for a custom quantizer" as soon as the dispatch above it stopped raising. Both names are accepted now, keyed off the buffer the storage actually carries. Nothing covered this. The module and ops grouped paths only build custom operands for the discrete (list) form, so the reference `qgemm` implementations were never exercised by the suite and the defect above only surfaced once the path was driven by hand on a device. The new `tests/pytorch/test_custom_recipe_grouped.py` drives both entry points with the reference NVFP4/MXFP4 quantizer and pins what has no other coverage: per-expert operand roles and m_splits slicing, per-group offsets and block scales, the per-expert bias, the `qresult_x`/`qresult_w` hand-off that non-pow-2 NVFP4 needs, the documented refusals (dgrad/wgrad, non-TN layout, malformed bias list, non-pow-2 recipe on the grouped-tensor layout), and the reference quantizer's in-place update and fp32-epilogue-bias contracts. Signed-off-by: Nolan Ho --- tests/pytorch/test_custom_recipe_grouped.py | 293 ++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 20 +- 2 files changed, 309 insertions(+), 4 deletions(-) create mode 100644 tests/pytorch/test_custom_recipe_grouped.py diff --git a/tests/pytorch/test_custom_recipe_grouped.py b/tests/pytorch/test_custom_recipe_grouped.py new file mode 100644 index 00000000000..aa14459985a --- /dev/null +++ b/tests/pytorch/test_custom_recipe_grouped.py @@ -0,0 +1,293 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped GEMM with custom (reference) quantizers. + +These drive the two custom-recipe grouped entry points with the reference NVFP4/MXFP4 quantizer, +which is the only in-tree way to reach the Python `qgemm` implementations: the module and ops +grouped paths only ever build custom operands for the discrete (list) form, and the reference +quantizer is what makes a `qgemm` callable without a compiled custom kernel. + +Covered here, because nothing else in the suite reaches it: + * `_custom_grouped_gemm` (list form) — operand roles, m_splits slicing, per-expert bias, and the + `qresult_x`/`qresult_w` hand-off that non-pow-2 NVFP4 needs. + * `_custom_grouped_tensor_gemm` (GroupedTensorStorage form) — per-group offsets and block scales. + * The refusals: dgrad/wgrad, non-TN layouts, malformed bias lists, and the non-pow-2 recipe that + this layout cannot serve. +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.cpp_extensions import ( + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) +from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.custom_recipes.gemm import GEMMType, MMParams +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import GroupedTensorStorage + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +SPLITS = [16, 32, 8, 24] +K, N = 256, 128 +TOTAL_M = sum(SPLITS) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="grouped GEMM requires a CUDA device" +) + + +def _quantizer(pow_2_scales: bool) -> NVFP4QuantizerRef: + # MXFP4 (E8M0 scales) is defined on 1x32 tiles, NVFP4 on 1x16. + return NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, + pow_2_scales=pow_2_scales, + quant_tile_shape=(1, 32) if pow_2_scales else (1, 16), + ) + + +def _operands(): + torch.manual_seed(0) + weights = [torch.randn(N, K, device=DEVICE, dtype=DTYPE) for _ in SPLITS] + inputs = [torch.randn(m, K, device=DEVICE, dtype=DTYPE) for m in SPLITS] + biases = [torch.randn(N, device=DEVICE, dtype=DTYPE) for _ in SPLITS] + return weights, inputs, biases + + +def _per_group_reference(quantizer, qweights, qinputs, biases=None) -> torch.Tensor: + """One qgemm per expert, placed at that expert's row offset: what the dispatch must produce.""" + expected = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + offset = 0 + for i, m in enumerate(SPLITS): + expected[offset : offset + m] = quantizer.qgemm( + qinputs[i].data[:m], + qweights[i].data, + MMParams(out_dtype=DTYPE, use_split_accumulator=False), + DTYPE, + qinputs[i].scale[:m], + qweights[i].scale, + biases[i] if biases is not None else None, + gemm_type=GEMMType.FPROP, + qresult_x=qinputs[i], + qresult_w=qweights[i], + ) + offset += m + return expected + + +def _assert_close_to_unquantized(actual: torch.Tensor, weights, inputs, biases=None) -> None: + """4-bit operands are coarse, so only gross (operand-role) mistakes are ruled out here.""" + offset = 0 + for i, m in enumerate(SPLITS): + expected = torch.nn.functional.linear( + inputs[i], weights[i], biases[i] if biases is not None else None + ) + block = actual[offset : offset + m] + deviation = (block - expected).norm() / expected.norm().clamp_min(1e-6) + assert deviation < 0.25, f"expert {i}: quantized result deviates {deviation:.1%}" + offset += m + + +@pytest.mark.parametrize("use_bias", [False, True], ids=["no-bias", "bias"]) +@pytest.mark.parametrize("pow_2_scales", [True, False], ids=["mxfp4", "nvfp4"]) +def test_custom_grouped_gemm_list_form(use_bias: bool, pow_2_scales: bool) -> None: + """`general_grouped_gemm` dispatches per-expert custom operands to the quantizer's qgemm.""" + quantizer = _quantizer(pow_2_scales) + weights, inputs, biases = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + + out = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + general_grouped_gemm( + qweights, + qinputs, + [out], + [None] * len(SPLITS), + DTYPE, + single_output=True, + m_splits=SPLITS, + bias=biases if use_bias else None, + use_bias=use_bias, + ) + + torch.testing.assert_close( + out, _per_group_reference(quantizer, qweights, qinputs, biases if use_bias else None) + ) + _assert_close_to_unquantized(out, weights, inputs, biases if use_bias else None) + + +def test_custom_grouped_tensor_gemm() -> None: + """`general_grouped_gemm_for_grouped_tensor` slices groups out of GroupedTensorStorage.""" + quantizer = _quantizer(pow_2_scales=True) + weights, inputs, _ = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + + row_offsets = [0] + for m in SPLITS: + row_offsets.append(row_offsets[-1] + m) + weight_offsets = [i * N for i in range(len(SPLITS) + 1)] + + grouped_weights = GroupedTensorStorage( + shape=(len(SPLITS) * N, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qweights], dim=0), + scale_inv=torch.cat([t.scale for t in qweights], dim=0), + offsets=weight_offsets, + scale_inv_offsets=weight_offsets, + ) + grouped_inputs = GroupedTensorStorage( + shape=(TOTAL_M, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qinputs], dim=0), + scale_inv=torch.cat([t.scale for t in qinputs], dim=0), + offsets=row_offsets, + scale_inv_offsets=row_offsets, + ) + grouped_out = GroupedTensorStorage( + shape=(TOTAL_M, N), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=None, + data=torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE), + offsets=row_offsets, + ) + + general_grouped_gemm_for_grouped_tensor( + grouped_weights, grouped_inputs, grouped_out, layout="TN" + ) + + torch.testing.assert_close( + grouped_out.rowwise_data, + _per_group_reference(quantizer, qweights, qinputs), + ) + _assert_close_to_unquantized(grouped_out.rowwise_data, weights, inputs) + + +def test_custom_grouped_tensor_gemm_refuses_nvfp4() -> None: + """This layout carries no per-group quantized tensors, so a non-pow-2 recipe must refuse.""" + quantizer = _quantizer(pow_2_scales=False) + weights, inputs, _ = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + row_offsets = [0] + for m in SPLITS: + row_offsets.append(row_offsets[-1] + m) + weight_offsets = [i * N for i in range(len(SPLITS) + 1)] + + grouped_weights = GroupedTensorStorage( + shape=(len(SPLITS) * N, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qweights], dim=0), + scale_inv=torch.cat([t.scale for t in qweights], dim=0), + offsets=weight_offsets, + scale_inv_offsets=weight_offsets, + ) + grouped_inputs = GroupedTensorStorage( + shape=(TOTAL_M, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qinputs], dim=0), + scale_inv=torch.cat([t.scale for t in qinputs], dim=0), + offsets=row_offsets, + scale_inv_offsets=row_offsets, + ) + grouped_out = GroupedTensorStorage( + shape=(TOTAL_M, N), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=None, + data=torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE), + offsets=row_offsets, + ) + with pytest.raises(ValueError, match="qresult"): + general_grouped_gemm_for_grouped_tensor( + grouped_weights, grouped_inputs, grouped_out, layout="TN" + ) + + +def test_custom_grouped_gemm_refuses_unsupported_contracts() -> None: + """dgrad/wgrad, non-TN layouts and malformed bias lists are refused before any work.""" + quantizer = _quantizer(pow_2_scales=True) + weights, inputs, biases = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + out = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + + def call(**kwargs): + general_grouped_gemm( + qweights, + qinputs, + [out], + [None] * len(SPLITS), + DTYPE, + single_output=True, + m_splits=SPLITS, + **kwargs, + ) + + with pytest.raises(NotImplementedError, match="fprop only"): + call(grad=True) + with pytest.raises(NotImplementedError, match="TN layout"): + call(layout="NN") + with pytest.raises(ValueError, match="per-group bias"): + call(use_bias=True) + with pytest.raises(ValueError, match="expected 4 biases"): + call(bias=biases[:2], use_bias=True) + + +def test_reference_quantizer_in_place_update_contract() -> None: + """`quantize_` refreshes every field a later GEMM reads, honours noop_flag, returns self.""" + quantizer = _quantizer(pow_2_scales=False) + tensor = torch.randn(32, K, device=DEVICE, dtype=DTYPE) + qresult = quantizer(tensor) + before_amax = qresult.global_amax_row.clone() + + assert qresult.quantize_(tensor * 4.0) is qresult + assert not torch.equal(before_amax, qresult.global_amax_row), "stale global amax after refresh" + assert tuple(qresult.original_shape) == (32, K) + + snapshot = qresult.data.clone() + qresult.quantize_(tensor * 0.25, noop_flag=torch.tensor([1.0], device=DEVICE)) + assert torch.equal(snapshot, qresult.data), "noop_flag was ignored" + + +def test_reference_qgemm_applies_bias_in_fp32_epilogue() -> None: + """The bias is an epilogue term, so qgemm(bias=x) == qgemm() + x rather than an error.""" + quantizer = _quantizer(pow_2_scales=True) + activations = quantizer(torch.randn(16, K, device=DEVICE, dtype=DTYPE)) + weights = quantizer(torch.randn(N, K, device=DEVICE, dtype=DTYPE)) + bias = torch.randn(N, device=DEVICE, dtype=DTYPE) + params = MMParams(out_dtype=DTYPE, use_split_accumulator=False) + + plain = quantizer.qgemm( + activations.data, + weights.data, + params, + DTYPE, + activations.scale, + weights.scale, + gemm_type=GEMMType.FPROP, + ) + with_bias = quantizer.qgemm( + activations.data, + weights.data, + params, + DTYPE, + activations.scale, + weights.scale, + bias, + gemm_type=GEMMType.FPROP, + ) + torch.testing.assert_close(with_bias, plain + bias, rtol=2e-2, atol=2e-2) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 93273ad3c66..469f3ea771d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1490,12 +1490,24 @@ def scale_offsets_of(x, name): qx = B.rowwise_data[b0:b1] qw = A.rowwise_data[a0:a1] + # GroupedTensorStorage names its block scales `scale_inv`; `rowwise_scale_inv` is the name + # the per-tensor storages use, so accept either rather than only the one this entry point + # never actually receives. + def block_scales(x): + for field in ("scale_inv", "rowwise_scale_inv"): + buffer = getattr(x, field, None) + if buffer is not None: + return buffer + return None + sx = None - if getattr(B, "rowwise_scale_inv", None) is not None and b_s_off is not None: - sx = B.rowwise_scale_inv[b_s_off[i] : b_s_off[i + 1]] + b_scales = block_scales(B) + if b_scales is not None and b_s_off is not None: + sx = b_scales[b_s_off[i] : b_s_off[i + 1]] sw = None - if getattr(A, "rowwise_scale_inv", None) is not None and a_s_off is not None: - sw = A.rowwise_scale_inv[a_s_off[i] : a_s_off[i + 1]] + a_scales = block_scales(A) + if a_scales is not None and a_s_off is not None: + sw = a_scales[a_s_off[i] : a_s_off[i + 1]] if sx is None or sw is None: raise ValueError(f"group {i}: missing rowwise scale_inv for a custom quantizer")