[None][feat] Load static FP8 Cosmos3 Nano/Super checkpoints without re-quantizing them by ishovkun · Pull Request #17476 · NVIDIA/TensorRT-LLM

ishovkun · 2026-08-10T20:38:00Z

Description

The pre-quantized (ModelOpt-calibrated) Cosmos3 Nano and Super FP8 checkpoints
already loaded, but they did not load faithfully.

Static FP8 calibrates one scale per projection. The Cosmos3 topology fused the
generation tower's Q/K/V into one Linear and both towers' gate/up into
another. A per-tensor FP8 Linear holds exactly one weight scale, so the loader
kept max(scales) and re-quantized the other members of each group onto it:

qWi' = FP8-round((si * qWi) / max(s)) instead of keeping si * qWi

That is a second rounding on top of ModelOpt's. It affected 141 of 252 fused
shards in Nano and 252 of 448 in Super (~56% in both), and on affected shards
the added error was 7-18x the numerical floor measured without the extra
conversion. The checkpoints were fine; the fusion was discarding their
calibration.

This PR keeps those projections separate under static FP8, so every tensor and
scale loads exactly as calibrated. One predicate,
transformer_cosmos3.uses_static_fp8(), gates all of it. BF16 Cosmos3,
dynamically quantized Cosmos3, and every other model keep the fused topology
unchanged.

vLLM-Omni's Cosmos3 implementation also keeps these projections separate; the
single-scale fusion was specific to this topology.

It is also faster

Unfusing was expected to cost a little. It does not — the split topology is
2.13% faster end to end (Nano, default T2V, one B200: 104.62 s vs 106.90 s,
variants interleaved at the load level, 4 measured runs each after a discarded
warmup, ranges non-overlapping).

The win is not in the GEMMs (CUPTI puts split QKV at 1.03x fused and gate/up at
1.00x). Fusing QKV hands downstream code qkv.split() views that keep the
fused row stride, so per-head QK RMSNorm, RoPE and the und/gen K/V concatenation
all read strided. Split projections are contiguous Linear outputs. A
kernel-level profile diff attributes the delta to exactly those ops.

What changed

Shared infrastructure (opt-in, inert for every existing caller):

swiglu.py / torch_custom_ops.py — a new two-input SwiGLU,
trtllm::silu_and_mul_2in, computing SiLU(G) * U from two separate tensors
and quantizing the result directly for the FP8 down_proj. Bit-exact against
the shipped fused kernel on every tested BF16/FP16 shape, with swiglu_limit,
with swiglu_alpha/swiglu_beta, and when emitting FP8. It is 0.79x the
shipped fused activation on the FP8 path Cosmos3 takes (1,047 -> 824 us Nano,
2,211 -> 1,725 us Super), DRAM utilization 68% -> 86-87% (B200). Full
custom_op contract: register_fake, opcheck,
torch.compile(fullgraph=True), CUDA-graph capture.

Its launch parameters (get_silu_b200_tuning_params) are constants from a
B200 sweep, deliberately scoped rather than generalized: a B300 (sm_103)
sweep puts them within 0.4% of the best configuration measured there, and the
docstring directs a future architecture to add its own table and dispatch on
device capability rather than widening this one.

torch.cat([gate, up]) was rejected — it adds a full intermediate-sized GPU
copy — as was writing both GEMMs into column slices of one [M, 2I] buffer,
which is silently mis-written because cublas_gemm_caller hardcodes
ldc = n (cublasScaledMM.cpp) while cublas_scaled_mm_out still validates
the output stride.
gated_mlp.py — opt-in GatedMLP(split_gate_up=...). Quantizes the shared
activation once and hands both projections the FP8 tensor
(FP8QDQLinearMethod.apply passes a pre-quantized input straight through).
Rejects LoRA and force_dynamic_quantization explicitly rather than by
accident.
visual_gen/modules/attention.py — opt-in
Attention(share_qkv_input_quant=...), the FP8 analog of the existing NVFP4
dedup, with the same single-quantization property for Q/K/V.

Both carry a post_load_weights() validating the invariant the sharing rests
on — that the group's projections carry the same calibrated input_scale.
Checked at load, never in forward(): reading a scale tensor on the hot path
would sync the device and break fullgraph compilation. The invariant was
verified exhaustively across the shipped checkpoints — 108 groups in Nano, 192
in Super, zero with differing input scales — but a checkpoint that violated
it would otherwise produce silently wrong output.

