GEMM + Swiglu fused Grouped MLP for MXFP8

[ksivaman]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR introduces a fused GEMM + SwiGLU kernel for MXFP8 Grouped MLP on SM100 (Blackwell) GPUs, using CuTe DSL kernels from cuDNN front-end. It adds ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 and BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8 fused ops, supporting both stacked-weight and discrete-weight variants, along with GroupedTensor / GroupedTensorStorage extensions for swizzled-scale handling.

Many correctness concerns (non-contiguous .view() on quant kernel outputs, stale grouped_fc1_dy scale layout, mark_grouped_tensor assertion on frozen weights, overwrite_main_grad + single_grouped_weight silent wgrad drop, and meta-device AttributeError) remain open from earlier review rounds and are not yet fixed in the current head.

Confidence Score: 2/5

Not safe to merge — multiple prior-round P1 correctness issues remain present in the current head with no fix.

Several confirmed runtime-crash or silent-wrong-result bugs flagged in earlier review rounds remain unaddressed: (1) fc1_dgrad_kernel_out[d_tensor].view(in_shape) will raise RuntimeError because the quant kernel output requires .permute(2,0,1) before .view(); (2) grouped_fc1_dy data fields are 2-D and scales are unpermuted, producing wrong wgrad; (3) mark_grouped_tensor hard-asserts columnwise_data is not None, crashing on frozen-weight + grad-input; (4) overwrite_main_grad=True + single_grouped_weight=True silently writes wgrad to a scratch buffer, never updating main_grad.

transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py (lines 504-516, 611) and transformer_engine/pytorch/utils.py (mark_grouped_tensor) need the most attention before merge.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	New file implementing the fused MXFP8 GEMM+SwiGLU forward pass; several correctness issues remain open (non-contiguous .view() on kernel outputs, unconditional bias_tensor kwarg, c_dtype hardcoded to bfloat16).
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py	New file implementing the fused MXFP8 dGLU + dgrad backward pass; multiple correctness issues remain open from prior rounds (non-contiguous d_tensor .view(), stale grouped_fc1_dy scale layout, overwrite_main_grad dropping wgrad).
transformer_engine/pytorch/ops/_common.py	Adds fuse_grouped_mlp_ops sliding-window fusion helper and validate_grouped_mlp_dims; logic is correct but bias guard only blocks fusion when has_bias=True, leaving a TypeError path for has_bias=False + missing kwarg.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds single_grouped_weight/bias support, make_grouped_weights, and delay_wgrad hooks; meta-device + single_grouped_weight + delay_wgrad AttributeError remains open.
transformer_engine/pytorch/utils.py	Adds get_cached_ones_tensor and mark_grouped_tensor helpers; the unconditional assertion in mark_grouped_tensor still fires when columnwise_data is None (frozen-weight + grad-input scenario).
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Adds with_gemm_swizzled_scales flag and extended constructor; unreachable line after return in new was flagged and fixed, storage now correctly initialised via _initialize_storage_fields.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	New C++ helpers for per-tensor and grouped scale swizzling; grouped_swizzle_for_gemm correctly operates on a Python-side copy and nullifies the unused scale direction.
transformer_engine/pytorch/csrc/extensions/utils.cpp	New file exposing get_device_pointer_for_data_and_scales for the discrete-weight kernel path; swizzle-and-pointer-gather logic looks correct.
tests/pytorch/test_fusible_ops.py	Adds tests for the new fused MLP ops; is_supported() check is now properly guarded with pytest.skip rather than a hard assert (prior concern addressed).

Sequence Diagram

sequenceDiagram
    participant Input
    participant FC1_GLU as FC1 GLU Kernel
    participant FC2_QUANT as FC2 QUANT Kernel
    participant Output
    participant FC2_DGLU as FC2 dGLU Kernel
    participant FC1_DGRAD as FC1 dGrad Kernel

    Note over Input,Output: FORWARD PASS
    Input->>FC1_GLU: a_tensor (MXFP8), b_tensor (weight), prob_tensor (scales), bias_tensor
    FC1_GLU-->>FC2_QUANT: d_tensor (MXFP8 act), sfd_row_tensor
    FC1_GLU-->>FC1_GLU: c_tensor (swiglu_in bf16), d_col_tensor, sfd_col_tensor
    FC2_QUANT-->>Output: d_tensor permute(2,0,1).view()

    Note over Input,Output: BACKWARD PASS
    Output->>FC2_DGLU: a_tensor (dy MXFP8), c_tensor (swiglu_in), b_tensor (FC2 weight col)
    FC2_DGLU-->>FC1_DGRAD: d_row_tensor (fc1_dy), sfd_row_tensor
    FC2_DGLU-->>FC2_DGLU: dprob_tensor (grad_scales), dbias_tensor
    FC1_DGRAD-->>Input: d_tensor .view() needs permute(2,0,1) first
    FC2_DGLU-->>FC2_DGLU: wgrad via general_grouped_gemm
    FC1_DGRAD-->>FC1_DGRAD: wgrad via general_grouped_gemm

