Cosmos3 fp8 mixed precision - #36380
Merged
mickqian merged 12 commits intoSep 4, 2026
Merged
Conversation
kediwu0331
requested review from
AgainstEntropy,
BBuf,
HaiShaw,
OrangeRedeng,
mickqian,
ping1jing2 and
yichiche
as code owners
August 25, 2026 21:29
kediwu0331
force-pushed
the
cosmos3-fp8-high-precision
branch
from
August 25, 2026 22:41
1f524a4 to
68c6799
Compare
ishovkun
added a commit
to ishovkun/TensorRT-LLM
that referenced
this pull request
Aug 26, 2026
…ions
A ModelOpt checkpoint carries one activation scale per projection, calibrated
as a max over the whole sampling trajectory ('calib_cfg': {'method': 'max'}).
That single scale fits the first and last denoising steps worst. Those steps
can instead run the resident FP8 weights through a 16-bit GEMM: the weight is
dequantized with its own weight_scale and input_scale goes unused, while the
middle steps keep the checkpoint's fully quantized path.
Nothing extra is read from the checkpoint -- same weights, same scales, no
second checkpoint and no persistent dequantized copy. first_steps/last_steps
are a runtime policy, not a calibrated quantity: the checkpoint records no
per-step information of any kind (no step/timestep/schedule tensors, one
scale per module across all 896 quantized modules).
vLLM-Omni (vllm-project/vllm-omni#6560) and SGLang (sgl-project/sglang#36380)
both landed the same mitigation with the same 3/3 windows and the same
one-step warmup carve-out. This matches their semantics; the defaults follow
SGLang, which enables it for every ModelOpt FP8 Cosmos3 checkpoint.
Unlike either of those, this topology quantizes shared activations *above*
the Linear -- gate/up and q/k/v each quantize once and hand the same tensor to
their projections, and swiglu_2in emits FP8 straight into down_proj. All three
must stand down while a 16-bit step is selected, or the step still runs on FP8
activations and the feature is silently absent. A quantization method advertises
that by publishing 'high_precision'; the sharing sites consult it, and
apply_fp8_w8a16_linear raises rather than accept an already-quantized
activation.
Precision is selected once per step from a pure function of the step index, so
a step's conditional and unconditional CFG branches cannot disagree.
Enabled by default for static FP8, off via
VisualGenArgs.step_precision_config.enable; only static FP8 qualifies, since
dynamic quantization derives its scale per call and has no calibration
mismatch for the outer steps to avoid.
Signed-off-by: Igor Shovkun <igshov@gmail.com>
1 task
…ision # Conflicts: # python/sglang/multimodal_gen/runtime/loader/fsdp_load.py # python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
Contributor
Author
|
@mickqian could you help review this PR? |
mickqian
approved these changes
Sep 2, 2026
Collaborator
|
/tag-and-rerun-ci |
Contributor
Author
|
@mickqian I think the requested test case failure is due to transient network issue. Could you help rerun the failed test cases? |
mickqian
approved these changes
Sep 4, 2026
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
…ct#36380) Co-authored-by: Kedi Wu <kediw@nvidia.com>
Merged
10 tasks
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
3 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.
Motivation
Improve the quality of generated video for Cosmos3 fp8 checkpoint. Parse diffusion_step_policy from modelopt fp8 checkpoint. Use W8A16 for specified steps of denoising stage.
Modifications
Accuracy Tests
Accuracy validated for Nano, Super, Super i2v.
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #33800700111
Latest PR Test (Extra): ❌ Run #33800699651
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33800699908