[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 Factor 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 launch configurations and multi-stage chunk processing.

• Uses direct chunk mapping for uniform/common-last-dimension layouts and per-tensor persistent scheduling for descriptor-backed layouts.
• Adds device-resident tensor dimensions and offsets alongside dynamic TMA descriptors.
• Extends numerical regression coverage for multi-chunk execution, dBias/dAct, all shape representations, and 64 descriptors.

Confidence Score: 5/5

The PR appears safe to merge because no concrete changed-code failure remains after checking mapping boundaries, graph-safe capacity, descriptor usage, and reachable dBias configurations.

The new schedulers keep accesses within capacity-sized buffers, the reachable multi-stage paths preserve tensor and workspace ownership, and the suspected shared-descriptor concurrency issue was not introduced or materially worsened by this change.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Introduces layout-specific traits, direct and persistent work mapping, multi-axis staging, descriptor acquisition, and revised dBias workspace handling; no actionable defect was established.
transformer_engine/common/cast/core/grouped_tma.cuh	Extends grouped TMA device storage with per-tensor rows, columns, and offsets populated during descriptor updates.
tests/cpp/operator/test_cast_mxfp8_grouped.cu	Adds focused numerical coverage for multi-chunk layouts, dBias/dAct execution, and the maximum supported descriptor count.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Grouped MXFP8 request] --> B{Shape representation}
  B -->|SAME_BOTH_DIMS| C[Direct tensor-local CTA mapping]
  B -->|VARYING_FIRST_DIM| D[Direct vertically stacked CTA mapping]
  B -->|VARYING_LAST_DIM| E[Update per-tensor TMA descriptors and metadata]
  B -->|VARYING_BOTH_DIMS| E
  E --> F[Per-tensor persistent scheduling]
  C --> G[Process 32x128 stages]
  D --> G
  F --> G
  G --> H[Write quantized outputs and scales]
  H --> I{dBias enabled?}
  I -->|Yes| J[Write workspace and reduce]
  I -->|No| K[Complete]
  J --> K

Loading

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