[None][perf] cuDNN FP8 attention path for MiniMax-H3 on SM107 under Ulysses by karljang · Pull Request #21485 · NVIDIA/TensorRT-LLM

karljang · 2026-11-02T23:34:26Z

Summary

The cuDNN FP8 attention path for MiniMax-H3 on SM 107, including what it needs under Ulysses:

a trailing key_padding_mask in the cuDNN backend (each batch item runs on its valid prefix; Ulysses pads 37,820 and 104,076 tokens at degree 8), with the padded-prefix output written in place (H3_CUDNN_PREFIX_INPLACE=1, no pad + cat copies);
QK-norm, RoPE and the FP8 per-tensor cast fused into one amax pass and one cast pass (fused_fp8_qk.py, Triton), replacing the compiled norm + RoPE pair and three per-tensor scale kernels; dynamic scales by default, a static scale table as an opt-in (H3_FP8_SCALE_FUSION=dynamic|static, H3_FP8_STATIC_SCALES);
the Ulysses exchange can carry the e4m3 codes the fused cast produces, so the all-to-all moves 37.5 % fewer bytes (bit-identical with receiver-exact scales).

Measured internally on SM 107 with the static NVFP4 checkpoint, one GPU, 544×960×124 frames: cuDNN FP8 attention is faster than FA4 2-CTA at 19k, 38k and 104k tokens and the gap grows with the sequence length; the fused QK cast and the in-place prefix write (Ulysses degree 4) each pay on top. Timings are not published for this part. The fused cast passed a blind pairwise quality check (0 of 10 prompts lost); the other changes are bit-identical.

Stacked on #21089 (loader), #21275 (FA4 2-CTA default) and #19738 (cuDNN FP8 on SM 107 by @wu6u3tw): review only the commits after those. The norm + RoPE half of the fused cast overlaps #19645's BF16 kernel; before this leaves draft the fused cast will be rebuilt as an extension of that kernel rather than a second one, and the env switches will become config fields.

Test Coverage

tests/unittest/_torch/visual_gen/test_attention_cudnn.py and multi_gpu/test_attn2d_attention.py carry the [None][fix] cuDNN attention on Rubin (SM107): allow FP8 and fix LSE layout #19738 and padding-mask tests.
The fused cast and the in-place prefix write were verified bit-identical / pairwise-checked end to end on SM 107 at three shapes and Ulysses degrees 1, 4 and 8; unit tests for fused_fp8_qk against the unfused path will be added before the PR leaves draft.

PR Checklist

PR title follows the TensorRT-LLM format
PR description explains what and why
Test cases added
Documentation: module docstrings in cudnn.py and fused_fp8_qk.py

🤖 Generated with Claude Code

…ckpoints Accept a static NVFP4 or per-tensor FP8 quantization_config in the MiniMax-H3 transformer so a ModelOpt export loads as exported: quantized weights, scales, calibrated input_scale and the ignore list. Split only the row-indexed ff.net.0.proj tensors (weight, bias, a 2-D weight_scale) into gate/up and share the per-tensor scalars and input-side tensors. Cast each media projection's input from its own quant state instead of the global algorithm so excluded FP32 projections keep FP32 inputs. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

- Test helper passed quant_config twice; it now takes the override directly. - A one-row ff.net.0.proj tensor (a [1]-shaped per-tensor FP8 weight_scale) is shared between gate and up instead of being split. - Checkpoint-backed test: a ModelOpt-style NVFP4 and FP8 export of the tiny model goes through load_weights and one forward on the static path. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

_TINY_CONFIG has no transformer blocks, so transformer_blocks[0] raised IndexError on the GPU run. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…d name the first non-finite Linear Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…tic-checkpoint Linear amplitudes The media-projection input cast reads has_any_quant on every projection, so the recording stub in the dynamic-FP8 input-dtype test needs the attribute. The static-checkpoint forward hooks now record each quantized Linear's output amplitude and the test rejects implausibly large values alongside non-finite ones, so a GEMM that returns garbage of finite magnitude is named too. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…exporting it TRT-LLM Linears allocate their parameters with torch.empty, so the BF16 reference model's non-quantized tensors, the FP32 media projections among them, carried uninitialised memory into the synthetic checkpoint. The test now gives the reference the deterministic golden initialisation used by the Diffusers parity tests. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…in case With the reference model initialised, the FP8 forward is finite on SM 107 as well, so the expected failure no longer applies. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…M107 FA4 fixes its 2-CTA switch from FA_DISABLE_2CTA when flash_attn.cute.utils is imported, and the VisualGen FA4 backend disables the kernel process-wide because it regresses long sequences on SM100. MiniMax-H3's transformer blocks run attention at the packed clip length (19k to 104k tokens, 56 heads x 128), where the 2-CTA kernel is the faster BF16 attention on SM103 and SM107. - FlashAttn4Attention takes use_2cta; it steers FA4's module flag around the forward and restores it. An explicit FA_DISABLE_2CTA in the environment wins. - The shared Attention module forwards backend_kwargs to create_attention. - MiniMaxH3Transformer3DModel requests the kernel for its transformer blocks on compute capability 10.3 and 10.7 only; the token refiner and other parts keep FA4's default. - Unit tests for the capability default, the plumbing to the block backends and the per-call flag toggle; an FA4 output-equivalence test on a 2-CTA-eligible shape. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

_TINY_CONFIG has no transformer blocks, so indexing transformer_blocks[0] raised IndexError on GPU runs. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

CuDNNAttention.check_hardware_compatibility rejected quantized attention on any SM other than 100/103, so FP8/MXFP8 cuDNN attention raised at graph build on Rubin (SM107) even though cuDNN ships a native SM107 FP8 flash-attention kernel (cudnn_generated_fort_native_sdpa_sm107_flash_fprop_f8). Add SM107 to the allowed set, name Rubin in the error message, and add host-independent tests for the gate. Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

Ring and Attention2D sequence-parallel attention consume forward_with_lse output as [B, H, S], which FA4 and FlashInfer return, but CuDNNAttention transposed it to [B, S, H]. Under an Attention2D mesh with more than one rank per row, flash_attn_combine then weighted partial outputs with LSE from the wrong query/head positions. Return the cuDNN stats as [B, H, S], update the unit test, add an Attention2D test with the cuDNN inner backend, and drop the SM 107 code comment per review. Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

check_hardware_compatibility read the SM version of CUDA device 0 instead of the device the tensors live on, so on a mixed-GPU host it could reject a supported GPU or admit quantized attention on an unsupported one. Query the given device, skip the check for unquantized attention, and add a mixed-device regression test. Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…attention backend Run each batch item on its valid prefix and zero-fill the padded output rows (the discard_padded_queries contract), with the per-request prefix plan cached by mask identity and mutation counter. The FP8 and MXFP8 recipes run on the prefix unchanged. MiniMax-H3 lists CUDNN among the backends that honour the mask so Ulysses padding works with cuDNN attention. Ported from the rc1-based branch feat/minimax-h3-cudnn-padding-rc1 (5373af8433) onto main c8ee74b for the §657 sweep: the pipeline keeps main's CUDA-graph rejection and the FA4 2-CTA default stays in its own PR (NVIDIA#21275). Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…t for cuDNN FP8 attention Lever 8 of the sweep-nvfp4 trial (board 657): on the cuDNN per-tensor FP8 recipe, one Triton pass normalises and rotates Q and K from the packed QKV projection and casts Q, K and V to FP8 E4M3 with per-tensor scales (dynamic amax pass + cast pass, or static scales); the cuDNN backend accepts pre-quantized inputs with their descales. Opt-in through H3_FP8_SCALE_FUSION=dynamic|static. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…oPE + FP8 cast norm_rope_tensor(x, norm_weight, cos, sin, eps, fp8=...) runs one [B, S, H, D] tensor (norm + RoPE for Q / K, identity for V) to BF16 or FP8 with its per-tensor descale, for the async Ulysses path that projects Q, K and V separately (board 657, lever 13 request). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…d MiniMax-H3 FP8 cast H3_FP8_FUSION_SPLIT=1 runs the packed QKV through the per-tensor kernels (three amax and three cast launches with the packed row stride) for the launch-shape sweep; the static-scale table accepts the calibration log written by the dynamic runs. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…e fused FP8 cast op A torch custom op may not return a tensor that aliases one of its inputs; the static-scale branch returned the table itself. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

