Add MXFP8 attention unit test with linear and rope layers

[layalir]

Add a DSv3-shaped MXFP8 attention unit test covering the training path:

• Adds MLA RoPE utilities for the DSv3 671B attention shape.
• Adds an end-to-end MXFP8 path: Linear(QKV) -> MLA RoPE -> DotProductAttention -> Linear(out).
• Exercises MXFP8 forward and backward through TE's real DotProductAttention wrapper.
• Runs BF16 reference comparison by default.
• Runs the performance benchmark by default and reports fprop and bprop timing separately from the same benchmark collection.

Validation

Local checks:

• python -m py_compile tests/pytorch/attention/test_linear_mxfp8_attention.py tests/pytorch/attention/mla_rope_utils.py
• git diff --check

GB300 dlcluster validation:

• Job: 1062811
• GPU: NVIDIA GB300
• CUDA capability: (10, 3)
• cuDNN: (9, 21, 1)
• MXFP8 available: (True, '')
• Command: python -m pytest tests/pytorch/attention/test_linear_mxfp8_attention.py -v -s
• Result: 3 passed

Perf output:

[PERF] b=1 s=4096:
  BF16 fprop:  7.219 ms  (567397 tok/s)
  BF16 bprop:  15.179 ms  (269844 tok/s)
  MXFP8 fprop: 4.718 ms  (868181 tok/s)
  MXFP8 bprop: 9.215 ms  (444492 tok/s)
  Fprop speedup: 1.53x
  Bprop speedup: 1.65x

[greptile-apps]

Greptile Summary

This PR adds a DSv3-shaped MXFP8 attention unit test covering the full training path (Linear(QKV) → MLA RoPE → DotProductAttention → Linear(out)) plus the companion Triton/PyTorch MLA RoPE utilities. The implementation is well-structured and validated on GB300 hardware.

• mla_rope_utils.py: New Triton kernels (rotary_fwd/bwd_q_kernel, rotary_fwd/bwd_kv_kernel) implement the interleaved-to-NeoX RoPE rotation for DSv3 MLA; a pure-PyTorch fallback is provided when Triton is unavailable. The backward correctly reconstructs the key-positional-embedding gradient by accumulating across all heads and routing it only to head 0's rope slice.
• test_linear_mxfp8_attention.py: Three test cases exercise accuracy (loose FP8-appropriate tolerances), gradient flow (NaN/Inf checks), and throughput (speedup assertion vs BF16 reference), parameterised over the fixed b=1, s=4096 DSv3 shape.

Confidence Score: 5/5

This PR is safe to merge — it adds test-only files with no changes to library source code, and the Triton kernel logic is correct for the fixed 128-head DSv3 configuration it targets.

Both files are purely test infrastructure; no production library code is modified. The Triton RoPE kernels are logically correct for the specific 128-head count used in the test (128 is divisible by every autotune BLOCK_H candidate, so the partial-block boundary condition never triggers). The overall forward/backward autograd wiring is sound, and the benchmark and accuracy thresholds are appropriate for FP8 arithmetic.

No files require special attention for merging; both are new test-only additions with no impact on library users.

Important Files Changed

Filename	Overview
tests/pytorch/attention/mla_rope_utils.py	New Triton implementation of MLA RoPE for DSv3 671B attention, providing both forward and backward kernels with a pure-PyTorch fallback; the partial-head-block mask bug (already flagged) is the main latent risk, but is harmless for the fixed 128-head DSv3 shape.
tests/pytorch/attention/test_linear_mxfp8_attention.py	End-to-end MXFP8 accuracy, backward, and benchmark tests covering the QKV→RoPE→DPA→out pipeline; three separate fp8_autocast contexts (already flagged) and warmup-backward cache-invalidation (already flagged) are the open concerns.

Sequence Diagram

