[Common] Improved performance of Group MXFP8 kernels

[Oleg-Goncharov]

Description

This PR optimizes grouped MXFP8 quantization by introducing layout-specific kernel configurations and work scheduling strategies.

Groups with uniform shapes or a common last dimension use direct CTA-to-chunk mapping, while layouts requiring per-tensor TMA descriptors use per-tensor persistent scheduling. The kernel processes larger chunks through multiple stages, reducing scheduling and tensor-mapping overhead for large grouped workloads.

The change also ensures that dynamically updated TMA descriptors are acquired before use, handles graph-safe logical capacity without out-of-bounds accesses, and preserves correct dBias workspace behavior across multi-chunk execution.

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

• Extend grouped TMA storage with per-tensor row, column, and offset metadata.
• Introduce layout-specific grouped MXFP8 kernel configurations and compile-time traits.
• Use direct CTA-to-chunk mapping for SAME_BOTH_DIMS and VARYING_FIRST_DIM layouts.
• Use per-tensor persistent scheduling for VARYING_LAST_DIM and VARYING_BOTH_DIMS layouts.
• Process 128x128 and 128x256 chunks through multiple X/Y stages.
• Acquire dynamically updated TMA descriptors before they are consumed by the quantization kernel.
• Avoid out-of-bounds accesses when graph-safe logical capacity exceeds the active tensor range.
• Preserve correct rowwise, columnwise, and bidirectional scaling behavior.
• Preserve correct dBias accumulation and workspace reduction for multi-chunk execution.
• Add focused regression coverage for all grouped shape representations, multi-chunk dBias/dAct execution, and the maximum of 64 dynamic TMA descriptors.
• Keep all existing dispatch paths, operator tests, and CUDA architecture targets enabled.

Runtime Speedup Factors on GB300

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

The PR restructures grouped MXFP8 quantization around layout-specific kernel traits and scheduling strategies.

• Uses direct chunk mapping for uniform/common-last-dimension layouts and tensor-local persistent scheduling for fully dynamic descriptors.
• Adds multi-stage processing for wider chunks and per-tensor dimension/offset metadata alongside TMA descriptors.
• Expands C++ regression coverage for multi-chunk execution, dBias paths, grouped shape representations, and the 64-descriptor boundary.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Introduces layout-specific traits, direct and persistent work mappers, multi-stage chunk processing, descriptor acquisition, and revised dBias workspace handling.
transformer_engine/common/cast/core/grouped_tma.cuh	Extends global grouped TMA storage with per-tensor rows, columns, and offsets populated during descriptor updates.
tests/cpp/operator/test_cast_mxfp8_grouped.cu	Adds focused grouped MXFP8 coverage for multi-chunk scaling and dBias execution plus the maximum dynamic descriptor count.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Grouped MXFP8 input] --> B{Shape representation}
  B -->|SAME_BOTH_DIMS| C[Direct CTA-to-chunk mapping]
  B -->|VARYING_FIRST_DIM| C
  B -->|VARYING_LAST_DIM| D[Per-tensor persistent scheduling]
  B -->|VARYING_BOTH_DIMS| D
  D --> E[Acquire dynamic TMA descriptor]
  C --> F[Process staged 128-row tiles]
  E --> F
  F --> G[Write rowwise and/or columnwise output]
  F --> H[Accumulate dBias workspace when enabled]
  H --> I[Grouped dBias reduction]

Loading

_{Reviews (5): Last reviewed commit: "Fix" | Re-trigger Greptile}

[Oleg-Goncharov]

/te-ci