Repository navigation
Conversation
…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>
3 of 4 tasks
karljang
added a commit
to karljang/TensorRT-LLM
that referenced
this pull request
Nov 6, 2026
…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>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
FA4 reads
FA_DISABLE_2CTAonce, whenflash_attn.cute.utilsis imported, and the VisualGen FA4 backend sets it to1process-wide because the 2-CTA kernel 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(use_2cta=...)steers FA4's module flag around the forward and restores it; an explicitFA_DISABLE_2CTAin the environment wins; FA4's own CUDA 12 opt-out stays.Attentionmodule forwardsbackend_kwargstocreate_attention.MiniMaxH3Transformer3DModelrequests the kernel for its transformer blocks on compute capability 10.3 and 10.7 only (fa4_use_2cta_default); the token refiner and SM100 keep FA4's default.No public API change: a backend kwarg and a model-level default only.
Measurements
Speedup ratio S = T(1-CTA) / T(2-CTA). Block fixture at 19k / 38k / 104k tokens, compile on, 100 CUDA-graph replays, same container and commit per part.
End to end on SM107 (same-job A/B, BF16 T2VA 544x960x124, 19,394 packed tokens, compile on, graphs off) the two clips are bit-identical (the 2-CTA kernel keeps the per-row accumulation order) and the 2-CTA step is the faster one; the timings stay internal. SM103 end to end: to be run before leaving draft.
Test Coverage
tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py::test_fa4_2cta_default_follows_device_capability,::test_fa4_2cta_default_without_cuda,::test_fa4_2cta_request_reaches_transformer_block_backends,::test_fa4_2cta_request_is_not_sent_to_other_backends(CPU; file listed inl0_b200.yml)tests/unittest/_torch/visual_gen/test_attention_fa4.py::test_use_2cta_request_toggles_fa4_flag_for_the_call_only,::test_use_2cta_forward_matches_default_kernel[True|False](GPU; file listed inl0_b200.yml)PR Checklist
🤖 Generated with Claude Code