From 60b0ee461e891e2712116c5b38156c703d7163d5 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Wed, 28 Oct 2026 01:24:15 -0700 Subject: [PATCH 1/5] [PyTorch] Quantize fused GroupedMLP fc2 input row-wise only when it is recomputed in backward With activation recompute in the fused NVFP4 GroupedMLP, the backward rebuilds fc2_x from the dSReLU kernel and re-quantizes it column-wise for wgrad. The forward still quantized fc2_x row- and column-wise and then dropped the column-wise copy at save time. Decide recompute_srelu_fc2_x before the fc2 input quantization, request only row-wise data when the backward will rebuild the tensor, and mark for paged stashing only the tensor that is actually saved for backward. Rowwise FP4 data and the fc2 GEMM result are unchanged; the group quantize kernel for the fc2 input drops from ~570us to ~330us per fused call on GB300 for a 130B hybrid MoE proxy. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- .../pytorch/ops/fused/grouped_mlp.py | 35 +++++++++++-------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d29ae1fd43..3581f25583 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1936,9 +1936,24 @@ def fuser_forward( ): grouped_fc2_weight._with_gemm_swizzled_scales = False + # Decide up front whether the backward rebuilds fc2_x from the dSReLU kernel + # (which applies the tanh clamp when regenerating fc2_x). When it does, the + # forward only needs row-wise fc2_x for the fc2 GEMM; the column-wise copy + # for wgrad is produced in the backward from the recomputed tensor. + recompute_srelu_fc2_x = ( + isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) + and bool(getattr(activation_op, "activation_recompute_in_mlp", False)) + and weight_requires_grad + and _grouped_gemm_dsrelu_backward_supported() + and _nvidia_cudnn_frontend_supports_wgrad() + ) + # Unpack kernel outputs if use_nvfp4: - fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc2_input_quantizer.set_usage( + rowwise=True, + columnwise=weight_requires_grad and not recompute_srelu_fc2_x, + ) fc2_input_quantizer.optimize_for_gemm = True if kernel_impl == "gemm_act": # Quantize to NVFP4 @@ -2266,22 +2281,12 @@ def fuser_forward( # Save state for backward pass if requires_grad: saved_fc1_x = grouped_fc1_x - mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) - activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() - # The dSReLU kernel applies the tanh clamp when regenerating fc2_x. - activation_is_srelu = isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) - activation_recompute_in_mlp = bool( - getattr(activation_op, "activation_recompute_in_mlp", False) - ) - recompute_srelu_fc2_x = ( - activation_is_srelu - and activation_recompute_in_mlp - and weight_requires_grad - and _grouped_gemm_dsrelu_backward_supported() - and _nvidia_cudnn_frontend_supports_wgrad() - ) + # recompute_srelu_fc2_x was decided before the fc2 input quantization above. + # A recomputed fc2_x is row-wise only and is not saved, so only mark what + # actually goes to save_for_backward. saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x + mark_grouped_tensor(saved_fc1_x, activation_in, scales, saved_grouped_fc2_x) # wgrad only needs column-wise data, so clear row-wise data for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): From 7acab7de30b3869bf7a006b574b40d8a5c55afb2 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Wed, 28 Oct 2026 20:38:36 -0700 Subject: [PATCH 2/5] [PyTorch] Reuse the forward NVFP4 amax of fc2 input when it is recomputed in backward With activation recompute in the fused NVFP4 GroupedMLP, the backward regenerates fc2_x with the dSReLU kernel and re-quantizes it column-wise, which first runs the group Hadamard-amax kernel over the whole tensor and then the cast. The regenerated tensor is bit-identical to the forward one (the plain SReLU kernel rounds the accumulator through C dtype before the activation; the Hadamard-fused kernel does so with the matching cuDNN frontend fix), so its per-group amax is the forward amax. Keep the forward (pre-RHT, post-RHT) amax of fc2_x on the fc2 context and quantize the regenerated tensor with nvfp4_group_quantize_with_amax, which drops the backward amax pass. The Hadamard-fused fc1 kernel already returns both amax values; for the plain SReLU kernel add nvfp4_group_hadamard_amax, a binding that runs only the group Hadamard-amax kernel, so the forward computes both values in the single pass it needs for the row-wise cast. Without the binding the previous behavior is kept. NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX=0 disables the reuse. Verified on GB300 with a fused grouped MLP (8 experts, NVFP4 + RHT): with recompute on and off, output, dgrad, dprob and both weight gradients match bitwise for tanh-SReLU and plain SReLU, and the backward launches one fewer Hadamard-amax kernel per fused call where the forward amax is available. Co-Authored-By: Claude Fable 5.1 The recompute test now asserts a bitwise-equal FC2 weight gradient for tanh-SReLU with NVFP4; plain SReLU keeps its tolerance until the cuDNN frontend fix for the Hadamard-fused kernel is required. Signed-off-by: Xiaowei Ren --- tests/pytorch/test_grouped_mlp.py | 15 ++- transformer_engine/pytorch/csrc/extensions.h | 14 ++ .../pytorch/csrc/extensions/cast.cpp | 71 ++++++++++ .../pytorch/csrc/extensions/pybind.cpp | 4 + .../pytorch/ops/fused/grouped_mlp.py | 123 +++++++++++++++--- 5 files changed, 204 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 18f9d3f772..6adb641372 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2854,9 +2854,18 @@ def test_grouped_mlp_activation_recompute( if bias: torch.testing.assert_close(on["fc1_db"], off["fc1_db"], **tols) torch.testing.assert_close(on["fc2_db"], off["fc2_db"], **tols) - assert_close_rms( - on["fc2_dw"], off["fc2_dw"], rtol=self._RECOMPUTE_FC2_WGRAD_RMS_TOL[quantization] - ) + # The dSReLU kernel regenerates fc2_x bit-exactly from the saved pre-activation and + # the NVFP4 backward re-quantizes it with the forward amax, so FC2's weight gradient + # matches bitwise. Plain SReLU needs a cuDNN frontend whose Hadamard-fused kernel + # rounds the accumulator through C's dtype before the activation; older versions + # regenerate a tensor that differs by a few BF16 ulps, hence the tolerance. MXFP8 + # quantizes the regenerated tensor inside the dSReLU kernel with its own scales. + if quantization == "nvfp4_rht" and activation == "scaled_tanh_srelu": + torch.testing.assert_close(on["fc2_dw"], off["fc2_dw"], rtol=0, atol=0) + else: + assert_close_rms( + on["fc2_dw"], off["fc2_dw"], rtol=self._RECOMPUTE_FC2_WGRAD_RMS_TOL[quantization] + ) @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) def test_grouped_mlp_caller_buffers( diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 7bcd8430ce..4f38b6fd5e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -427,6 +427,20 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q const at::Tensor &columnwise_amax, std::optional tensor_offsets); +/*! @brief Per-group NVFP4 amax of a grouped BF16 tensor without quantizing it. + * + * Returns (pre-RHT amax, post-RHT amax), one float32 value per group, as the group + * Hadamard-amax kernel computes them for the RHT recipe. Lets a caller quantize row-wise + * now and column-wise later from the same amax, e.g. when the column-wise copy is + * regenerated in the backward pass. + */ +std::tuple nvfp4_group_hadamard_amax(const at::Tensor &tensor, + py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + std::optional last_dims, + std::optional tensor_offsets); + py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8c9b6b1079..a04dfd0668 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -691,6 +691,77 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q return py::reinterpret_borrow(grouped_output_py); } +std::tuple nvfp4_group_hadamard_amax(const at::Tensor &tensor, + py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + std::optional last_dims, + std::optional tensor_offsets) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(!last_dims.has_value(), + "nvfp4_group_hadamard_amax: varying last dim is not supported with NVFP4."); + NVTE_CHECK(first_dims.has_value() && tensor_offsets.has_value(), + "nvfp4_group_hadamard_amax: first_dims and tensor_offsets are required."); + NVTE_CHECK(first_dims->is_cuda() && first_dims->scalar_type() == at::kLong && + first_dims->numel() == static_cast(num_tensors), + "nvfp4_group_hadamard_amax: first_dims must be an int64 CUDA tensor with one " + "entry per group."); + NVTE_CHECK(tensor_offsets->is_cuda() && tensor_offsets->scalar_type() == at::kLong && + tensor_offsets->numel() == static_cast(num_tensors + 1), + "nvfp4_group_hadamard_amax: tensor_offsets must be an int64 CUDA tensor with " + "num_tensors + 1 entries."); + + auto quantizer_cpp = convert_quantizer(quantizer); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "nvfp4_group_hadamard_amax only supports NVFP4 quantizers."); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + NVTE_CHECK(nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax, + "nvfp4_group_hadamard_amax expects the NVFP4 RHT recipe with post-RHT amax."); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto float_opts = at::TensorOptions().dtype(at::kFloat).device(tensor.device()); + auto rowwise_amax = at::zeros({static_cast(num_tensors)}, float_opts); + auto columnwise_amax = at::zeros({static_cast(num_tensors)}, float_opts); + if (logical_shape[0] == 0 || logical_shape[1] == 0) { + return {rowwise_amax, columnwise_amax}; + } + + const auto input_dtype = GetTransformerEngineDType(tensor.scalar_type()); + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); + + // The amax launcher derives which amax to compute from the storage the output declares: + // row-wise data requests the pre-RHT amax and column-wise data the post-RHT (transposed) + // amax. It only ever writes the amax buffers, so alias the input as both and allocate no + // quantized storage. + auto grouped_amax_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_amax_tensor.set_rowwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); + grouped_amax_tensor.set_columnwise_data(tensor.data_ptr(), input_dtype, + getTensorShape(tensor)); + grouped_amax_tensor.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(rowwise_amax)); + grouped_amax_tensor.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(columnwise_amax)); + grouped_amax_tensor.set_first_dims(first_dims->data_ptr(), DType::kInt64, + getTensorShape(*first_dims)); + grouped_amax_tensor.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_amax_graph_safe( + grouped_input_tensor.data(), grouped_amax_tensor.data(), 0, + nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, at::cuda::getCurrentCUDAStream()); + }); + + return {rowwise_amax, columnwise_amax}; +} + py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 6fbed1436f..2c25aae11b 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -187,6 +187,10 @@ void bind_quantize_with_amax_extensions(py::module_ &m) { py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), py::arg("tensor_offsets") = py::none()); + m.def("nvfp4_group_hadamard_amax", nvfp4_group_hadamard_amax, py::arg("tensor"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), + "Per-group (pre-RHT, post-RHT) NVFP4 amax of a grouped tensor without quantizing it"); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 3581f25583..018f9696b8 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -384,6 +384,41 @@ def _group_quantize_with_amax_for_grouped_mlp( ) +def _fc2_x_amax_reuse_enabled() -> bool: + """Whether the backward reuses the forward NVFP4 amax of a recomputed fc2 input. + + The grouped dSReLU kernel regenerates fc2_x bit-exactly from the saved pre-activation, + so the per-group amax computed in the forward is also the amax of the regenerated + tensor. Reusing it turns the backward re-quantization into a single cast kernel. + ``NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX=0`` disables the reuse. + """ + return int(os.environ.get("NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX", "1")) > 0 + + +def _nvfp4_group_hadamard_amax( + tensor: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: Optional[torch.Tensor], + *, + tensor_offsets: Optional[torch.Tensor] = None, +) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Per-group (pre-RHT, post-RHT) NVFP4 amax of ``tensor`` without quantizing it. + + Returns ``None`` when the extension lacks the kernel binding or the grouped metadata + is incomplete; callers then quantize with amax computed on the fly. + """ + amax_fn = getattr(tex, "nvfp4_group_hadamard_amax", None) + if ( + amax_fn is None + or not isinstance(quantizer, NVFP4Quantizer) + or split_sizes is None + or tensor_offsets is None + ): + return None + return amax_fn(tensor, quantizer, num_groups, split_sizes, tensor_offsets=tensor_offsets) + + def _weight_quantizers(op: GroupedLinear, num_groups: int) -> list[Quantizer]: """Per-expert weight quantizers of ``op``. @@ -1947,6 +1982,10 @@ def fuser_forward( and _grouped_gemm_dsrelu_backward_supported() and _nvidia_cudnn_frontend_supports_wgrad() ) + # Forward (pre-RHT, post-RHT) amax of fc2_x, kept for its column-wise re-quantization + # in the backward when the regenerated tensor is bit-identical to this one. + fc2_x_amaxes_for_recompute = None + reuse_fc2_x_amax = recompute_srelu_fc2_x and _fc2_x_amax_reuse_enabled() # Unpack kernel outputs if use_nvfp4: @@ -1956,31 +1995,59 @@ def fuser_forward( ) fc2_input_quantizer.optimize_for_gemm = True if kernel_impl == "gemm_act": - # Quantize to NVFP4 + # Quantize to NVFP4. When the backward rebuilds fc2_x, compute the amax + # here with the kernel the row-wise quantization needs anyway and keep + # it, so the backward re-quantization is only the column-wise cast. fc2_in = fc1_kernel_out["d_tensor"] fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_in, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - use_dense_single_group=use_dense_single_group, - ) + fc2_x_amaxes = None + if reuse_fc2_x_amax: + fc2_x_amaxes = _nvfp4_group_hadamard_amax( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + ) + if fc2_x_amaxes is not None: + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + *fc2_x_amaxes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) + fc2_x_amaxes_for_recompute = fc2_x_amaxes + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) elif kernel_impl == "gemm_act_rht_amax": - # Quantize to NVFP4 using precomputed amax + # Quantize to NVFP4 using the amax the fc1 kernel computed fc2_in = fc1_kernel_out["d_tensor"] fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + fc2_x_amaxes = ( + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), + ) grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( fc2_in, fc2_input_quantizer, num_groups, split_sizes, - fc1_kernel_out["amax_tensor"].view(-1), - fc1_kernel_out["post_rht_amax_tensor"].view(-1), + *fc2_x_amaxes, tensor_offsets=fc2_x_tensor_offsets, use_dense_single_group=use_dense_single_group, ) + if reuse_fc2_x_amax: + fc2_x_amaxes_for_recompute = fc2_x_amaxes elif kernel_impl == "gemm_act_rht_quant": # Unpack NVFP4 output fc2_in_row_data = fc1_kernel_out["d_tensor"] @@ -2351,6 +2418,7 @@ def fuser_forward( fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x + fc2_ctx.fc2_x_amaxes_for_recompute = fc2_x_amaxes_for_recompute return fc2_out, [(), (), ()] @@ -2838,14 +2906,29 @@ def fuser_backward( fc2_input_quantizer = fc2_ctx.input_quantizers[0] fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) fc2_input_quantizer.optimize_for_gemm = True - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_x_bf16, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - use_dense_single_group=use_dense_single_group, - ) + fc2_x_amaxes = getattr(fc2_ctx, "fc2_x_amaxes_for_recompute", None) + fc2_ctx.fc2_x_amaxes_for_recompute = None + if fc2_x_amaxes is not None: + # The regenerated fc2_x is bit-identical to the forward one, so its + # forward amax applies and the Hadamard amax pass is skipped. + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_x_bf16, + fc2_input_quantizer, + num_groups, + split_sizes, + *fc2_x_amaxes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_x_bf16, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") if sfd_col_d_srelu_tensor is None: From 67d47bc34f3fdea400e8a081f5cc8395dae1aab0 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Wed, 28 Oct 2026 22:18:27 -0700 Subject: [PATCH 3/5] [PyTorch] Use the Hadamard-fused fc1 kernel for tanh-SReLU when the cuDNN frontend supports it ScaledTanhSReLU was routed to the plain grouped SReLU kernel because the Hadamard-fused kernel had no soft-clamp support. With NVFP4 RHT that costs a separate Hadamard-amax pass over the fc2 input in every forward. The cuDNN frontend now accepts tanh_clamp_scale on grouped_gemm_glu_hadamard_wrapper_sm100; detect that by signature and let the clamped op take the same fused path as the unclamped one. The kernel then returns both NVFP4 amax values, so the forward amax pass disappears and the backward re-quantization of the recomputed fc2 input reuses them. Verified on GB300 with a fused grouped MLP (8 experts, NVFP4 + RHT): the tanh-SReLU forward launches one Hadamard-amax kernel fewer per fused call, the backward one fewer with recompute, and recompute on/off give bitwise- equal output and gradients. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- .../pytorch/ops/fused/grouped_mlp.py | 35 ++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 018f9696b8..b56bf8007b 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -126,6 +126,32 @@ def _cudnn_frontend_supports_grouped_gemm_srelu_tanh() -> bool: return False +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard() -> bool: + """Whether the cuDNN FE grouped GEMM + Hadamard wrapper accepts tanh_clamp_scale. + + With it, the clamped SReLU uses the fused RHT-amax kernel like the unclamped one + instead of the plain SReLU kernel plus a separate amax pass. Detected by signature + for the same reason as _cudnn_frontend_supports_grouped_gemm_srelu_tanh. + """ + if not ( + _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() + and _cudnn_frontend_supports_grouped_gemm_srelu_tanh() + ): + return False + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_glu_hadamard_wrapper_sm100, + ) + except ImportError: + return False + try: + signature = inspect.signature(grouped_gemm_glu_hadamard_wrapper_sm100) + return "tanh_clamp_scale" in signature.parameters + except (TypeError, ValueError): + return False + + @functools.lru_cache(maxsize=None) def _cudnn_frontend_supports_grouped_gemm_situglu() -> bool: """Feature-detect complete cuDNN frontend grouped SiTU-GLU support.""" @@ -1472,7 +1498,14 @@ def fuser_forward( if unit_activation_scale and num_groups != 1: unit_activation_scale = False - activation_is_srelu = isinstance(activation_op, ScaledSReLU) + # The hadamard kernel is the GLU one; it applies the SReLU soft clamp only in + # cuDNN frontends whose wrapper accepts tanh_clamp_scale. Without that, + # ScaledTanhSReLU keeps the plain SReLU kernel and a separate RHT amax pass -- + # a performance limitation, not a correctness one. + activation_is_srelu = isinstance(activation_op, ScaledSReLU) or ( + isinstance(activation_op, ScaledTanhSReLU) + and _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard() + ) supports_single_group_runtime_offsets = ( _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) From c351077bec003dd496afaf11bc51af103cc476f1 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Wed, 28 Oct 2026 23:22:15 -0700 Subject: [PATCH 4/5] [PyTorch] Pass the NVFP4 RHT sign mask to the Hadamard-fused grouped GEMM kernel The fused fc1 GEMM + activation + Hadamard kernel returns the post-RHT amax that sets the global scale of the column-wise NVFP4 cast of fc2 input. The NVFP4 quantizer applies a randomized Hadamard transform with a per-quantizer sign mask, but the kernel computed the amax for the plain transform, so the scale was off by the sign pattern: up to ~17% too large for SReLU-shaped activations and possibly too small for GLUs. cuDNN frontend now accepts rht_sign_mask on grouped_gemm_glu_hadamard_wrapper_sm100 and also rounds its amax values like the quantizer does; pass the quantizer mask when the wrapper takes it. On frontends without it, use the fused kernel only when the mask is zero, so a wrong amax is never fed to the cast. Verified on GB300 (8 experts, NVFP4 + RHT): with the matching cuDNN frontend the fused amax values are bit-identical to the standalone group Hadamard-amax kernel for SReLU, tanh-SReLU and SwiGLU, and tanh-SReLU gives bitwise-equal output and gradients on the fused and the plain kernel path. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- .../pytorch/ops/fused/grouped_mlp.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index b56bf8007b..4fed03a2f4 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -152,6 +152,29 @@ def _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard() -> bool: return False +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask() -> bool: + """Whether the cuDNN FE grouped GEMM + Hadamard wrapper accepts rht_sign_mask. + + The kernel's post-RHT amax must be computed for the same randomized Hadamard + transform the NVFP4 quantizer applies; without the mask it describes the plain + transform and the column-wise cast gets a wrong global scale. + """ + if not _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): + return False + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_glu_hadamard_wrapper_sm100, + ) + except ImportError: + return False + try: + signature = inspect.signature(grouped_gemm_glu_hadamard_wrapper_sm100) + return "rht_sign_mask" in signature.parameters + except (TypeError, ValueError): + return False + + @functools.lru_cache(maxsize=None) def _cudnn_frontend_supports_grouped_gemm_situglu() -> bool: """Feature-detect complete cuDNN frontend grouped SiTU-GLU support.""" @@ -1726,6 +1749,13 @@ def fuser_forward( if self.grouped_gemm_act_hadamard_kernel() is None: # Kernel is not available pass + elif ( + fc2_input_quantizer.rht_matrix_random_sign_mask_t != 0 + and not _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask() + ): + # Kernel would compute the post-RHT amax for the plain Hadamard + # transform while the quantizer applies a randomized one + pass elif self._cudnn_act_func in ("swiglu", "situglu"): kernel_impl = "gemm_act_rht_amax" elif activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): @@ -1853,6 +1883,10 @@ def fuser_forward( elif kernel_impl == "gemm_act_rht_amax": fc1_activation_kwargs["d_dtype"] = torch.bfloat16 fc1_activation_kwargs["use_tmem_post_rht_amax"] = _use_tmem_post_rht_amax() + if _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask(): + fc1_activation_kwargs["rht_sign_mask"] = ( + fc2_input_quantizer.rht_matrix_random_sign_mask_t + ) elif kernel_impl == "gemm_act_rht_quant": fc1_activation_kwargs["d_dtype"] = torch.float4_e2m1fn_x2 fc1_activation_kwargs["rht_colwise_dtype"] = torch.float4_e2m1fn_x2 From 5eace13b0fc4ab933e8c6c7000bc46a89f9c25fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 10:30:51 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions.h | 10 ++++------ transformer_engine/pytorch/csrc/extensions/cast.cpp | 13 +++++-------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 4f38b6fd5e..02bf677344 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -434,12 +434,10 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q * now and column-wise later from the same amax, e.g. when the column-wise copy is * regenerated in the backward pass. */ -std::tuple nvfp4_group_hadamard_amax(const at::Tensor &tensor, - py::handle quantizer, - const size_t num_tensors, - std::optional first_dims, - std::optional last_dims, - std::optional tensor_offsets); +std::tuple nvfp4_group_hadamard_amax( + const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets); py::object group_dequantize(const py::handle &input, DType otype); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index a04dfd0668..3191d1b7bd 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -691,12 +691,10 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q return py::reinterpret_borrow(grouped_output_py); } -std::tuple nvfp4_group_hadamard_amax(const at::Tensor &tensor, - py::handle quantizer, - const size_t num_tensors, - std::optional first_dims, - std::optional last_dims, - std::optional tensor_offsets) { +std::tuple nvfp4_group_hadamard_amax( + const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -742,8 +740,7 @@ std::tuple nvfp4_group_hadamard_amax(const at::Tensor &t // quantized storage. auto grouped_amax_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_amax_tensor.set_rowwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); - grouped_amax_tensor.set_columnwise_data(tensor.data_ptr(), input_dtype, - getTensorShape(tensor)); + grouped_amax_tensor.set_columnwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); grouped_amax_tensor.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, getTensorShape(rowwise_amax)); grouped_amax_tensor.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32,