Optimize single-group grouped MLP quantization

[sraman-rgb]

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

[sraman-rgb]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR extends the single-group quantization shortcut (previously NVFP4-only) to also cover MXFP8, and avoids a GPU prefix-sum kernel for split-metadata when num_groups==1 by computing the metadata directly and caching it on the module instance.

• Quantizer generalization: _wrap_single_nvfp4_as_grouped, _nvfp4_single_tensor_from_grouped, and _nvfp4_single_group_wgrad_gemm are renamed and broadened to handle both MXFP8Quantizer and NVFP4Quantizer; MXFP8-specific data-shape logic is added in the reconstruction path.
• Split-metadata cache: A new _single_group_split_metadata helper computes split_sizes/split_points/tensor_offsets directly (no GPU cumsum kernel), and the results are memoized in self._single_group_split_metadata_cache keyed by (device, num_tokens, fc1/fc2 feature sizes).

Confidence Score: 4/5

The change is functionally correct for the single-group MXFP8 path and the numerical equivalence with the original GPU prefix-sum is exact. The two findings are non-blocking quality concerns that would not cause incorrect results in normal use.

The metadata precomputation is mathematically equivalent to the original tex.splits_to_offsets_multi call for num_groups==1, and the MXFP8 tensor reconstruction path correctly handles shape differences versus NVFP4. The instance-level cache introduces an unbounded dict that accumulates entries for every distinct token count seen, which can slowly leak GPU memory in variable-length workloads. The quantizer guard in the wgrad path uses an asymmetric isinstance pattern that could silently skip the optimization if subclasses of MXFP8Quantizer are involved. Neither issue causes wrong outputs today.

transformer_engine/pytorch/ops/fused/grouped_mlp.py — specifically the cache management logic around line 997 and the quantizer guard at line 667.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Extends the single-group NVFP4 quantization shortcut to also cover MXFP8; renames the associated helpers and adds an instance-level cache to avoid a GPU prefix-sum kernel for the split-metadata computation when num_groups==1.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[fuser_forward] --> B{num_groups == 1?}
    B -- Yes --> C[Check metadata cache]
    C -- Hit --> D[Return cached tensors]
    C -- Miss --> E[_single_group_split_metadata\ncompute directly, no GPU cumsum]
    E --> F[Store in cache]
    F --> D
    B -- No --> G[tex.splits_to_offsets_multi\nGPU prefix-sum kernel]
    D --> H{quantizer type?}
    G --> H
    H -- MXFP8 or NVFP4 --> I[_group_quantize_for_grouped_mlp\nsingle-group path]
    H -- Other --> J[tex.group_quantize]
    I --> K{num_groups==1 AND\nMXFP8 or NVFP4?}
    K -- Yes --> L[tex.quantize +\n_wrap_single_quantized_as_grouped]
    K -- No --> J
    L --> M[GroupedTensor]
    J --> M
    M --> N{wgrad needed?\nnum_groups==1?}
    N -- Yes, MXFP8 or NVFP4 --> O[_single_group_wgrad_gemm\nregular GEMM]
    N -- No or multi-group --> P[general_grouped_gemm_for_grouped_tensor]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[fuser_forward] --> B{num_groups == 1?}
    B -- Yes --> C[Check metadata cache]
    C -- Hit --> D[Return cached tensors]
    C -- Miss --> E[_single_group_split_metadata\ncompute directly, no GPU cumsum]
    E --> F[Store in cache]
    F --> D
    B -- No --> G[tex.splits_to_offsets_multi\nGPU prefix-sum kernel]
    D --> H{quantizer type?}
    G --> H
    H -- MXFP8 or NVFP4 --> I[_group_quantize_for_grouped_mlp\nsingle-group path]
    H -- Other --> J[tex.group_quantize]
    I --> K{num_groups==1 AND\nMXFP8 or NVFP4?}
    K -- Yes --> L[tex.quantize +\n_wrap_single_quantized_as_grouped]
    K -- No --> J
    L --> M[GroupedTensor]
    J --> M
    M --> N{wgrad needed?\nnum_groups==1?}
    N -- Yes, MXFP8 or NVFP4 --> O[_single_group_wgrad_gemm\nregular GEMM]
    N -- No or multi-group --> P[general_grouped_gemm_for_grouped_tensor]

Loading

_{Reviews (1): Last reviewed commit: "Optimize single-group grouped MLP quanti..." | Re-trigger Greptile}

[greptile-apps]

Unbounded cache growth for variable token counts

_single_group_split_metadata_cache is a plain dict with no eviction policy. The key includes in_shape[0] (the token count), which changes on every call with a different batch/sequence length. Even though tokens must be multiples of 128, a long-running inference service with variable-length inputs will accumulate one cache entry per unique token count encountered — each entry holding six GPU tensors. Over time this is a steady memory leak.

[greptile-apps]

The quantizer-type guard for the single-group wgrad optimization uses isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__). This is asymmetric: if grouped_x carries a subclass of MXFP8Quantizer, the check requires grouped_dy to be that same subclass (or a further subclass), not just a compatible MXFP8Quantizer. The symmetrical pattern used everywhere else in the file is more robust.

Suggested change

	and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer))
	and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__)
	and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer))
	and type(grouped_x.quantizer) == type(grouped_dy.quantizer)

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!