GEMM + Swiglu fused Grouped MLP for MXFP8#2769
GEMM + Swiglu fused Grouped MLP for MXFP8#2769ksivaman wants to merge 11 commits intoNVIDIA:mainfrom
Conversation
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Greptile SummaryThis PR adds a fused GEMM + SwiGLU forward/backward kernel for MXFP8 Grouped MLP on SM100+ (Blackwell) hardware, wiring together experimental CuTe DSL kernels from the cuDNN front-end with the existing Key concerns remaining from previous review rounds (not yet resolved per the thread history):
New findings from this review:
Confidence Score: 2/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Input as Input Tensor<br/>(M×K, BF16)
participant FwdFused as ForwardGroupedMLP<br/>CuTeGEMMSwiGLU_MXFP8
participant CuDNN as cudnn<br/>grouped_gemm_swiglu<br/>_wrapper_sm100
participant FC2GEMM as general_grouped_gemm<br/>FC2 (cuBLAS)
participant BwdFused as BackwardGroupedMLP<br/>CuTeGEMMDSwiGLU_MXFP8
participant CuDNN2 as cudnn<br/>grouped_gemm_dswiglu<br/>_wrapper_sm100
Note over Input,FC2GEMM: Forward Pass (MXFP8, SM100+)
Input->>FwdFused: x [M×K], split_sizes, scales
FwdFused->>FwdFused: group_quantize(x) → grouped_fc1_x (FP8)
FwdFused->>FwdFused: quantize FC1/FC2 weights (FP8)
FwdFused->>CuDNN: fc1_x_data, fc1_w_data,<br/>fc1_x_scales, fc1_w_scales,<br/>split_points, alpha_tensor, prob_scales
CuDNN-->>FwdFused: c_tensor (BF16, pre-SwiGLU)<br/>d_tensor (FP8, post-SwiGLU row)<br/>d_col_tensor (FP8, post-SwiGLU col)<br/>sfd_row/col (MXFP8 scales)
FwdFused->>FwdFused: pack grouped_fc2_x<br/>from d_tensor + sfd scales
FwdFused->>FC2GEMM: grouped_fc2_weight × grouped_fc2_x
FC2GEMM-->>FwdFused: fc2_out [M×N]
FwdFused->>FwdFused: save swiglu_in, fc1/fc2 weights,<br/>fc1/fc2 inputs for backward
Note over Input,CuDNN2: Backward Pass (MXFP8, SM100+)
FwdFused->>BwdFused: grad_output [M×N]
BwdFused->>BwdFused: group_quantize(dy) → grouped_fc2_dy (FP8)
BwdFused->>CuDNN2: fc2_dy_data, fc2_w_col_data,<br/>swiglu_in, scales, split_points
CuDNN2-->>BwdFused: d_row/col_tensor (FC1 grad, FP8)<br/>sfd_row/col (MXFP8 scales)<br/>dprob_tensor (grad_scales)
BwdFused->>FC2GEMM: FC2 wgrad: grouped_fc2_x^T × grouped_fc2_dy
BwdFused->>FC2GEMM: FC1 dgrad: grouped_fc1_weight × grouped_fc1_dy
BwdFused->>FC2GEMM: FC1 wgrad: grouped_fc1_x^T × grouped_fc1_dy
FC2GEMM-->>BwdFused: grad_input, fc1_wgrad, fc2_wgrad
Reviews (5): Last reviewed commit: "Merge branch 'main' into fused_mxfp8_gro..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| import os | ||
| import functools | ||
| import math | ||
| from pickle import TRUE |
There was a problem hiding this comment.
from pickle import TRUE
TRUE is imported from the pickle module but is never used anywhere in this file. pickle.TRUE is an internal pickle opcode byte string (b'I01\n'), not a Python True value. This is almost certainly a leftover from development and should be removed.
| from pickle import TRUE | |
| from typing import Optional |
| return instance | ||
|
|
||
| self.with_gemm_swizzled_scales = with_gemm_swizzled_scales | ||
|
|
There was a problem hiding this comment.
return
The line self.with_gemm_swizzled_scales = with_gemm_swizzled_scales is placed after return instance in __new__, so it will never execute. The intent seems to be handled already in _initialize_storage_fields (where instance.with_gemm_swizzled_scales = with_gemm_swizzled_scales is correctly set), so this line is both unreachable and redundant.
| return instance | |
| self.with_gemm_swizzled_scales = with_gemm_swizzled_scales | |
| return instance | |
| def has_data(self) -> bool: |
| global global_alpha_tensor | ||
| alpha_tensor = self._mxfp8_alpha_tensor | ||
| norm_const_tensor = self._mxfp8_norm_const_tensor | ||
| if ( | ||
| alpha_tensor is None | ||
| or alpha_tensor.numel() != num_groups | ||
| or alpha_tensor.dtype != dtype | ||
| or alpha_tensor.device != device | ||
| ): | ||
| if global_alpha_tensor is None: | ||
| global_alpha_tensor = torch.ones(num_groups, dtype=dtype, device=device) | ||
| alpha_tensor = global_alpha_tensor | ||
| norm_const_tensor = alpha_tensor[:1] | ||
| self._mxfp8_alpha_tensor = alpha_tensor | ||
| self._mxfp8_norm_const_tensor = norm_const_tensor | ||
| elif ( | ||
| norm_const_tensor is None |
There was a problem hiding this comment.
global_alpha_tensor stale for multiple instances with different num_groups
In _get_kernel_constants, the module-level global_alpha_tensor is created only once (when it is None). If a second ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 instance is created with a different num_groups, the stale global (wrong size) is assigned to self._mxfp8_alpha_tensor. On every subsequent call, the condition alpha_tensor.numel() != num_groups will remain True, and the wrong-sized tensor from the global will keep being assigned—without ever being recreated.
Concretely:
- Instance A (4 groups): creates
global_alpha_tensorwith 4 elements ✓ - Instance B (8 groups):
global_alpha_tensor is NoneisFalse→ skips creation → assigns the 4-element tensor asalpha_tensor→ wrong size passed to the kernel.
The fix is to always recreate the global when the cached size/dtype/device doesn't match:
if (
global_alpha_tensor is None
or global_alpha_tensor.numel() != num_groups
or global_alpha_tensor.dtype != dtype
or global_alpha_tensor.device != device
):
global_alpha_tensor = torch.ones(num_groups, dtype=dtype, device=device)
alpha_tensor = global_alpha_tensorThe same issue exists in backward_grouped_mlp.py's _get_kernel_constants.
| assert hasattr(weight_param, "main_grad"), "MAIN GRAD NOT FOUND !!!!" | ||
| assert weight_param.main_grad is not None, "MAIN GRAD IS NONE !!!!" |
There was a problem hiding this comment.
These assert statements with !!!! in the messages look like debugging leftovers. While the assertions themselves are reasonable sanity checks, the message style is not production-appropriate. Consider using RuntimeError or a cleaner assert message:
| assert hasattr(weight_param, "main_grad"), "MAIN GRAD NOT FOUND !!!!" | |
| assert weight_param.main_grad is not None, "MAIN GRAD IS NONE !!!!" | |
| if self._accumulate_into_main_grad: | |
| if not hasattr(weight_param, "main_grad"): | |
| raise RuntimeError( | |
| f"Expected 'main_grad' attribute on weight parameter, but it was not found." | |
| ) | |
| if weight_param.main_grad is None: | |
| raise RuntimeError( | |
| f"'main_grad' on weight parameter is None." | |
| ) |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| norm_const_tensor=norm_const_tensor, | ||
| d_dtype=torch.float8_e4m3fn, | ||
| cd_major="n", | ||
| sf_vec_size=32, |
There was a problem hiding this comment.
sf_vec_size hardcoded instead of using the shared constant
The backward kernel call uses sf_vec_size=32 hardcoded, while the forward kernel in forward_grouped_mlp.py correctly uses sf_vec_size=MXFP8_BLOCK_SCALING_SIZE (which equals 32). Both values happen to be the same today, but for consistency and maintainability, the backward should also import and use MXFP8_BLOCK_SCALING_SIZE.
| sf_vec_size=32, | |
| sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, |
(add from ...constants import MXFP8_BLOCK_SCALING_SIZE to the imports)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
for more information, see https://pre-commit.ci
| dtype=torch.float32, | ||
| ) | ||
| fc1_weight.main_grad.fill_(value) | ||
| fc2_weight.main_grad.fill_(value) | ||
|
|
||
| def _collect_main_grads() -> tuple[torch.Tensor, torch.Tensor]: | ||
| if single_grouped_parameter: | ||
| fc1_main_grad = fc1.weight.main_grad.detach().clone() | ||
| fc2_main_grad = fc2.weight.main_grad.detach().clone() | ||
| else: | ||
| fc1_main_grad = torch.stack( | ||
| [ | ||
| getattr(fc1, f"weight{group_idx}").main_grad.detach().clone() | ||
| for group_idx in range(group_size) | ||
| ], | ||
| dim=0, | ||
| ) | ||
| fc2_main_grad = torch.stack( | ||
| [ | ||
| getattr(fc2, f"weight{group_idx}").main_grad.detach().clone() | ||
| for group_idx in range(group_size) |
There was a problem hiding this comment.
is_supported() causes test failure instead of skip
assert te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported() will raise AssertionError on any Blackwell system where the cudnn front-end package is installed for MXFP8 quantization but the specific grouped-GEMM-SwiGLU kernel is not available. In that scenario the test should be skipped, not failed. The maybe_skip_quantization guard only checks whether MXFP8 quantization in general is supported — it does not gate on the availability of the fused kernel.
Consider wrapping the fusion-presence block in a skip guard:
if (
quantization == "mxfp8"
and dtype in (torch.bfloat16, torch.float16)
and not bias
and glu_interleave_size == 32
):
if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported():
pytest.skip("MXFP8 fused grouped MLP kernel not available on this system")
forward_ops = module._module_groups[0]._forward_ops
...| window = [op] | ||
| else: | ||
| # Shift window if window doesn't match pattern | ||
| out.extend(window[:-2]) | ||
| window = window[-2:] | ||
|
|
||
| # Adjust window to expected size | ||
| out.extend(window[:-3]) | ||
| window = window[-3:] | ||
| while ops and len(window) < 3: | ||
| window.append(ops[0]) | ||
| ops = ops[1:] | ||
|
|
||
| # Return list of ops | ||
| out.extend(window) | ||
| return out | ||
|
|
||
|
|
||
| # Register fusion if available | ||
| if BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): | ||
| register_backward_fusion(fuse_backward_ops, prepend=True) |
There was a problem hiding this comment.
norm_const_tensor not refreshed when only alpha_tensor is stale
The backward's _get_kernel_constants is missing the elif guard that the forward version has. If self._mxfp8_alpha_tensor is somehow valid but self._mxfp8_norm_const_tensor is None (e.g., after a partial object state restore or future refactoring), the function returns None for norm_const_tensor, which will crash the kernel call.
The forward version correctly handles this with:
elif (
norm_const_tensor is None
or norm_const_tensor.numel() != 1
or norm_const_tensor.dtype != dtype
or norm_const_tensor.device != device
):
norm_const_tensor = alpha_tensor[:1]
self._mxfp8_norm_const_tensor = norm_const_tensorAdd the same defensive elif branch here for consistency and safety.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
| fc1_x_data = grouped_fc1_x.rowwise_data.view(in_shape[0], in_shape[1]) | ||
| fc1_x_data = fc1_x_data.view(dtype=torch.float8_e4m3fn) | ||
| fc1_x_data = fc1_x_data.unsqueeze(0).permute(1, 2, 0) | ||
| fc1_x_scales = grouped_fc1_x.scale_inv | ||
| fc1_x_scales = fc1_x_scales.view(dtype=torch.float8_e8m0fnu) | ||
| fc1_x_scales = fc1_x_scales.view( | ||
| 1, | ||
| in_shape[0] // 128, | ||
| in_shape[1] // 128, | ||
| 32, | ||
| 4, | ||
| 4, | ||
| ) | ||
| fc1_x_scales = fc1_x_scales.permute(3, 4, 1, 5, 2, 0) |
There was a problem hiding this comment.
The scale tensor view at lines 272–279 uses integer division in_shape[0] // 128 to reshape the MXFP8 scale buffer. If in_shape[0] (i.e., sum(split_sizes)) is not divisible by 128, the view shape product will not match the actual buffer size and either produce incorrect behavior (wrong permute dimensions) or a runtime error with a confusing message.
The constructor checks that in_features % 256 == 0 and out_features % 256 == 0, but nothing validates that the token dimension sum(split_sizes) is divisible by 128 (required by the MXFP8 block-scaling layout). A user passing split sizes like [64, 65] would hit this silently.
The same assumption appears in the backward pass at backward_grouped_mlp.py lines 243–250.
Consider adding a guard before the view:
if in_shape[0] % 128 != 0:
raise ValueError(
f"Total token count must be divisible by 128 for MXFP8 fused kernel, "
f"but got sum(split_sizes)={in_shape[0]}."
)
| bool with_gemm_swizzled = false; | ||
| if (py::hasattr(tensor, "_with_gemm_swizzled_scales")) { | ||
| with_gemm_swizzled = tensor.attr("_with_gemm_swizzled_scales").cast<bool>(); | ||
| } else if (py::hasattr(tensor, "with_gemm_swizzled_scales") && |
There was a problem hiding this comment.
Why do we need a second attribute name? Shouldn't we just use the _with_gemm_swizzled_scales attribute?
There was a problem hiding this comment.
This probably came from merging the changes. These file changes are not needed
There was a problem hiding this comment.
I think we should get rid of the "_" and just use with_gemm_swizzled_scales everywhere
| ) -> None: | ||
| super().__init__(name) | ||
|
|
||
| # Temporary for quick testing. |
ptrendx
left a comment
ptrendx
left a comment
There was a problem hiding this comment.
Not yet done with the full review, but cursory glance shows some leftover debugging code and some other random things that should be cleaned up.
| Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte); | ||
| Tensor cublas_ws("cublas_ws", std::vector<size_t>{cublas_ws_bytes}, DType::kByte); | ||
|
|
||
| nvte_grouped_gemm_with_discrete_out(grouped_A.get_handle(), |
There was a problem hiding this comment.
Not a fan of this name, but it was added in another PR, so not a problem here.
| getTensorShape(*tensor_offsets)); | ||
| } | ||
| nvte_set_grouped_tensor_swizzled_scales(out_cpp.data(), | ||
| static_cast<uint8_t>(with_gemm_swizzled_scales)); |
There was a problem hiding this comment.
set_with_gemm_swizzled_scales below already swizzles the scales. This is no longer needed. MR509 was not fully cleaned up/synced with main. I guess thats why the left over.
None of this file's changes are needed really.
There was a problem hiding this comment.
Agreed, please clean it up
| del warmup_out, warmup_x, warmup_probs, warmup_dy | ||
| gc.collect() |
There was a problem hiding this comment.
Do we expect users to do the same things when capturing this?
| * | ||
| * \param[in] tensor Grouped tensor. | ||
| * \param[in] val 1 if scales are swizzled, 0 otherwise. |
There was a problem hiding this comment.
This File change is also not needed. Got rid of this experimental API now.
| return 0; | ||
| } | ||
| const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); | ||
| return t.with_gemm_swizzled_scales ? 1 : 0; |
There was a problem hiding this comment.
Changes in this file can be reverted
| } // namespace | ||
|
|
||
| // MXFP8 support for grouped GEMM requires cuBLAS 13.3+ | ||
| #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 |
There was a problem hiding this comment.
Need to revert this changes
| * - scale_inv is stored in row-major per group. | ||
| * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. | ||
| * - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. | ||
| */ |
There was a problem hiding this comment.
Already added in previous PR. not needed
| dst.columnwise_scale_inv_offsets = src.columnwise_scale_inv_offsets | ||
| dst.logical_shape = src.logical_shape | ||
| dst.quantized_tensors = src.quantized_tensors | ||
| dst.with_gemm_swizzled_scales = src.with_gemm_swizzled_scales |
There was a problem hiding this comment.
There needs to be _ in here. Attribute for swizzled scales used _
|
|
||
| // MXFP8 support for grouped GEMM requires cuBLAS 13.3+ | ||
| #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 | ||
| #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130200 |
There was a problem hiding this comment.
This should not be changed - cublas 13.2 has issues with wgrad.
| } | ||
|
|
||
| TEST(SwizzleGroupedTestSuite, TestGroupedSwizzleMXFP8) { | ||
| performTestGroupedSwizzleMXFP8(3, 256, 256); |
There was a problem hiding this comment.
This is very limited coverage - what about the cases where
- we have different M or K?
- M/K are not nice numbers?
- Some of M/K are 0?
| quantizer = self.get_quantizer("forward", 1) | ||
|
|
||
| recipe = None if quantizer is None else quantizer._get_compatible_recipe() | ||
| if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): |
There was a problem hiding this comment.
Shouldn't this error out instead of silently returning?
| maybe_dequantize, | ||
| ) | ||
|
|
||
| global_alpha_tensor = None |
There was a problem hiding this comment.
Why do we have this as global? This effectively makes it impossible to run multiple instances of those ops in parallel without silent data corruption.
| maybe_dequantize, | ||
| ) | ||
|
|
||
| global_alpha_tensor = None |
There was a problem hiding this comment.
Same comment as in the other file. Why do we need to have this?
There was a problem hiding this comment.
To ensure that the alpha pointer is constant for cuda graphs. Without caching this tensor, the pointer changes and we get NaN in CG
| assert te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported() | ||
| assert te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported() |
There was a problem hiding this comment.
We should skip those tests instead I think? Or do we expect this to be a core requirement to have those fusions supported?
| # TODO(ksivaman): Proper support for meta device. | ||
| # We do not want to reset params later as it wipes off | ||
| # main_grad and related attributes. |
There was a problem hiding this comment.
Is this going to be done in this PR?
| * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. | ||
| * - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. | ||
| */ | ||
| void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, |
There was a problem hiding this comment.
I don't see the code of this function in this PR?
| instance.quantized_tensors = None | ||
| instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales | ||
|
|
||
| instance.with_gemm_swizzled_scales = with_gemm_swizzled_scales |
There was a problem hiding this comment.
same comment with @vthumbe1503 , but I think we should get rid of _with_gemm_swizzled_scales
| @staticmethod | ||
| def make_tensor_offsets(first_dims: torch.Tensor, logical_last_dim: int) -> torch.Tensor: | ||
| """Calculate GPU offsets from first dim splits.""" | ||
| return torch.cat( |
There was a problem hiding this comment.
This op can leverage the new fused kernel from @ksivaman
| else noop_cat([w._rowwise_scale_inv for w in grouped_fc1_weight]) | ||
| ) | ||
| fc1_w_scales = fc1_w_scales.view(dtype=torch.float8_e8m0fnu) | ||
| fc1_w_scales = fc1_w_scales.view( |
There was a problem hiding this comment.
why don't we use the grouped swizzle kernel here?
| fc1_w_data = ( | ||
| grouped_fc1_weight.rowwise_data | ||
| if fc1_op.single_grouped_parameter | ||
| else noop_cat([w._rowwise_data for w in grouped_fc1_weight]) |
There was a problem hiding this comment.
how about the discreate weight version?
| fc2_in_col_data = fc2_in_col_data.permute(2, 0, 1) | ||
| fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]).contiguous() | ||
| fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] | ||
| fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) |
There was a problem hiding this comment.
why do we choose to swizzle here and then call make_grouped_tensor_from_buffers and mark with_gemm_swizzled_scales=True, we should be able to leverage grouped swizzle kernel from @vthumbe1503 ?
| fc2_dy_data = fc2_dy_data.unsqueeze(0).permute(1, 2, 0) | ||
| fc2_dy_scales = grouped_fc2_dy.scale_inv | ||
| fc2_dy_scales = fc2_dy_scales.view(dtype=torch.float8_e8m0fnu) | ||
| fc2_dy_scales = fc2_dy_scales.view( |
There was a problem hiding this comment.
same ^ we should be able to use grouped quantize swizzle fusion here
| logical_last_dim=fc1_weight_shape[1], | ||
| ) | ||
|
|
||
| general_grouped_gemm_for_grouped_tensor( |
There was a problem hiding this comment.
For FC1 DGRAD, would it be better if we use contiguous layout gemm with activation fusion disabled?
| weight_param.dtype, | ||
| zero=getattr(weight_param, "zero_out_wgrad", False), | ||
| ) | ||
| else: |
There was a problem hiding this comment.
Nit: this if else is pretty long, please add some comments to the else
4 participants
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: