[None][feat] feat: VisualGen TE-FP8 attention backend + per-layer quant - #15668
Merged
chang-l merged 3 commits intoJun 29, 2026
Merged
Conversation
wu6u3tw
requested review from
JunyiXu-nv,
danielafrimi,
syuoni and
yuxianq
and removed request for
a team
June 26, 2026 17:59
wu6u3tw
changed the title
wu6u3tw
changed the title
wu6u3tw
force-pushed
the
feat/1.3-mlpinf-vg-te-fp8-attn
branch
from
June 26, 2026 18:34
358e788 to
b3361d9
Compare
wu6u3tw
changed the title
torch.load inside __init__ triggers aten.set_.source_Storage_storage_offset on a meta tensor when TRT-LLM initialises weights under MetaInitMode, raising MetaInitException. Defer the load to the first real inference call in generate() where the model is fully materialised. Only the path is stored in __init__; the tensor is loaded and cached on the first generate() invocation, then reused across requests (same behaviour as before for steady-state throughput). Fixes: FL1/FL2 (jobs 2140007/2140008) crash at server startup. Validated: FL72 (job 2191174) DP72 GB300×72, Mean(6)=0.6946 ✅ Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com>
…ayer YAML key)
Add support for mixed-precision quantization in the VisualGen diffusion
pipeline via a `per_layer` key in quant_config YAML. Enables configs like:
quant_config:
quant_algo: NVFP4 # default for all layers
per_layer:
"blocks.*.attn1.*": FP8_BLOCK_SCALES # attention linears use FP8
"blocks.*.attn2.*": FP8_BLOCK_SCALES
Motivation: full-NVFP4 gives +6.3% throughput vs FP8_BLOCK_SCALES but
fails VBench accuracy (dynamic_degree 0.861 < 0.948 threshold). Keeping
attention linear projections at FP8 precision may recover accuracy while
retaining partial NVFP4 throughput gain from FFN layers.
Changes:
- config.py: parse `per_layer` dict in load_diffusion_quant_config() (fills
existing TODO); switch get_quant_config() from exact-match to fnmatch
pattern lookup; enable NVFP4 tunable quantize if any per-layer algo is NVFP4
- transformer_wan.py: apply per-layer quant config overrides to Linear modules
in apply_quant_config_exclude_modules() before create_weights() is called
- loader.py: pass per-layer group_size to _maybe_dynamic_quantize() so
FP8_BLOCK_SCALES attn layers use block_size=128 not the global NVFP4 value
- tests: 5 unit tests (no GPU) + integration smoke (verified on GB300)
New config: exp/configs/wan22_dp_te_fp8_attn_nvfp4_attnfp8block_gb300.yml
VBench accuracy job: 2108769 (x16 GB300, pending results)
Signed-off-by: Tin-Yin Lin <tinyinl@nvidia.com>
Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com>
wu6u3tw
force-pushed
the
feat/1.3-mlpinf-vg-te-fp8-attn
branch
from
June 26, 2026 19:05
402fc86 to
0f21d7e
Compare
wu6u3tw
changed the title
…NC/libx264 encode ## TE FP8 attention (attention_backend: TE) Adds a new attention backend that routes through TransformerEngine's DotProductAttention under fp8_autocast with DelayedScaling(fp8_dpa=True, fp8_mha=True). On Blackwell (GB300), this dispatches to cuDNN Flash Attention FP8, which is significantly more optimized for SM100a tensor cores than the TRTLLM and Sage backends. The module is decorated with @torch.compiler.disable because TE's fp8_autocast context manager causes graph breaks under torch.compile. Input layout is NHD (as used by WanPipeline); the backend reshapes to the bshd layout expected by TE. Wired as backend: "TE" in AttentionConfig; lazy-imported to avoid a hard TransformerEngine dependency when the backend is not selected. Measured on GB300 (Wan2.2-T2V-A14B offline, MLPerf): - Sage (TRTLLM kernel): 0.0254 QPS @ DP4 (-30.6% vs v6.0) - TE FP8 (cuDNN): 0.0368 QPS @ DP4 (+0.5% vs v6.0, parity) Per-step time: 7.6 s (Sage) -> 5.4 s (TE), per-video: 155 s -> 110 s. ## NVENC + libx264 fallback (encoding.py) Refactors _FfmpegCliEncoder to extract a _run_ffmpeg() helper, then adds TRTLLM_VIDEO_NVENC=1 opt-in that attempts h264_nvenc first and falls back to libx264 on failure (e.g. "unsupported device" on datacenter GPUs that have no NVENC engine). Default behavior (libx264 only) is unchanged. Also adds a "raw" output format that dumps uncompressed uint8 RGB bytes to disk, bypassing ffmpeg entirely. Signed-off-by: Tin-Yin Lai <tinyinl@nvidia.com>
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
2 participants
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.
Summary
Two commits on top of
feat/1.3-mlpinf-vg:358e7880da—[feat] VisualGen: TE FP8 attention backend + fixed-latent env varChanges
attention_backend/te.py):TEAttentionusing TransformerEngineDotProductAttentionunderfp8_autocast/DelayedScaling(fp8_dpa, fp8_mha). Registered as"TE"inAttentionConfig.backend.Test plan
🤖 Generated with Claude Code