Repository navigation
[Quantization][Qwen3-Omni] Enable NVFP4 W4A4 serving on Blackwell - #4025
david6666666 merged 1 commit into
Conversation
7e6f7d9 to
df4dda7
Compare
|
Hi @lishunyang12 @baonudesifeizhai — flagging this NVFP4 W4A4 enablement PR for visibility given your prior work on the unified quantization framework (#1764) and the ModelOpt FP8 / mixed FP8-NVFP4 paths (#2913, #3570). The change here is a load-time clamp in |
df4dda7 to
d5f296c
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
| ) | ||
|
|
||
| def _patched_nvfp4_pwal(self, layer, *args, **kwargs): | ||
| _original_nvfp4_pwal(self, layer, *args, **kwargs) |
There was a problem hiding this comment.
Clamping after the original PWAL is too late for the Marlin fallback: prepare_fp4_layer_for_marlin() converts/permutates weight_scale before this runs, so NaN bytes can already have been folded into the Marlin scale format and may no longer be detected here. Please clamp before calling the original PWAL when the upstream method is not already patched.
There was a problem hiding this comment.
Thanks. Will test the Marlin fallback path
| quant_params = [(QUANTIZED_MODEL, _stage_config())] | ||
|
|
||
|
|
||
| @hardware_test(res={"cuda": "H100"}, num_cards=2) |
There was a problem hiding this comment.
This is marked for H100, but the PR/docs say the native W4A4 FP4 path requires sm_100+ Blackwell. On H100 init_nvfp4_linear_kernel() falls back to Marlin weight-only FP4, so this smoke does not exercise the FlashInfer/CUTLASS W4A4 path being enabled. Please add/select a Blackwell resource, or keep this as an explicit fallback test and add a Blackwell smoke.
|
|
||
| current = ModelOptNvFp4LinearMethod.process_weights_after_loading | ||
| name = getattr(current, "__qualname__", "") or getattr(current, "__name__", "") | ||
| assert "patched_nvfp4_pwal" in name or "_patched_nvfp4_pwal" in name, ( |
There was a problem hiding this comment.
This assertion will fail as soon as the self-extinguish path works: when upstream vLLM adds its own clamp, patch.py intentionally skips installing _patched_nvfp4_pwal, so the qualname will not contain this marker. Please accept the upstream-patched case as passing, or skip this assertion when _already_patched_upstream is true.
|
is acc maintained? |
02fc0f9 to
98937b4
Compare
|
Hi @hsliuustc0106 — thanks for the question. Yes. Updated the Test Result section with the relevant evidence:
The clamp doesn't degrade accuracy because it's mathematically equivalent to clamping at the calibration write-side — same |
@lishunyang12 does qwen3-omni support online/offline fp8 quantization now? |
98937b4 to
ba4d056
Compare
| # for-shape with the original tensor; pin that here so a future | ||
| # packed scale format would fail loudly rather than silently corrupt | ||
| # unrelated bytes. | ||
| byte_view = weight_scale.data.view(torch.uint8) |
There was a problem hiding this comment.
.view(torch.uint8) raises on a non-contiguous weight_scale, so a checkpoint whose per-block scale tensor isn't contiguous would crash here before ever reaching the shape assert below. A weight_scale.contiguous() (or an is_contiguous() guard) before taking the byte view would make this robust.
There was a problem hiding this comment.
Good catch — fixed in 5a83ec3. The clamp now normalizes a non-contiguous weight_scale with .contiguous() and writes the copy back to the Parameter before the uint8 view (a bare .contiguous() copy would be discarded, leaving the original tensor — the one the kernel reads — un-clamped). Added a test_clamps_non_contiguous_weight_scale regression test that plants NaN bytes in a strided view.
| _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading = _patched_nvfp4_pwal | ||
| # Fail loudly if install dropped silently (e.g. another plugin | ||
| # patched the same class) rather than degrade to `!!!!` at decode. | ||
| assert _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading is _patched_nvfp4_pwal, ( |
There was a problem hiding this comment.
This assert is compiled out under python -O / PYTHONOPTIMIZE, so in optimized runs the "install failed" guard silently disappears and a conflicting plugin would degrade straight to !!!! at decode — exactly the case this block exists to catch. Since the intent here is to fail loudly, raise RuntimeError(...) would be safer than assert (same goes for the shape assert at L156).
There was a problem hiding this comment.
Agreed — switched both the install self-check (L223) and the byte-view shape check (L156) from assert to raise RuntimeError in 5a83ec3, so they survive python -O / PYTHONOPTIMIZE. These guards exist precisely to fail loudly before a conflicting plugin or a changed scale layout degrades to !!!! at decode, so they are the last things that should be stripped under -O.
| _nan_clamp_import_err, | ||
| ) | ||
| else: | ||
| _original_nvfp4_pwal = _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading |
There was a problem hiding this comment.
On a module reload (e.g. importlib.reload in a test, or a second import path) this captures the already-patched method, so _patched_nvfp4_pwal ends up wrapping itself and the clamp runs twice per load. It is idempotent so not a correctness bug, but stashing a sentinel attribute on the class (or checking _original_nvfp4_pwal.__name__ != "_patched_nvfp4_pwal") would keep the install truly idempotent.
There was a problem hiding this comment.
Fixed in 5a83ec3. The wrapper now stamps a _vllm_omni_wrapped_pwal sentinel pointing at the genuine upstream method, and the install reads it via getattr(..., "_vllm_omni_wrapped_pwal", _current) before re-wrapping — so a module reload recovers the real upstream PWAL instead of capturing our own wrapper and nesting the clamp (which would also have made the co_names self-extinguish heuristic introspect our wrapper rather than upstream). Added test_reload_does_not_nest_wrapper asserting the chain stays flat after an importlib.reload.
lishunyang12
left a comment
lishunyang12
left a comment
There was a problem hiding this comment.
Thanks for this — the load-time NaN-clamp is a well-scoped, well-documented fix for the !!!! failure on community NVFP4 W4A4 checkpoints, and the self-extinguish heuristic plus the clamp-before-PWAL ordering rationale are carefully thought through. My inline notes are all minor robustness/hygiene points, not blockers: (1) the weight_scale.view(torch.uint8) byte view should guard contiguity; (2) the assert-based install/shape guards are compiled out under python -O/PYTHONOPTIMIZE, so the intended loud failure is lost in optimized runs — prefer raise; and (3) the override is not reload-idempotent (it can wrap itself on re-import). None affect the happy path — addressing them just makes the guard hold under -O and on module reload.
ba4d056 to
5a83ec3
Compare
|
Pushed 5a83ec3 addressing @lishunyang12's three review nits (replies inline):
Re-validated on 2x H100 (vllm 0.22.0): 10/10 CPU unit tests pass (incl. new non-contiguous + reload tests), and the e2e smoke on the NaN-reproducer ckpt ( |
5a83ec3 to
bb5d38c
Compare
david6666666
left a comment
david6666666
left a comment
There was a problem hiding this comment.
One more test issue.
| # Falsy values must leave the install in place — this is the case | ||
| # the repo-wide bool-env idiom in vllm_omni.patch guards against | ||
| # (naive `if os.environ.get(...)` would wrongly skip on "0"/"false"). | ||
| ("0", "INSTALLED"), |
There was a problem hiding this comment.
This still conflicts with the self-extinguish path. Once upstream vLLM has its own clamp, "0" / "false" should mean "do not use the escape hatch", but the wrapper should still be skipped because _already_patched_upstream is true. This test would then print NOT_INSTALLED and fail even though the behavior is correct. Please make the falsey cases accept the upstream-patched state too, similar to test_install_is_self_asserting.
There was a problem hiding this comment.
Thanks. Verified the falsey + upstream-patched case you flagged by simulating an upstream clamp before import for the most recent commit.
a78dd4e to
079d0c1
Compare
david6666666
added
ready
|
please fix ci, thx |
Adds a defensive runtime override in vllm_omni.patch of
ModelOptNvFp4LinearMethod.process_weights_after_loading that clamps NaN
bytes in per-block weight_scale tensors at load time.
ModelOpt 0.44's float32 -> torch.float8_e4m3fn cast of per-block
weight_scale occasionally emits literal NaN bytes (E4M3 encoding 0x7F /
0xFF) when the pre-cast scale rounds above the FP8 E4M3 max of 448
after the global-scale division. A single NaN byte in any weight_scale
propagates through the FlashInfer FP4 GEMM and collapses the served
model output to `!!!!`.
The load-time clamp is mathematically equivalent to clamping at
calibration write-side: by the time the kernel reads weight_scale,
every position that would have been NaN is the FP8 E4M3 max byte
(0x7E = 448.0). The proper fixes sit upstream — at the ModelOpt cast
sites, so newly calibrated checkpoints export cleanly, and in vLLM's
own PWAL, so vllm-omni can drop this override. Both are planned as
follow-up PRs after this one merges. Until those land in vllm-omni's
pinned vLLM, any NVFP4 W4A4 Qwen3-Omni checkpoint published by the
community — exported with stock ModelOpt 0.44 — would silently serve
only `!!!!`, which looks like a vllm-omni inference bug to the user.
Newly calibrated clean checkpoints pay no runtime cost (no-op when no
NaN bytes are present).
The override clamps NaN bytes BEFORE delegating to the original PWAL.
The non-Blackwell Marlin fallback's kernel PWAL casts weight_scale
FP8 -> bf16 and permutes; clamping after would trip the byte-view
shape assertion (FP8 1 B/elem vs bf16 2 B/elem) or operate on already-
transformed bytes. Blackwell Cutlass/TRTLLM/cuDNN PWALs only swizzle
FP8 bytes, so clamp-before is equivalent there. Clamping uses a uint8
view + masked_fill_ (PyTorch does not implement masked_fill_ for
float8_e4m3fn directly).
Self-extinguish: an __code__.co_names check requires all three of
{masked_fill_, isnan, weight_scale} in vLLM's PWAL — the structural
signature of an in-place NaN clamp on weight_scale — before declaring
upstream patched and skipping install. Env var
VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP={1,true,yes,on} disables the override
for diagnostics (repo-wide bool-env idiom; "0"/"false" correctly mean
"do not skip"). A dtype guard skips the clamp when weight_scale.dtype
!= torch.float8_e4m3fn, defending against re-invocation on a Marlin-
transformed layer.
Robustness: a non-contiguous weight_scale is made contiguous and
written back to the Parameter before the uint8 view (a bare
.contiguous() copy would be discarded, leaving the original un-clamped);
the install self-check and byte-view shape check raise RuntimeError
rather than assert so they survive python -O / PYTHONOPTIMIZE; and the
wrapper stamps a sentinel attribute so a module reload recovers the
genuine upstream method instead of nesting the clamp.
Hardware validation:
B200 / FlashInfer Cutlass NvFp4 Linear + FLASHINFER_TRTLLM NvFp4 MoE
backend, YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-
awqclip (clean ckpt; clamp is a no-op):
- accuracy: Daily-Omni n=1197 W4A4 0.675 vs BF16 0.694 (-1.9 pp);
OmniBench n=200 W4A4 0.450 vs BF16 0.455 (-0.5 pp)
- throughput: at concurrency=32, output-len=512 random workload,
W4A4 2766 tok/s vs BF16 2083 tok/s (+33%);
TPOT W4A4 9.12 ms vs BF16 12.76 ms (-29%)
- stability: 30-min sustained concurrency=32 produces 3200 successful
responses with 0 errors and 0 `!!!!` collapses
2x H100 SXM / Marlin fallback path, YihongJin/Qwen3-Omni-30B-A3B-
Instruct-NVFP4-W4A4-full-thinker-mse-preview (NaN-containing
reproducer ckpt; clamp active):
- e2e smoke passes (no `!!!!` collapse)
- Daily-Omni n=200: 0.665
2x RTX PRO 6000 WS (sm_120) / FlashInfer Cutlass NvFp4 — same kernel
class as B200, full-thinker-mse-preview (clamp active):
- e2e smoke passes; load logs show 4 NaN bytes detected and clamped
- Daily-Omni n=200: 0.670
The 0.5 pp gap on the same kernel surface (Pro 6000 Cutlass 0.670 vs
B200 Cutlass clean awq_clip 0.675) is attributable to calibration
algorithm (mse vs awq_clip), not the clamp.
Tests:
- tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py:
CPU-only unit tests covering clamp behavior on a planted-NaN
weight_scale (contiguous and non-contiguous), no-op on finite input,
no-op when weight_scale is absent, the install self-check with 3-way
branching on {_already_patched_upstream, _clamp_installed, neither},
reload idempotency (a module reload does not nest the wrapper), and
the env-var escape hatch parametrized over {"1","true"} (expect
NOT_INSTALLED) and {"0","false"} (expect our wrapper INSTALLED OR
upstream already clamping — the self-extinguish state is accepted so
the test still passes once the upstream fix lands) via subprocess,
reporting state from vllm_omni.patch's module flags rather than method
qualname introspection.
- tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py:
hardware-gated e2e smoke on the H100 Marlin fallback path that
loads the published W4A4 checkpoint and asserts the response does
not match the `!!!!` degenerate-output signature.
Docs: docs/user_guide/quantization/modelopt.md adds a Qwen3-Omni W4A4
section with the HF repo link, hardware requirement (sm_100+),
calibration recipe pointer, and a note describing the ModelOpt 0.44
NaN regression and the workaround installed by vllm_omni.patch.
Signed-off-by: Yihong Jin <yihongjin404@gmail.com>
b1138d5 to
debb2ad
Compare
hsliuustc0106
added
ready
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
can you help check the profiling with w4a4 quant to see which part/ops benefits most
|
@david6666666 the current AMD/Intel failures are unrelated to this PR's changes (NVFP4 clamp + a test file) — CUDA (
Could a maintainer retry the AMD/Intel builds? Happy to dig into any specific one if it looks related to the change. |
…lm-project#4025) Signed-off-by: Yihong Jin <yihongjin404@gmail.com>
…lm-project#4025) Signed-off-by: Yihong Jin <yihongjin404@gmail.com>
…lm-project#4025) Signed-off-by: Yihong Jin <yihongjin404@gmail.com>
Tracks: #1854 (Continuous Quantization Support — Qwen3-Omni / ModelOpt NVFP4 W4A4 row).
Purpose
Add support for serving ModelOpt NVFP4 W4A4 quantized checkpoints of
Qwen3-Omni-30B-A3B-Instruct through vllm-omni.
The runtime change is a single defensive override in
vllm_omni/patch.pyof
ModelOptNvFp4LinearMethod.process_weights_after_loading. ModelOpt0.44's
float32 -> torch.float8_e4m3fncast of per-blockweight_scaleoccasionally emits literal NaN bytes (E4M3 encoding
0x7F/0xFF) whenthe pre-cast scale rounds above the FP8 max of 448 after the global-scale
division. A single NaN byte in any
weight_scalepropagates through theFlashInfer FP4 GEMM into the residual stream and collapses the served
model output to
!!!!. The override scanslayer.weight_scalefor NaNbytes after the original PWAL runs and clamps them to the FP8 E4M3 max in
place (via a
uint8view, since PyTorch does not implementmasked_fill_for
float8_e4m3fn).The override is defensive against checkpoints already in the wild.
Public NVFP4 W4A4 Qwen3-Omni checkpoints exported with vanilla
nvidia-modelopt==0.44.0(e.g. the precursor preview checkpointYihongJin/...-W4A4-experts-oproj-mse-previewand several community checkpoints) contain literal NaN bytes in
per-layer
weight_scaletensors (e.g. that preview ckpt has 4 NaNbytes spread across
o_projlayers 4 and 9). Without this runtimeclamp, those checkpoints currently serve only
!!!!. The clamp letsvllm-omni serve them coherently while the underlying ModelOpt 0.44 bug
gets fixed upstream.
The override self-extinguishes once vllm-omni's vllm pin moves to a
release containing the upstream fix (detected by inspecting the upstream
PWAL's
__code__.co_namesformasked_fill_), and can be disabled atload time with
VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1for diagnostics.Sibling upstream fixes against
vllm-project/vllm(the same PWAL clampin vLLM's own ModelOptNvFp4LinearMethod) and
NVIDIA/TensorRT-Model-Optimizer(clamping the FP32 -> FP8 cast at the calibration-export source so future
checkpoints ship clean) are in progress and will be linked here once
filed.
How we identified the bug
The root cause and fix in this PR are the result of a checkpoint-by-checkpoint
bisection. We published the intermediate checkpoints to HF as we went so each
step is reproducible.
...-W4A4-experts-mse...-W4A4-experts-oproj-mse-previewo_proj!!!!within the first decode stepo_projto the FP4 surface breaks it. Initial hypothesis was a "small-alpha NaN" in the FlashInfer FP4 apply path....-W4A4-full-thinker-mse-preview!!!!o_proj— any wider scope triggers it.*.weight_scalein step-2'ssafetensors, replaced with the FP8 E4M3 max byte (0x7E). Loaded the edited file on stock vllm-omni. Result: coherent text, no!!!!.weight_scale. The FlashInfer FP4 apply path was not generating it — it was reading it. Initial "small-alpha NaN in apply" hypothesis was wrong; the real bug is upstream in calibration export.modelopt/torch/quantization/qtensor/nvfp4_tensor.pyandmodelopt/torch/export/quant_utils.py. Both files call.to(torch.float8_e4m3fn)onfloat32per-block scales without first clamping to the FP8 max of 448. When the pre-cast value rounds above 448, PyTorch's FP32→FP8 cast emits NaN bytes. Added.clamp_(max=torch.finfo(torch.float8_e4m3fn).max)before each cast (6 sites total)....-W4A4-experts-oproj-mseo_proj...-W4A4-full-thinker-awqclipWhy this fix sits in vllm-omni
Steps 1-4 establish that the bug is emitted at calibration time — the
exported
safetensorsalready carry NaN bytes; the FlashInfer FP4 applypath is just a faithful reader. The load-time clamp added here is
mathematically equivalent to clamping at calibration write-side: by the
time the FP4 kernel reads
weight_scale, every position that would havebeen NaN is the FP8 E4M3 max byte (
0x7E = 448.0), regardless of whichpoint in the pipeline did the clamping. The proper fix sits upstream in
two places: at the ModelOpt cast sites (step 5 demonstrates the
calibration-side fix works) and at the corresponding PWAL in vLLM itself.
Both are planned as follow-up PRs after this one merges so we can
keep one coordinated narrative across vllm-omni / vLLM / ModelOpt.
Until those upstream fixes land in vllm-omni's pinned vLLM, any NVFP4 W4A4
Qwen3-Omni checkpoint published by the community — exported with stock
ModelOpt 0.44 — will silently serve only
!!!!on vllm-omni. Usershitting that failure will reasonably file bugs against vllm-omni, not
against ModelOpt, because the visible symptom is degenerate output from a
vllm-omni server. The load-time clamp added here keeps those community
checkpoints serving coherently in the interim.
Design choice (PWAL vs per-forward
apply()): the clamp runs once atworker init — O(num_layers) — and prevents NaN from entering the FP4 GEMM
at all. Newly calibrated, clean checkpoints (e.g. steps 6-7) pay no
runtime cost (no-op when no NaN bytes are present). The
_already_patched_upstreamcheck inspects vLLM's own PWAL for amasked_fill_call — the structure the upstream fix will take — so thisoverride skips itself once that upstream fix lands in vllm-omni's vllm
pin, and the patch block can be deleted.
Test Plan
core_model + cpumarks) attests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.pycovering the override's clamp behavior, no-op-on-finite-input behavior,
install assertion, and env-var escape hatch.
full_model + omni,H100x2) attests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.pyloading the published W4A4 checkpoint and asserting the response does
not match the
!!!!degenerate-output signature.YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip:output_len=[128, 256, 512, 1024] with
--dataset-name random.guard under sustained load.
vllm serve YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip \ --omni --port 8000Test Result
Unit tests: 8 passed, 0 skipped, 0 failed (4 base + 4 parametrized
env-var cases locking in the
{1, true, yes, on}skip semantics and{0, false}install-still-active semantics).Accuracy — Daily-Omni n=1197 + OmniBench n=200 (B200, CUDA graphs)
W4A4 stays within 2 pp of BF16 on Daily-Omni and within 0.5 pp on
OmniBench.
Hardware coverage — load-clamp ≡ calibration-clamp, verified
To verify the equivalence claim above, the same NaN-containing reproducer
checkpoint (
...-full-thinker-mse-preview)was loaded with the load-time clamp active on two additional hardware
classes covering both NVFP4 kernel paths. On Blackwell the test exercises
the same FlashInfer Cutlass NvFp4 Linear class as the B200 production
numbers; on H100 it exercises the Marlin weight-only FP4 fallback.
awqclip(clean)mse-preview(NaN)mse-preview(NaN)The 0.5 pp gap on the same kernel surface (Pro 6000 Cutlass 0.670 vs B200
Cutlass clean awq_clip 0.675) is attributable to the calibration algorithm
(mse vs awq_clip), not the clamp — consistent with the equivalence claim.
Checkpoint size
(Audio encoder, vision encoder, talker, code2wav, and lm_head stay in BF16.)
Throughput / latency — vLLM serve on B200, random workload
Synthetic random workload (
--random-input-len 256,--random-output-len 512,--num-prompts 200, no--enforce-eager). CUDA graphs captured for prefilland decode batch sizes 1..256.
W4A4 wins on every metric at concurrency >= 8. The single-stream TTFT
loss reflects the small FP4 prefill dequant overhead that BF16 native
tensor cores don't pay; once concurrency saturates the GEMM, FP4 weight
bandwidth savings dominate and W4A4 pulls ahead by 15-33% in output
tok/s and 15-31% in TPOT across batch sizes.
Kernel selection (verified at server boot)
No fallback to Marlin / Triton dequant. End-to-end native FP4.
Stability
30-minute stress at concurrency=32 against the W4A4 server:
3,200 successful requests, 0 errors, 0
!!!!collapses — the patchholds under sustained load (this is the regression guard for the failure
mode that motivated this PR).
Hardware footnote
All B200 numbers above use a single GPU at
--gpu-memory-utilization 0.85.B100 / B200 with NVLink + TP=2 will scale differently and is intentionally
out of scope for this PR.
The PR's earlier Pro 6000 Blackwell WS draft is preserved in the commit
history; we've replaced it with B200 numbers because (a) B200 is the
production target, and (b) the Pro 6000's GDDR7 memory makes it
bandwidth-bound on this workload, masking the FP4 throughput win that
HBM3e exposes.
How the +4 pp early reading turned into a real -1.9 pp number
An earlier in-process smoke check on Daily-Omni n=50 had W4A4 at
0.76 vs BF16 0.72 (apparent +4 pp for the quantized model). The full
n=1197 run lands the same W4A4 ckpt at 0.675 vs BF16 0.694
— well within the binomial 95% confidence band the n=50 slice allows.
Recording this here so reviewers don't have to triage the same surprise.
Files
vllm_omni/patch.py(+119 / -88 net): PWAL-time NaN-clamp overridewith WHY / WHY NOT / SCOPE / FRAGILITY / TODO block comment, env-var
escape hatch, self-extinguishing upstream-already-patched check, and
install assert.
tests/model_executor/quantization/__init__.py: new test subpackage.tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py(+176): CPU-only unit test of the override.
tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py(+76): hardware-gated e2e smoke test asserting non-degenerate output
on the published W4A4 checkpoint.
docs/user_guide/quantization/modelopt.md(+37): adds Qwen3-Omni W4A4section with HF repo link, hardware requirement, calibration recipe
pointer, and a note describing the ModelOpt 0.44 NaN regression and
the workaround installed by
vllm_omni.patch.