H3_FP8_SAT_LOG=<json> counts the elements clipped at the E4M3 range per layer and tensor (static-scale calibration check, board 657 lever 8 job 4). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…x-H3 FP8 cast; define interior_padding When the NVFP4 producer fusion hands the attention an Fp4QuantizedTensor, the fused QK-norm + RoPE + FP8 cast reads device, rank and sequence length from the codes tensor (board 657, stack-2 precheck 922175). The cherry-pick of the cuDNN padding commits onto the stack branch had dropped the interior_padding definition in the padded-sequence check; U1 runs never reach it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…tic FP8 cast variant The three tables measured in board 657 lever 8 (in-calibration, held-out margin 1.0, held-out margin 1.194), the two held-out prompt sets, the saturation logs and the flags that reproduce the alternative row of the stack post. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…iMax-H3 cast Under Ulysses the ranks hold shards of one sequence and the FP8 codes travel through the all-to-all; the dynamic per-tensor amax is now all-reduced (MAX) over the Ulysses group before the cast, so every rank casts and dequantizes with the same scale (board 657, stack-2 U4 arms). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

With Ulysses trailing padding the padded tokens carry the AdaLN shift and are never read by the attention; the fused QK-norm + RoPE + FP8 cast now takes the per-item valid length (from key_padding_mask and the shard rank) and leaves those rows out of the per-tensor amax (board 657, lever 13b finding). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…-H3 FP8 cast Keeps MiniMaxH3Attention.__init__ unchanged so the async-Ulysses overlay anchors there as before (board 657, composed-set check 15083). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…sses Document the all-reduce count per entry point and the ruling that an asynchronous-Ulysses arm carrying lever 8 uses the static table. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <kanghwanj@users.noreply.github.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…e and 2 warps Lever 8c (board §657 ¶15221): the same launch shape wins at 19k, 38k and 104k rows, bit-identical to the 8-head / 4-warp launch (heads per warp stays 2) and 5-7 % faster on the amax + cast pair at the long shapes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <kanghwanj@users.noreply.github.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

… 8 norm + RoPE kernel forward_async fuses each projection's per-head RMSNorm and RoPE with the lever 8 per-tensor kernel when its overlay is present (fused_fp8_qk.py) and, with H3_L13_PAIR=fp8 on an FP8 attention backend, casts to e4m3 before the push so the all-to-all carries 1-byte codes viewed as bf16; UlyssesAttention views the received tensors back through the new recv_view_dtype kwarg. The plain async path is unchanged without the overlay. Measured internally at Ulysses degree 4 on SM 107 (sweep/nvfp4 lever 13): NCCL knobs are inert; the async overlay loses at 19k tokens and is a small gain at 104k on both rows; the FP8-codes pairing in dynamic mode loses at 19k (150 small all-reduces per step) and pays over the sync lever 8 path at 104k. Timings are not published for this part. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…t in place The trailing-padding path ran attention on the valid prefix, zero-filled a full-size copy (F.pad) and copied it again (torch.cat) per layer. Allocate the padded [B, S, H, D_v] output once, zero only the padding rows and let the cuDNN graph write the prefix through an output view with the same strides, so the plan and the values are unchanged. Study switch H3_CUDNN_PREFIX_INPLACE=1. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: kanghwanj <861393+karljang@users.noreply.github.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…ntion path Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…tes from the cuDNN path Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

…DIA#21275 NVIDIA#19738 NVIDIA#21485 NVIDIA#21483 NVIDIA#21484 One branch that is the MiniMax-H3 recipe for SM 107, so it can be checked out and built instead of composed from per-lever patches at job time. Base: NVIDIA/TensorRT-LLM main aea80a0 (2026-11-06). The PR branches are merged in this order, every original commit kept: NVIDIA#21089 feat/minimax-h3-static-nvfp4-load static NVFP4 / FP8 checkpoint loader NVIDIA#21275 feat/minimax-h3-fa4-2cta-default FA4 2-CTA default on SM 103 / SM 107 NVIDIA#19738 wu6u3tw:fix/cudnn-fp8-attn-sm107 cuDNN FP8 attention on SM 107, LSE layout NVIDIA#21485 feat/minimax-h3-cudnn-fp8-ulysses cuDNN padding mask, fused FP8 QK cast, prefix output in place, async-Ulysses FP8 pairing NVIDIA#21483 feat/minimax-h3-fused-nvfp4-quant producer quantize fusion, residual fold, attention-output quantize NVIDIA#21484 feat/minimax-h3-vae-decoder-fp8-tiles schedule, FP8 tiled VAE decoder This commit applies the pre-commit fixes on the merged tree (ruff import order in pipeline_minimax_h3.py, ruff-format of vae_decoder_fp8.py), drops a measured kernel share from the vae_decoder_fp8 docstring, and routes the fused FP8 QK attention path's output through _out_proj: NVIDIA#21485 added that return after NVIDIA#21483 was cut, so the merge left it calling to_out directly and the attention-output quantize of NVIDIA#21483 never ran on the cuDNN FP8 row. How to turn the recipe on (NVFP4 checkpoint, request flags --attention-backend CUDNN --qk-dtype fp8 --v-dtype fp8): # the composed arm, every layer on (U1 and U4; drop H3_CUDNN_PREFIX_INPLACE only if the shape divides by U) export H3_FUSED_QUANT=1 H3_FUSED_QUANT_WARPS=4 H3_LEVER10_FOLD=1 export H3_FP8_SCALE_FUSION=dynamic H3_FP8_FUSION_SPLIT=1 export H3_OUT_QUANT=1 H3_CUDNN_PREFIX_INPLACE=1 export H3_L14_SCHED=1 H3_L14_VAE_FP8=1 H3_L14_VAE_TILE=288x272x32 # cuDNN FP8 attention: apply PR NVIDIA#19738 (it allows FP8 on SM 107 and fixes the LSE layout) to backends/cudnn.py first; # the measurements carried it as a shim ahead on PYTHONPATH, switched by H3_UNGATE_CUDNN=1; that shim is not in this folder # do NOT set H3_FP8_FUSION_BLOCK_H / H3_FP8_FUSION_WARPS: the module default (4 / 2) is the tuned launch # attention selection on the request: --attention-backend CUDNN --qk-dtype fp8 --v-dtype fp8 # U8 at 104k tokens: static FP8 scales instead of dynamic (table of 12 held-out prompts, in this folder) # (H3_FP8_STATIC_MARGIN defaults to 1.0) export H3_FP8_SCALE_FUSION=static H3_FP8_STATIC_SCALES=patches/fp8-static-scales-heldout.json # U8 at 19k tokens: all ten patches applied, step fusions off export H3_FUSED_QUANT=0 H3_OUT_QUANT=0; unset H3_LEVER10_FOLD H3_FP8_SCALE_FUSION H3_FP8_FUSION_SPLIT # fallback row (BF16 FA4 2-CTA): skip the NVIDIA#19738 step, unset H3_FP8_SCALE_FUSION H3_FP8_FUSION_SPLIT H3_CUDNN_PREFIX_INPLACE, # request FA4 and let the 2-CTA default of patch 2 apply (section 6 for the worker spawn) # the FP8 option (section 2.2): point the model at the static FP8 export (composite over the BF16 checkpoint), # same ten patches and the NVIDIA#19738 step; the NVFP4 fusions off, the static scale table never export H3_FUSED_QUANT=0 H3_OUT_QUANT=0; unset H3_LEVER10_FOLD export H3_FP8_SCALE_FUSION=dynamic H3_FP8_FUSION_SPLIT=1 H3_CUDNN_PREFIX_INPLACE=1 export H3_L14_SCHED=1 H3_L14_VAE_FP8=1 H3_L14_VAE_TILE=288x272x32 # request: --attention-backend CUDNN --qk-dtype fp8 --v-dtype fp8 The block above is the recipe report's "how to turn it on" verbatim; on this branch three of its lines are already satisfied by the code: NVIDIA#19738 is merged (no shim, no H3_UNGATE_CUDNN), H3_LEVER10_FOLD has no switch (the residual fold is part of H3_FUSED_QUANT, which defaults to on, as does H3_OUT_QUANT), and "patch 2" is the FA4 2-CTA default of NVIDIA#21275. The package ships one static FP8 descale table (calibrated at 544x960x124) under tensorrt_llm/_torch/visual_gen/models/minimax_h3/fp8_scale_tables/; the held-out table the report names for the U8 104k row stays beside the report. Signed-off-by: Kanghwan Jang <kanghwanj@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

karljang · 2026-11-10T08:46:54Z

Closing this draft: the work continues on the integration branch karljang/TensorRT-LLM:feat/minimax-h3-r200-stack and will be re-proposed from the ModelingV2 scenario layout.

karljang and others added 29 commits November 2, 2026 15:33

[None][test] Static-checkpoint load test builds one transformer block

ca8da09

_TINY_CONFIG has no transformer blocks, so transformer_blocks[0] raised IndexError on the GPU run. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][test] Widen the static NVFP4 / FP8 load test to GEMM-legal shapes

1d1bc77

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][test] Derive the static-checkpoint fixture scales from amax an…

63a8798

…d name the first non-finite Linear Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][test] Expect the FP8 static-checkpoint forward to fail on SM 107

664ae0c

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][chore] Format the static-checkpoint test

0f88380

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][test] Run the FP8 static-checkpoint forward on SM 107 as a pla…

db2cb21

…in case With the reference model initialised, the FP8 forward is finite on SM 107 as well, so the expected failure no longer applies. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][test] Build the 2-CTA plumbing tests with one transformer block

8e8c600

_TINY_CONFIG has no transformer blocks, so indexing transformer_blocks[0] raised IndexError on GPU runs. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][chore] Format the FA4 2-CTA tests

ef220bd

Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

github-actions Bot assigned karljang Nov 2, 2026

github-actions Bot added the VisualGen label Nov 2, 2026

karljang and others added 4 commits November 5, 2026 12:27

[None][chore] Apply pre-commit fixes to the MiniMax-H3 cuDNN FP8 atte…

79f6f3b

…ntion path Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

[None][chore] Keep one static FP8 descale table and drop the study no…

4b001a0

…tes from the cuDNN path Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>

karljang force-pushed the feat/minimax-h3-cudnn-fp8-ulysses branch from eeeb503 to 4b001a0 Compare November 5, 2026 20:32

karljang closed this Nov 10, 2026

Navigation Menu

Repository navigation

[None][perf] cuDNN FP8 attention path for MiniMax-H3 on SM107 under Ulysses - #21485

[None][perf] cuDNN FP8 attention path for MiniMax-H3 on SM107 under Ulysses#21485
karljang wants to merge 33 commits into
NVIDIA:mainfrom
karljang:feat/minimax-h3-cudnn-fp8-ulysses

karljang commented Nov 2, 2026 •

edited

Loading

Uh oh!

karljang commented Nov 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

karljang commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Coverage

PR Checklist

Uh oh!

karljang commented Nov 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

karljang commented Nov 2, 2026 •

edited

Loading