[Pytorch] Add support for row-wise quanted input for grouped gemm

[YangFei1990]

Description

To support the NCCL EP FP8 dispatch, this PR introduce a capability to allow grouped gemm to take input that is already quantized in row-wise.

• For forward, we directly pass the row-wise quantized for forward grouped gemm and dequant + requant with col-wise to save for backward.
• For backward, we pass the row-wise quantized for dgrad computation and dequant + requant with col-wise for wgrad

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

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 adds row-wise prequantized MXFP8 input and gradient support for grouped GEMM.

• Adds a native in-place grouped requantization helper that swizzles rowwise scales and optionally rebuilds columnwise storage.
• Integrates the helper into basic grouped linear and fused grouped MLP forward/backward paths.
• Extends grouped-scale swizzling metadata and adds CUDA graph, numerical, shape, and edge-case tests.

Confidence Score: 4/5

The PR does not yet appear safe to merge because the previously reported frozen-weight backward and fused-forward failures remain outstanding.

Row-wise prequantized tensors still encounter unsupported frozen-weight states: the fused forward can save rowwise-only activation storage through a helper that requires columnwise data, while the ordinary and scaled-bias backward scenarios have not been made valid end to end.

Files Needing Attention: transformer_engine/pytorch/ops/basic/grouped_linear.py; transformer_engine/pytorch/ops/fused/grouped_mlp.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds validation, dequantization, rowwise-scale swizzling, and optional columnwise reconstruction for grouped MXFP8 tensors.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Propagates varying per-tensor dimensions into grouped-scale swizzle descriptors.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds prequantized grouped-tensor handling to grouped linear forward and backward.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds prequantized grouped-tensor handling to fused grouped MLP forward and backward.
tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py	Adds direct coverage for requantization, graph-safe layouts, zero-token groups, paged stashing, and invalid states.
tests/pytorch/test_grouped_mlp.py	Adds numerical coverage for prequantized inputs and gradients in grouped linear and fused grouped MLP operations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Row-wise MXFP8 grouped tensor] --> B[group_requantize_inplace]
  B --> C[Swizzle row-wise scales]
  B --> D{Columnwise needed?}
  D -->|Yes| E[Dequantize and requantize columnwise]
  D -->|No| F[Keep row-wise storage only]
  C --> G[Grouped GEMM]
  E --> G
  F --> G

Loading

_{Reviews (15): Last reviewed commit: "Merge branch 'main' into mxfp8_input_gro..." | Re-trigger Greptile}

[YangFei1990]

/te-ci pytorch

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[YangFei1990]

/te-ci L1 pytorch

[vthumbe1503]

Suggested change

	py::object group_requantize_columnwise_and_swizzle_rowwise_(
	py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
	std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
	bool return_dequantized) {
	py::object group_requantize(
	py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
	std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
	bool return_dequantized) {

I suggest to make this function generic.

The intent of this function as follows

1. If quantizer.rowwise_usage is true and grouped_x has rowwise data/scales, then grouped_swizzle
2. If quantizer.columnwise_usage is true and grouped_x has columnwise_data/scales, then grouped_swizzle
3. If quantizer.rowwise_usage is true and grouped_x doesnt have rowwise_data/scales, then dequant + requant with swizzle fusion
4. If quantizer.columnwise_usage is true and grouped_zx doesnt have columnwise_data/scales, then dequant + requant with swizzle fusion

3 --> is a unrealistic case and we can throw an error in that case as well

Also, we should make sure to gather all swizzling directions once(rowwise, colwise or both rowwise/colwise) and call grouped_swizzle once.

[YangFei1990]

I'm not sure if we should make this function generic, the purpose of it is to handle the mxfp8 output from dispatch, and we are also talking about make this a fused kernel with dispatch's permutation, that is the major reason to refactor it from python into c++ layer. cc @phu0ngng

[vthumbe1503]

So the thing is, for now it is now serving the purpose of
mxfp8 dispatch --> mxfp8 gemm ready input

But it can be this as well in the future
mxfp8 dispatch --> nvfp4 gemm ready input

Now if your intent is to retain the quantization between input and output
like mxfp8 dispatch --> mxfp8 gemm ready input
or nvfp4 dispatch --> nvfp4 gemm ready input
I am ok with that as well. But that should mentioned in the comment.

group_requantize_columnwise_and_swizzle_rowwise_ --> This as a name seems too specific and encodes too much information in the name of the function which isnt needed. Quantizer already has the information with optimize_for_gemm=True/False which tells whether to swizzle or not and which direction to swizzle. So just keeping the name as "group_requantize" should suffice.

As far as fused kernels are concerned, the special case where fusion is available, only that case can be replaced with the fused kernel

In general, I want to differentiate between use-case of a function and intent of the function. Use-case is dispatch --> gemm-ready handling. But the Intent of the function is --> requantize to the best of the ability. And that means if you have already quantized data along a direction, then to make gemm ready you just need to swizzle it. However if you dont have quantized data along a direction, then to make it gemm ready dequant + quant + swizzle(based on quantizer config)

[vthumbe1503]

Currently in your function 1 and 4 that I mentioned above is handled. We can throw error for 2 and 3 as well based on quantizer config and grouped_x.quantizer config. So I am not asking to implement the generic function. But keep the interface and name of the function generic. So that in future if we want to implement new feature we dont have to change the name and signature of the function

[YangFei1990]

I see. Renamed to group_requantize and added one additional assert. Currently the function only takes rowwise_data and column quantizer, but we can extend in the future.

[timmoon10]

Adding a special case for MXFP8 input is quite hacky. If we try to understand the code, this entire code section is converting the input (previously assumed to be BF16, but now may be grouped MXFP8) to the format needed by GGEMM. We can lift all of this into a clean helper function:

def _convert_input_to_grouped_tensor(input_, ...):

    # Do nothing if input is already in expected format
    if input_is_in_expected_format:
        return input_

    # Fast requantize impls
    if input_is_mxfp8 and compute_is_mxfp8:
        return tex.group_requantize...(x)
    if fancy_future_fused_impl_is_available:
        return tex.fancy_future_fused_impl(...)

    # Fallback: dequantize if needed and group quantize
    x = maybe_dequantize(input_)
    grouped_x = tex.group_quantize(x, ...)
    return grouped_x

[YangFei1990]

There are four sites that these checks are involved, fwd_grouped_linear (this one), bwd_grouped_linear, fwd_grouped_mlp, bwd_grouped_mlp, the paths have common parts, but also differ a lot, e.g. quantize entry point, handle of dbias, NVFP4 path, fallback paths, I feel if we create a single helper function to handle all cases, the function itself might be complicated with a lot of conditional branches. What are your recommendations?

[zhongbozhu]

It will be better if we let the group_requantize_columnwise_and_swizzle_rowwise_ API accepts regular 2D inputs or dY, convert to grouped tensor in C++ for less CPU overhead.

[vthumbe1503]

Agreed that It would reduce CPU overheads, however we would loose generality of the function. If we want to use it for NVFP4 later, then we would need to pass amax as well. I think for now, it is better to consume grouped tensor

[zhongbozhu]

same comment, it would be better if we just pass in plain torch tensor into this API (please refer to our API design for grouped_quantize, and the grouped tensor descriptor is created in C++ for less overhead

[YangFei1990]

/te-ci L1 pytorch

[YangFei1990]

/te-ci L1 pytorch

[vthumbe1503]

LGTM. Pending CI.

[YangFei1990]

/te-ci L1 pytorch