Cosmos3 wiring (transformer_cosmos3.py): uses_static_fp8() drives all
four sites — GEN cross-attention (SEPARATE_QKV + shared quant), UND causal
attention (shared quant; already separate), and both towers' GatedMLPs.
post_load_weights() gained a second pass over GatedMLP/Attention, because
named_modules() yields parents before children and the invariants must be
checked after the projections finalize.

The loader needed no change. DynamicLinearWeightLoader dispatches on each
module's own weight_mode, and the existing key remap already emits
per-projection checkpoint names, so a split (VANILLA) module matches its
checkpoint tensor directly while params_map continues to serve the fused BF16
path untouched.

Supported surface, and what is refused

Nano and Super; T2V, T2I, I2V and V2V; one GPU; static per-tensor
QuantAlgo.FP8 checkpoints only.

Static FP8 raises NotImplementedError for tp_size, ulysses_size,
cfg_size, cp_size or parallel_vae_size above 1, rather than running
untested. The guard is deliberate: putting cross-attention on SEPARATE_QKV
changes how the parallel wrappers treat it — with Attention2D engaged, Ulysses
is disabled rather than failing, announced only by a logger.debug — and a
quiet fallback is worse than a refusal. Existing distributed Cosmos3 tests use
synthetic BF16 weights and do not validate FP8 sharding, per-rank scales, or the
NVLink/NCCL collective paths.

Also out of scope, and documented as such in the example README:

Audio (T2AV/TI2AV) runs, and is not refused. Both checkpoints ship the
audio tower (sound_gen: true), so audio is quantized and generated like any
other task; the single-GPU guard covers only the parallel axes. It is simply
less exercised than the four video/image tasks and carries no quality claim.
No quality claim. There is no accuracy or generative-quality benchmark for
these checkpoints. This PR claims functional support, checkpoint fidelity, and
reproducible output — not that FP8 and BF16 are visually equivalent. The
LPIPS golden gates drift from this path's own pinned render; that is a
regression check, not a quality measurement.
No Hub IDs. The checkpoints live in subdirectories and are addressed by
local path. hf_ids registration is deferred until directly loadable official
FP8 repository IDs exist.

One interaction is documented but left unguarded: SEPARATE_QKV
cross-attention also falls back from CUTEDSL VSA to VANILLA, which would drop
sparse attention under FP8 with no diagnostic. No shipped Cosmos3 config enables
VSA, so the combination is unreachable today.

Test Coverage

New, 78 unit tests:

tests/unittest/_torch/modules/test_swiglu_2in.py (32) — torch.library.opcheck,
torch.compile(fullgraph=True), CUDA-graph capture, rank-N inputs,
non-contiguous rejection, and bit-exactness against the shipped kernel for
BF16 / FP16 / swiglu_limit / swiglu_alpha+swiglu_beta / FP8 output.
tests/unittest/_torch/modules/test_gated_mlp_split.py (19) — split topology,
single-quantization by storage identity (data_ptr, not dtype — three
independent quantizations would also yield FP8 everywhere), scale-mismatch
rejection, LoRA and dynamic-quantization refusals.
tests/unittest/_torch/visual_gen/test_attention_qkv_share.py (12) — the same
for Q/K/V, including the cross-attention case where k/v come from a different
tensor and sharing must not engage.
tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py (15) — checkpoint-gated.
test_previously_fused_groups_now_load_exactly compares FP8 weights bitwise
against the raw checkpoint tensors, and first asserts the group's weight
scales genuinely differ — otherwise fusion would have been lossless and
exactness would hold trivially under either topology, leaving the test unable
to discriminate. test_topology_follows_quantization and
test_dynamic_quantization_stays_fused pin all four directions so the BF16
and dynamic paths cannot drift silently.

Regression: the existing test_fused_activation_quant.py and
test_cosmos3_transformer.py suites stay green.

E2E: test_visual_gen_cosmos3.py::test_cosmos3_nano_fp8_t2i_lpips_against_golden
— one LPIPS self-golden, matching how fp8-blockwise and nvfp4 are covered.