sequenceDiagram
    participant X as Input x [s,b,H]
    participant QKV as te.Linear QKV
    participant SPL as _split_qkv
    participant RoPE as apply_mla_rope (Triton/PyTorch)
    participant DPA as te.DotProductAttention
    participant OUT as te.Linear out
    participant Y as Output [s,b,H]

    Note over QKV,OUT: fp8_autocast (3 separate contexts)

    X->>QKV: [s,b,H] fp8_autocast 1
    QKV-->>SPL: "qkv [s,b,2*H_QK+H_V]"
    SPL-->>RoPE: q[s,b,h,192], k[s,b,h,192], v[s,b,h,128]
    Note over RoPE: outside fp8_autocast
    RoPE-->>DPA: q_rot, k_rot, v fp8_autocast 2
    DPA-->>OUT: attn_out[s,b,h,128] fp8_autocast 3
    OUT-->>Y: [s,b,H]

    Note over X,Y: Backward (reverse order)
    Y->>OUT: d_attn_out
    OUT->>DPA: dq, dk, dv
    DPA->>RoPE: _MLARoPETriton.backward
    Note over RoPE: rotary_bwd_q_kernel in-place on dq, rotary_bwd_kv_kernel accumulates dEMB
    RoPE->>SPL: dq_in, dk_in, dv_in
    SPL->>QKV: d_qkv slice backward
    QKV->>X: d_x

Loading

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

[greptile-apps]

Mask ignores block offset for partial last head-block

The bound check x_off < head_num * stride_x_nheads compares only the intra-block head index i against the total head count, but the pointer Q has already been advanced by pid_head * BLOCK_H * stride_x_nheads. For the last block when head_num % BLOCK_H != 0, the absolute head index pid_head * BLOCK_H + i may exceed head_num while the mask still evaluates to True, causing an out-of-bounds load/store. The same pattern appears in rotary_bwd_q_kernel (line 178) and in the accumulation loop inside rotary_bwd_kv_kernel (line 332). For DSv3's 128 heads all BLOCK_H candidates (1–128) evenly divide 128, so the current test is unaffected, but any future caller with a non-aligned head count would silently corrupt memory.

[greptile-apps]

Three separate fp8_autocast scopes may reset per-layer FP8 statistics

The forward pipeline is split into three independent fp8_autocast contexts (QKV linear, DPA, out linear). For current MXFP8 block scaling the effect is benign because scales are computed per-block and don't depend on cross-layer statistics, but TE's FP8GlobalStateManager maintains per-forward-pass amax history used by some recipes. Exiting and re-entering the context between layers means each layer is treated as a separate forward pass for bookkeeping purposes, so any inter-layer scale propagation is lost. A single surrounding fp8_autocast context covering all three layers would match the standard training usage and would more faithfully exercise the end-to-end path.

[greptile-apps]

Warmup backward may invalidate the is_first_microbatch weight cache before timed iterations

The weight cache is seeded with is_first_microbatch=True inside torch.no_grad(), but the 10 warmup iterations inside _benchmark_training_step call backward() on a new computation graph. TE invalidates or refreshes the cached quantized weight buffer after a backward pass (the cache is keyed to the current "microbatch"). Consequently the timed iterations that pass is_first_microbatch=False may miss the cache on every call and silently fall back to per-iteration weight quantization, making the benchmark measure a different workload than described in the docstring.

[cyanguwa]

Thanks for the contribution! Could you please:

• fix the DCO (there are instructions in the DCO link)
• address Greptile comments
• add test_linear_mxfp8_attention.py to qa/L0_pytorch_unittest/test.sh, similar to this

[cyanguwa]

I think we should check "if not fused_attn_supported_fp8" here. FlashAttention doesn't support MXFP8, so it's kind of relevant here.

[cyanguwa]

Does this work for MXFP8 tensors? Should we do qkv.to() first before splitting?

[cyanguwa]

These tolerances are a bit too high, aren't they? For regular MXFP8 tests, I've been using atol = 5e-1, rtol = 5e-2. Would they work for these tests?