Loading

_{Reviews (35): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

No validation that total token count is divisible by 128

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]}."
    )

[ptrendx]

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.

[ptrendx]

Not a fan of this name, but it was added in another PR, so not a problem here.

[greptile-apps]

AttributeError when device="meta" + single_grouped_weight=True + delay_wgrad_compute=True

When device.type == "meta", line 168 skips reset_parameters(), so make_grouped_weights() is never called and self.weight is never registered. _apply_delay_wgrad_param_hooks() is then called unconditionally at line 173 and immediately accesses self.weight (line 180) → AttributeError at construction time.

This breaks any framework (e.g. Megatron-LM) that initialises models on device="meta" with delayed wgrad and single grouped weights. A simple guard fixes it:

Suggested change

	self._apply_delay_wgrad_param_hooks()
	def _apply_delay_wgrad_param_hooks(self) -> None:
	"""Set ``skip_backward_post_hook`` on weights when delaying wgrad (bias uses main backward)."""
	if not self.wgrad_store.delay_wgrad_compute():
	return
	if self.single_grouped_weight:
	self.weight.skip_backward_post_hook = True
	else:
	for group_idx in range(self.num_groups):
	getattr(self, f"weight{group_idx}").skip_backward_post_hook = True
	if device.type != "meta":
	self._apply_delay_wgrad_param_hooks()

pre_first_fuser_forward will call reset_parameters() (and re-invoke _apply_delay_wgrad_param_hooks) once real device params are materialised.

[vthumbe1503]

/te-ci L1 pytorch

[vthumbe1503]

/te-ci L1 pytorch

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

grouped_fc1_dy data tensors not flattened — inconsistency with forward may produce incorrect wgrad

data=fc1_dy_row_data and columnwise_data=fc1_dy_col_data are passed as 2-D tensors (after .view(out_shape[0], fc1_weight_shape[0])), while GroupedTensorStorage expects 1-D flattened buffers (documented at grouped_tensor_storage.py line 44: "ALL data fields are stored as 1D flattened arrays"). The forward's equivalent construction for grouped_fc2_x explicitly flattens every field:

# forward_grouped_mlp.py – what the forward does
data=fc2_in_row_data.reshape(-1),
columnwise_data=fc2_in_col_data.reshape(-1),
scale_inv=fc2_in_row_scale.reshape(-1),
columnwise_scale_inv=fc2_in_col_scale.reshape(-1),

The backward omits .reshape(-1) for both data tensors and omits the required permute(5, 2, 4, 0, 1, 3) + reshape(-1) for the scale tensors. When general_grouped_gemm_for_grouped_tensor indexes into the per-group data via element-level offsets it may read from the wrong memory locations, producing silently wrong FC1 weight gradients.

[greptile-apps]

overwrite_main_grad=True + single_grouped_weight=True silently drops wgrad into a scratch buffer

When weight_param.overwrite_main_grad is True, accumulate_into_main_grad is set to False (line 97). Because the if accumulate_into_main_grad: branch is skipped, grouped_wgrad remains None and the fallback at line 107 allocates a new scratch buffer entirely unrelated to main_grad. The GEMM writes the weight gradient into this temporary buffer, which is then discarded — main_grad is never updated.

Compare with the single_grouped_weight=False path (lines 116–128): w_list[idx] = wp.main_grad is set unconditionally before the accumulate_into_main_grad determination, so the GEMM always targets main_grad regardless of overwrite_main_grad.

The fix is to populate grouped_wgrad from main_grad before computing accumulate_into_main_grad:

# single_grouped_weight path, before the accumulate flag computation:
grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data(
    num_tensors=num_groups,
    tensor_shape=weight_shape,
    rowwise_data=main_grad,
    dtype=main_grad.dtype,
)
accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False)

[vthumbe1503]

CI is passing now.

Comments have been addressed and CI is green now.

[vthumbe1503]

/te-ci L1 pytorch