It is deliberately not referenced against BF16. Diffusion sampling is
chaotic, so an FP8 rounding difference at step 0 compounds into a different but
equally valid sample; measured FP8-vs-BF16 LPIPS runs 0.025 (T2I) to 0.175
(T2V), far above any threshold that would still catch a defect. Against its own
pinned render at a fixed seed that problem does not arise — two independent
renders under deterministic algorithms came out bit-identical (LPIPS
0.000000, max pixel delta 0), so the 0.05 threshold is entirely headroom for
cross-host kernel drift, the same allowance the sibling Cosmos3 goldens carry.

The golden is deliberately cheap: 256x256 at 4 steps — the smallest legal
resolution (VAE spatial factor 16 x latent patch 2, so multiples of 32), roughly
1/100th the cost of the 720x1280/35-step goldens beside it. The gate only has to
run every quantized projection in every layer once; it is not judging image
quality. It drives output_type="image", since the pipeline keys text-to-image
off output_type — which also selects COSMOS3_T2I_PARAMS, the T2I system
prompt and the image resolution template — so a num_frames=1 video run would
exercise none of the T2I path.

The generator asserts a GatedMLP actually has split gate/up projections before
rendering. Without that, a silent revert to the fused topology would produce a
nearly identical image, pass, and let a regenerated golden enshrine the
re-quantized weights this change exists to avoid.

CI (l0_b200.yml): the three checkpoint-free suites and test_cosmos3_fp8.py
run pre-merge beside the other visual_gen unit suites; the golden runs
post-merge beside the other Cosmos3 LPIPS gates. The FP8 checkpoints are not
staged in the CI llm-models/ tree, so the checkpoint-gated nodes skip
rather than run — a skip reports green, so they should not be read as evidence
that the FP8 path is exercised in CI today.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

silu_and_mul requires gate and up adjacent in one tensor. Models that keep them as separate projections have no such tensor, and concatenating one costs a full intermediate-sized copy of the activation. silu_and_mul_2in consumes the two tensors directly. It indexes both operands as flat contiguous runs over a 1-D grid, so it is rank-agnostic: models carry rank-3 [batch, seq, hidden] activations and requiring rank 2 would force callers to reshape. Strided inputs are rejected rather than copied, since the flat indexing would otherwise read the wrong addresses and silently return garbage. Launch parameters were tuned on B200 and live in get_silu_b200_tuning_params. Other architectures are unvalidated but not rejected; dispatch on device capability when a second architecture is tuned. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Statically quantized checkpoints calibrate one scale per projection. Fusing q/k/v or gate/up into a single Linear forces one scale on the group and re-quantizes the other members onto it, discarding their calibration. Add two opt-in flags, both inert by default so every existing caller keeps its fused topology: GatedMLP(split_gate_up=...) separate gate and up projections Attention(share_qkv_input_quant=...) quantize the activation once Unfusing must not mean quantizing the same activation two or three times. The members of a formerly fused group all consume it, so it is quantized once and handed to each projection; FP8QDQLinearMethod.apply passes a pre-quantized input straight through. Attention gains the FP8 analog of the dedup that already existed for NVFP4 in the async path. That sharing is only sound while the group's input scales agree, because each Linear applies its own scale in the GEMM epilogue. Both modules check this in post_load_weights() rather than forward(): reading scale tensors on the hot path would sync the device every call and break fullgraph compilation. Keeping the flags opt-in also makes enabling one a commitment to running that hook. Combinations without a static calibrated scale to share -- dynamic quantization, and a fused QKV that already quantizes once -- are rejected at construction instead of silently degrading. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Cosmos3 fused GEN QKV and both towers' gate/up pairs. On the static FP8 checkpoints that re-rounded 56% of shards onto a neighbour's scale -- 141 of 252 on Nano, 252 of 448 on Super -- at 7-18x the numerical floor. One predicate, uses_static_fp8(), now drives every affected site: GEN cross-attention takes SEPARATE_QKV, both towers share one quantized QKV activation, and both GatedMLPs split gate/up. BF16 and dynamically quantized Cosmos3 keep the fused topology, since neither has per-projection calibration to preserve. The loader needed no change. DynamicLinearWeightLoader dispatches on each module's own weight_mode and the existing key remap already emits per-projection checkpoint names, so a split (VANILLA) module matches its tensor directly while params_map continues to serve the fused path. Runtime Linear counts now equal the checkpoint's own projection counts: 252/252 on Nano, 448/448 on Super. post_load_weights() gains a second pass over GatedMLP and Attention. named_modules() yields parents before children, so folding it into the existing loop would check the shared-scale invariants before the projections had finalized. Static FP8 also rejects every multi-GPU axis outright. Moving cross-attention to SEPARATE_QKV changes how the parallel wrappers treat it -- with Attention2D engaged, Ulysses is silently disabled rather than failing -- and a quiet fallback is worse than a refusal. Measured 2.13% faster than the fused path it replaces at the default T2V request (720x1280, 189 frames, 35 steps) on one B200: 104.62 s against 106.90 s, medians of 4 interleaved runs each. The gain is not arithmetic. Fused QKV hands downstream code qkv.split() views that keep the fused row stride, so QK norm, RoPE and the und/gen K/V concatenation all read strided; split projections are contiguous Linear outputs. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

74 tests across four suites. The kernel and module suites assert bit-exact equality with the fused path rather than a tolerance: both do the same arithmetic in the same order and accumulate in fp32, so any difference is a defect. Shared quantization is checked by storage identity, not dtype -- three independent quantizations of one activation would also yield FP8 everywhere, so dtype alone proves nothing. test_previously_fused_groups_now_load_exactly replaces test_fused_shard_requantization_stays_within_measured_bound. The old test bounded relative RMS at 2e-3 against a global-max denominator that hid the smaller-scale shard, and sampled one of roughly 125 affected groups. The new one compares FP8 weights bitwise against the raw checkpoint tensors, and first asserts the group's weight scales actually differ -- were they equal, fusion would have been lossless and exactness would hold trivially under either topology, leaving the test unable to tell them apart. test_topology_follows_quantization and test_dynamic_quantization_stays_fused pin all four directions of the predicate, so BF16 and dynamic Cosmos3 cannot drift onto the split path unnoticed. The I2V and V2V collapse smokes now also assert on a frame no conditioning reaches. Both tasks copy structure from the reference into their leading frames, so a whole-video pixel standard deviation could clear the threshold even if every generated frame had collapsed. The three checkpoint-free suites are staged on l0_b200. The checkpoint-dependent ones stay developer-run-only until the FP8 checkpoints are published, since they skip when the directories are absent and a CI stage would report green without running anything. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Quantization is detected from the checkpoint's own metadata, so the FP8 builds are passed to --model exactly like a BF16 one. There are no FP8 Hub IDs yet, so the example uses a local path. Records the boundaries the tests actually establish: T2V/T2I/I2V/V2V on a single GPU. Audio and multi-GPU are untested and rejected, and FP8 output quality has not been benchmarked against BF16. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

The fused silu_and_mul applies alpha as a gain inside the sigmoid and beta as an offset on up. The two-input variant hardcoded the alpha=1, beta=0 case, so a GatedMLP built with split_gate_up and a swigluoai-shaped activation would have computed plain SwiGLU instead -- numerically wrong, and silently so. Also drop override_tp_sharding from the kwargs shared by both topologies. Linear asserts that a dict tp_sharding only ever reaches a fused weight mode, so passing it to the split VANILLA projections tripped that assertion; it now goes to the fused projection only, as upstream had it. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

…ghts Adds bit-exactness cases against the fused kernel for the swigluoai shape and for alpha and beta in isolation, plus one pinning that omitting both stays identical to plain SwiGLU. The split GatedMLP tests never initialized their FP8 weights, so they read whatever the caching allocator returned. That is NaN often enough to matter -- 11 of 12 constructions in one probe -- and any torch.equal assertion then failed regardless of the behaviour under test, because torch.equal is false for NaN even when both sides are bit-identical. Three tests failed this way; the same pattern appeared in six more that happened not to compare values. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

The checkpoint-gated FP8 tests were left out of every test database while their neighbours -- test_cosmos3_pipeline, test_cosmos3_distilled and the Cosmos3 LPIPS gates -- are all scheduled despite needing real checkpoints too. That made the FP8 loading/topology behaviour and the four advertised tasks the only Cosmos3 surface with no CI coverage. test_cosmos3_fp8.py now runs pre-merge beside the other visual_gen unit suites, and the eight decode smokes run post-merge beside the other end-to-end Cosmos3 examples. The README also claimed audio was "untested and rejected". Nothing rejects it: the single-GPU guard covers only the parallel axes, and both checkpoints ship the audio tower (sound_gen: true), so T2AV/TI2AV quantize and generate like any other task. Describe what the code does -- audio runs, it is simply less exercised than the four video/image tasks and carries no quality claim. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

ishovkun · 2026-08-10T21:40:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T21:45:49Z

PR_Github #65132 [ run ] triggered by Bot. Commit: e75a564 Link to invocation

The t2i case only set num_frames=1. The pipeline keys text-to-image off output_type, not frame count, and that flag also selects COSMOS3_T2I_PARAMS, the T2I system prompt and the image resolution template -- so the case was running a one-frame text-to-video and asserting on result.video. Thread output_type through the shared helper, return result.image for T2I, and let the collapse assertion accept the rank-4 (B, H, W, C) result via a view. Callers that omit output_type keep the video path, so the existing Cosmos3 LPIPS gates are unaffected. The checkpoint suite also set TRTLLM_DISABLE_COSMOS3_GUARDRAILS and TLLM_DISABLE_MPI at import and never restored them, so a combined run leaked both into every later module. Guardrails now come from the existing disable_cosmos3_guardrails fixture, and TLLM_DISABLE_MPI -- which must be set before the imports, so it cannot be a fixture -- is popped by a module-scoped teardown, matching test_cosmos3_pipeline.py. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

tensorrt-cicd · 2026-08-11T00:02:23Z

PR_Github #65132 [ run ] completed with state SUCCESS. Commit: e75a564
/LLM/main/L0_MergeRequest_PR pipeline #52928 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

ishovkun · 2026-08-11T04:18:38Z

/bot run --disable-fail-fast

The eight decode smokes only rejected a near-flat frame, which a subtly misloaded scale clears: the image stays plausible and the pixel spread stays high. Every quantization feature beside this one is covered by an LPIPS golden (fp8-blockwise, nvfp4), so match them. Replace the smokes with one self-golden. It is deliberately not referenced against BF16 -- diffusion sampling is chaotic, so an FP8 rounding difference at step 0 compounds into a different but equally valid sample, and measured FP8-vs-BF16 LPIPS runs 0.025 (T2I) to 0.175 (T2V), far above any useful gate. Against its own pinned render at a fixed seed the problem does not arise: two independent renders under deterministic algorithms came out bit-identical (LPIPS 0.000000, max pixel delta 0), so the 0.05 threshold is entirely headroom for cross-host kernel drift, matching the sibling goldens. The golden is deliberately tiny -- 256x256 (the smallest legal size: VAE spatial factor 16 x latent patch 2) at 4 steps, roughly 1/100th the cost of the 720x1280/35-step goldens next to it. The gate only has to run every quantized projection in every layer once; it is not judging image quality. The generator asserts a GatedMLP actually has split gate/up before rendering. Without that, a silent revert to the fused topology would render a nearly identical image, pass, and let a regenerated golden enshrine the re-quantized weights this feature exists to avoid. Net -34 lines. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

ishovkun · 2026-08-11T16:20:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T16:26:41Z

PR_Github #65330 [ run ] triggered by Bot. Commit: d4079aa Link to invocation

tensorrt-cicd · 2026-08-11T20:27:28Z

PR_Github #65330 [ run ] completed with state SUCCESS. Commit: d4079aa
/LLM/main/L0_MergeRequest_PR pipeline #53102 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

ishovkun added 7 commits August 10, 2026 13:13

github-actions Bot assigned ishovkun Aug 10, 2026

github-actions Bot added the VisualGen label Aug 10, 2026

Navigation Menu

Uh oh!

[None][feat] Load static FP8 Cosmos3 Nano/Super checkpoints without re-quantizing them - #17476

[None][feat] Load static FP8 Cosmos3 Nano/Super checkpoints without re-quantizing them#17476
ishovkun wants to merge 10 commits into
NVIDIA:mainfrom
ishovkun:cosmos3_fp8

ishovkun commented Aug 10, 2026 •

edited

Loading

Uh oh!

ishovkun commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

ishovkun commented Aug 11, 2026

Uh oh!

ishovkun commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ishovkun commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

It is also faster

What changed

Supported surface, and what is refused

Test Coverage

PR Checklist

Uh oh!

ishovkun commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

ishovkun commented Aug 11, 2026

Uh oh!

ishovkun commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ishovkun commented Aug 10, 2026 •

edited

Loading