Skip to content

[None][perf] Fuse strided MXFP8 packing for cuDNN attention - #19358

Closed
daichu-nv wants to merge 5 commits into
NVIDIA:mainfrom
daichu-nv:perf/wan-optimization-20260917
Closed

daichu-nv wants to merge 5 commits into
NVIDIA:mainfrom
daichu-nv:perf/wan-optimization-20260917

Conversation

@daichu-nv

Copy link
Copy Markdown
Contributor

Automatically pack eligible BF16 Q/K/V tensors directly into cuDNN's MXFP8 payload and scale layouts. The CuTe kernels remove intermediate layout copies while preserving the native quantization arithmetic. Dispatch selects the fast path on SM100 for supported shapes and strides; unsupported inputs keep the native implementation. No configuration flag is required.

The full-model comparison uses Wan2.2 A14B with the official static-NVFP4 checkpoint on one B200: 480×832, 33 frames, 40 steps, seed 42, three prompts, CUDNN MXFP8 attention, compilation on and CUDA graphs off. Both arms use TensorRT-LLM 1.3.0rc27.dev202609150000 in the same process.

Full-generation latency Native Candidate Reduction
Mean of prompt medians 21.57 s 20.05 s 7.06%

Validation: byte-level payload/scale and cuDNN consumer checks passed; production GPU tests: 34 passed and one expected SM103 hardware skip. Worst-prompt LPIPS across all 33 frames: 0.0 (limit 0.25).

Signed-off-by: Daisy Chu <daichu@nvidia.com>
Signed-off-by: Daisy Chu <daichu@nvidia.com>
Signed-off-by: Daisy Chu <daichu@nvidia.com>
Signed-off-by: Daisy Chu <daichu@nvidia.com>
Signed-off-by: Daisy Chu <daichu@nvidia.com>
@YihuiLu512

Copy link
Copy Markdown
Collaborator

Heads-up: #20214 moved visual_gen/attention_backend/ to visual_gen/attention/backends/; the old package is now a deprecated shim that only forwards package-root exports.

This PR merges cleanly, but the result breaks: your cudnn.py edits follow the rename to attention/backends/cudnn.py, while the new mxfp8_pack.py lands at the old attention_backend/mxfp8_pack.py. As a result, from .mxfp8_pack import ... raises ImportError, and pre-merge CI won't catch it because it only tests the PR head.

Could you rebase and move mxfp8_pack.py to tensorrt_llm/_torch/visual_gen/attention/backends/?

@daichu-nv

Copy link
Copy Markdown
Contributor Author

Let me close this PR first.

@daichu-nv daichu-nv closed this Oct 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants