From 9e08ba55d76f57a5e85f8c8e8e99a37a3d5fd95d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 02:15:31 +0000 Subject: [PATCH 1/8] b12x MoE: opt-in NVFP4-activation prefill for W4A16 experts, decode rows stay W4A16 With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, b12x serves W4A16 MoE calls of at least that many tokens with NVFP4 activations over the packed W4A16 weights. vLLM now passes the checkpoint's activation global scales in W4A16 mode for that path (the shared input scale is the smallest global scale, i.e. the widest calibrated range; intermediate scales stay per expert), skips layers whose scales are missing or invalid (they stay W4A16), and splits a mixed step so its leading decode rows always run on the exact W4A16 path and only the prefill rows take the faster one. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-b12x-w4a16-a4-prefill.json | 16 +++ vllm/model_executor/layers/fused_moe/b12x.py | 114 ++++++++++++++++++- 2 files changed, 129 insertions(+), 1 deletion(-) create mode 100644 .lil/changes/vllm-b12x-w4a16-a4-prefill.json diff --git a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json new file mode 100644 index 000000000000..4a82b4504c18 --- /dev/null +++ b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-b12x-w4a16-a4-prefill", + "category": "feature", + "summary": "Opt-in faster prefill for W4A16 MoE: NVFP4 activations for large prefill steps, decode stays BF16", + "models": ["GLM-5.3-Flash"], + "compatibility": "Opt-in through B12X_W4A16_A4_PREFILL_MIN_TOKENS; without it nothing changes.", + "details": [ + "With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, the b12x MoE in W4A16 mode receives the checkpoint's activation global scales, and b12x runs MoE calls of at least that many tokens with NVFP4 activations over the same packed W4A16 weights (B12X_W4A16_A4_PREFILL_TERMS=2 uses a two-term NVFP4 activation format with about 8-bit precision).", + "Decode rows always stay on the exact W4A16 path: a mixed step runs its leading decode rows through W4A16 and only its prefill rows through the faster path.", + "Layers without valid activation scales (for example a draft layer without calibration) stay W4A16." + ], + "authors": ["voipmonitor"], + "pull_requests": [], + "requires": [] +} diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index af958dca03ba..401227fffd56 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """b12x modular tensor-parallel fused MoE backend.""" +import os import weakref from collections.abc import Sequence from dataclasses import dataclass @@ -226,6 +227,54 @@ def _replace_parameter_with_empty( return getattr(layer, name) +def _w4a16_a4_prefill_min_tokens() -> int: + return max(int(os.environ.get("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "0") or 0), 0) + + +def _w4a16_a4_prefill_enabled() -> bool: + """b12x reads B12X_W4A16_A4_PREFILL_MIN_TOKENS at plan time; it needs the + activation global scales the A16 path otherwise ignores.""" + return _w4a16_a4_prefill_min_tokens() > 0 + + +_A4_SPLIT_LOGGED = False + + +def _num_leading_decode_tokens() -> int: + """Leading decode-row count of the current decodes-first batch (0 if unknown). + + Only called for steps at or above the A4 threshold, which never run inside a + captured CUDA graph: one host read of query_start_loc per step, cached on the + forward context. Requests with query length <= 4 (1 + 3 MTP tokens) count as + decode rows; KDA's reorder keeps them first. + """ + from vllm.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return 0 + context = get_forward_context() + cached = getattr(context, "_b12x_a4_decode_rows", None) + if cached is not None: + return cached + rows = 0 + metadata = context.attn_metadata + if isinstance(metadata, dict): + for value in metadata.values(): + starts = getattr(value, "query_start_loc", None) + if isinstance(starts, torch.Tensor) and hasattr(value, "num_actual_tokens"): + bounds = starts.tolist() + for lo, hi in zip(bounds, bounds[1:]): + if not 0 < hi - lo <= 4: + break + rows = hi + break + try: + context._b12x_a4_decode_rows = rows + except Exception: + pass + return rows + + def _normalize_expert_scale(scale: torch.Tensor) -> torch.Tensor: if scale.ndim == 2: if scale.shape[1] not in (1, 2): @@ -389,6 +438,36 @@ def _prepare_experts( raise ValueError("b12x NVFP4 MoE requires activation global scales") a1_gscale = _normalize_expert_scale(self.a1_gscale).to(w1.device) a2_gscale = _normalize_expert_scale(self.a2_gscale).to(w2.device) + elif ( + quant_mode == "w4a16" + and _w4a16_a4_prefill_enabled() + and self.a1_gscale is not None + and self.a2_gscale is not None + ): + # A4 prefill over the W4A16 weights quantizes each token once: share + # the widest calibrated input range (smallest global scale) across + # experts; the intermediate requant keeps per-expert scales. + a1 = _normalize_expert_scale(self.a1_gscale).to(w1.device) + a2 = _normalize_expert_scale(self.a2_gscale).to(w2.device) + if all( + bool((torch.isfinite(t) & (t > 0)).all()) for t in (a1, a2) + ): + a1_gscale = a1.amin().reshape(1) + a2_gscale = a2 + logger.info( + "b12x W4A16: A4 prefill for calls >= %s tokens (shared input " + "global scale %.4g, %d intermediate scales)", + os.environ.get("B12X_W4A16_A4_PREFILL_MIN_TOKENS"), + float(a1_gscale[0]), + int(a2_gscale.numel()), + ) + else: + # No calibrated activation scales (e.g. a draft layer): stay A16. + a1_gscale = a2_gscale = None + logger.warning( + "b12x W4A16: layer has no valid activation scales; A4 prefill " + "disabled for it (A16 only)" + ) else: a1_gscale = unit_scale a2_gscale = unit_scale @@ -446,7 +525,12 @@ def _refresh_quant_config(self, layer: torch.nn.Module) -> None: self.quant_config._w1.alpha_or_gscale = layer.w13_weight_scale_2 self.quant_config._w2.alpha_or_gscale = layer.w2_weight_scale_2 - if self._quant_mode in ("nvfp4", "w4a8_nvfp4"): + if self._quant_mode in ("nvfp4", "w4a8_nvfp4") or ( + self._quant_mode == "w4a16" + and _w4a16_a4_prefill_enabled() + and getattr(layer, "w13_input_scale", None) is not None + and getattr(layer, "w2_input_scale", None) is not None + ): self.quant_config._a1.alpha_or_gscale = 1.0 / layer.w13_input_scale self.quant_config._a2.alpha_or_gscale = 1.0 / layer.w2_input_scale @@ -881,6 +965,34 @@ def apply( if workspace2 is None or not workspace2.is_contiguous(): raise ValueError("b12x MoE requires contiguous caller-owned workspace2") scratch = workspace2.view(-1).view(torch.uint8) + tokens = int(hidden_states.shape[0]) + a4_min = _w4a16_a4_prefill_min_tokens() + if a4_min and tokens >= a4_min and self._quant_mode == "w4a16": + # Decode always stays A16: a mixed step runs its leading decode rows + # through W4A16 and only the prefill rows through the A4 path. + decode_rows = _num_leading_decode_tokens() + if 0 < decode_rows < tokens: + global _A4_SPLIT_LOGGED + if not _A4_SPLIT_LOGGED: + _A4_SPLIT_LOGGED = True + logger.info( + "b12x W4A16 A4 prefill split: %d decode rows A16, %d prefill rows", + decode_rows, + tokens - decode_rows, + ) + for lo, hi in ((0, decode_rows), (decode_rows, tokens)): + binding = _require_b12x_fused_moe().bind( + plan, + scratch=scratch, + a=hidden_states[lo:hi], + experts=prepared, + topk_weights=topk_weights[lo:hi], + topk_ids=topk_ids[lo:hi], + output=output[lo:hi], + input_scales_static=True, + ) + _require_b12x_fused_moe().run(binding=binding) + return binding = _require_b12x_fused_moe().bind( plan, scratch=scratch, From 023112b17b77f006153a7b9ab24c223e0ff55b74 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 02:20:55 +0000 Subject: [PATCH 2/8] tests: host contract of the opt-in b12x W4A16 A4 prefill (threshold, decode rows, split) Co-Authored-By: Claude Opus 5.5 --- tests/kernels/moe/test_b12x_a4_prefill.py | 124 ++++++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 tests/kernels/moe/test_b12x_a4_prefill.py diff --git a/tests/kernels/moe/test_b12x_a4_prefill.py b/tests/kernels/moe/test_b12x_a4_prefill.py new file mode 100644 index 000000000000..8ac6734bd793 --- /dev/null +++ b/tests/kernels/moe/test_b12x_a4_prefill.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-side contract of the opt-in A4 prefill for b12x W4A16 MoE. + +Covers the threshold variable, the leading decode-row count read from the +forward context, and the split that keeps decode rows on the W4A16 path. The +b12x library itself is replaced by a recorder; no GPU is needed. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.forward_context as forward_context +import vllm.model_executor.layers.fused_moe.b12x as b12x +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts + + +def test_min_tokens_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", raising=False) + assert b12x._w4a16_a4_prefill_min_tokens() == 0 + assert not b12x._w4a16_a4_prefill_enabled() + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + assert b12x._w4a16_a4_prefill_min_tokens() == 1536 + assert b12x._w4a16_a4_prefill_enabled() + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "-4") + assert b12x._w4a16_a4_prefill_min_tokens() == 0 + + +def _fake_context(monkeypatch, starts): + metadata = SimpleNamespace( + query_start_loc=torch.tensor(starts, dtype=torch.int32), num_actual_tokens=starts[-1] + ) + context = SimpleNamespace(attn_metadata={"layer.0": metadata}) + monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True) + monkeypatch.setattr(forward_context, "get_forward_context", lambda: context) + return context + + +@pytest.mark.parametrize( + "starts, rows", + [ + ([0, 4, 8, 12, 2312], 12), # three MTP3 decode requests, then a prefill chunk + ([0, 2300], 0), # prefill only + ([0, 1, 2, 3], 3), # decode only + ([0, 4, 2304, 2308], 4), # a decode request after the prefill is not leading + ], +) +def test_leading_decode_rows(monkeypatch: pytest.MonkeyPatch, starts, rows) -> None: + context = _fake_context(monkeypatch, starts) + assert b12x._num_leading_decode_tokens() == rows + # Cached on the forward context: one host read per step. + assert context._b12x_a4_decode_rows == rows + context.attn_metadata = {} + assert b12x._num_leading_decode_tokens() == rows + + +def test_leading_decode_rows_without_context(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: False) + assert b12x._num_leading_decode_tokens() == 0 + + +class _Recorder: + def __init__(self) -> None: + self.calls: list[tuple[int, int]] = [] + + def bind(self, plan, *, scratch, a, experts, topk_weights, topk_ids, output, **kwargs): + del plan, scratch, experts, kwargs + assert a.shape[0] == topk_weights.shape[0] == topk_ids.shape[0] == output.shape[0] + return SimpleNamespace(rows=int(a.shape[0]), first=int(a.data_ptr())) + + def run(self, *, binding) -> None: + self.calls.append((binding.rows, binding.first)) + + +def _experts(monkeypatch, quant_mode: str): + recorder = _Recorder() + monkeypatch.setattr(b12x, "_require_b12x_fused_moe", lambda: recorder) + experts = object.__new__(B12xExperts) + experts._apply_router_weight_on_input = False + experts._quant_mode = quant_mode + experts._prepared = lambda: object() + experts._plan_for_tokens = lambda tokens, **kwargs: object() + return experts, recorder + + +def _apply(experts, tokens): + hidden = torch.zeros(tokens, 8, dtype=torch.bfloat16) + experts.apply( + output=torch.zeros_like(hidden), hidden_states=hidden, w1=None, w2=None, + topk_weights=torch.ones(tokens, 2), topk_ids=torch.zeros(tokens, 2, dtype=torch.int32), + activation=MoEActivation.SILU, global_num_experts=4, expert_map=None, a1q_scale=None, + a2_scale=None, workspace13=None, workspace2=torch.zeros(64, dtype=torch.uint8), + expert_tokens_meta=None, apply_router_weight_on_input=False, + ) + return hidden + + +def test_mixed_step_splits_decode_rows(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + _fake_context(monkeypatch, [0, 4, 8, 2312]) + experts, recorder = _experts(monkeypatch, "w4a16") + hidden = _apply(experts, 2312) + row = hidden.element_size() * hidden.shape[1] + assert recorder.calls == [(8, hidden.data_ptr()), (2304, hidden.data_ptr() + 8 * row)] + + +@pytest.mark.parametrize( + "quant_mode, threshold, tokens", + [ + ("w4a16", "1536", 64), # below the threshold: one call + ("w4a16", "0", 2312), # option off + ("nvfp4", "1536", 2312), # not a W4A16 plan + ], +) +def test_no_split(monkeypatch: pytest.MonkeyPatch, quant_mode, threshold, tokens) -> None: + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", threshold) + _fake_context(monkeypatch, [0, 4, tokens]) + experts, recorder = _experts(monkeypatch, quant_mode) + _apply(experts, tokens) + assert [rows for rows, _ in recorder.calls] == [tokens] From e84fd3e675a5e07a38e6ddec5339a587934be8ec Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:51:39 +0000 Subject: [PATCH 3/8] b12x MoE: lint fixes for the A4 prefill split Co-Authored-By: Claude Opus 5.5 --- tests/kernels/moe/test_b12x_a4_prefill.py | 40 +++++++++++++++----- vllm/model_executor/layers/fused_moe/b12x.py | 12 +++--- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/tests/kernels/moe/test_b12x_a4_prefill.py b/tests/kernels/moe/test_b12x_a4_prefill.py index 8ac6734bd793..16ad8f9f17f9 100644 --- a/tests/kernels/moe/test_b12x_a4_prefill.py +++ b/tests/kernels/moe/test_b12x_a4_prefill.py @@ -32,7 +32,8 @@ def test_min_tokens_env(monkeypatch: pytest.MonkeyPatch) -> None: def _fake_context(monkeypatch, starts): metadata = SimpleNamespace( - query_start_loc=torch.tensor(starts, dtype=torch.int32), num_actual_tokens=starts[-1] + query_start_loc=torch.tensor(starts, dtype=torch.int32), + num_actual_tokens=starts[-1], ) context = SimpleNamespace(attn_metadata={"layer.0": metadata}) monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True) @@ -67,9 +68,13 @@ class _Recorder: def __init__(self) -> None: self.calls: list[tuple[int, int]] = [] - def bind(self, plan, *, scratch, a, experts, topk_weights, topk_ids, output, **kwargs): + def bind( + self, plan, *, scratch, a, experts, topk_weights, topk_ids, output, **kwargs + ): del plan, scratch, experts, kwargs - assert a.shape[0] == topk_weights.shape[0] == topk_ids.shape[0] == output.shape[0] + assert ( + a.shape[0] == topk_weights.shape[0] == topk_ids.shape[0] == output.shape[0] + ) return SimpleNamespace(rows=int(a.shape[0]), first=int(a.data_ptr())) def run(self, *, binding) -> None: @@ -90,11 +95,21 @@ def _experts(monkeypatch, quant_mode: str): def _apply(experts, tokens): hidden = torch.zeros(tokens, 8, dtype=torch.bfloat16) experts.apply( - output=torch.zeros_like(hidden), hidden_states=hidden, w1=None, w2=None, - topk_weights=torch.ones(tokens, 2), topk_ids=torch.zeros(tokens, 2, dtype=torch.int32), - activation=MoEActivation.SILU, global_num_experts=4, expert_map=None, a1q_scale=None, - a2_scale=None, workspace13=None, workspace2=torch.zeros(64, dtype=torch.uint8), - expert_tokens_meta=None, apply_router_weight_on_input=False, + output=torch.zeros_like(hidden), + hidden_states=hidden, + w1=None, + w2=None, + topk_weights=torch.ones(tokens, 2), + topk_ids=torch.zeros(tokens, 2, dtype=torch.int32), + activation=MoEActivation.SILU, + global_num_experts=4, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=None, + workspace2=torch.zeros(64, dtype=torch.uint8), + expert_tokens_meta=None, + apply_router_weight_on_input=False, ) return hidden @@ -105,7 +120,10 @@ def test_mixed_step_splits_decode_rows(monkeypatch: pytest.MonkeyPatch) -> None: experts, recorder = _experts(monkeypatch, "w4a16") hidden = _apply(experts, 2312) row = hidden.element_size() * hidden.shape[1] - assert recorder.calls == [(8, hidden.data_ptr()), (2304, hidden.data_ptr() + 8 * row)] + assert recorder.calls == [ + (8, hidden.data_ptr()), + (2304, hidden.data_ptr() + 8 * row), + ] @pytest.mark.parametrize( @@ -116,7 +134,9 @@ def test_mixed_step_splits_decode_rows(monkeypatch: pytest.MonkeyPatch) -> None: ("nvfp4", "1536", 2312), # not a W4A16 plan ], ) -def test_no_split(monkeypatch: pytest.MonkeyPatch, quant_mode, threshold, tokens) -> None: +def test_no_split( + monkeypatch: pytest.MonkeyPatch, quant_mode, threshold, tokens +) -> None: monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", threshold) _fake_context(monkeypatch, [0, 4, tokens]) experts, recorder = _experts(monkeypatch, quant_mode) diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 401227fffd56..1dfd5375e8fa 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """b12x modular tensor-parallel fused MoE backend.""" +import contextlib import os import weakref from collections.abc import Sequence @@ -268,10 +269,8 @@ def _num_leading_decode_tokens() -> int: break rows = hi break - try: + with contextlib.suppress(AttributeError): context._b12x_a4_decode_rows = rows - except Exception: - pass return rows @@ -449,9 +448,7 @@ def _prepare_experts( # experts; the intermediate requant keeps per-expert scales. a1 = _normalize_expert_scale(self.a1_gscale).to(w1.device) a2 = _normalize_expert_scale(self.a2_gscale).to(w2.device) - if all( - bool((torch.isfinite(t) & (t > 0)).all()) for t in (a1, a2) - ): + if all(bool((torch.isfinite(t) & (t > 0)).all()) for t in (a1, a2)): a1_gscale = a1.amin().reshape(1) a2_gscale = a2 logger.info( @@ -976,7 +973,8 @@ def apply( if not _A4_SPLIT_LOGGED: _A4_SPLIT_LOGGED = True logger.info( - "b12x W4A16 A4 prefill split: %d decode rows A16, %d prefill rows", + "b12x W4A16 A4 prefill split: %d decode rows stay A16, " + "%d prefill rows", decode_rows, tokens - decode_rows, ) From d6894b55302d5fd9b9be37f903d23ad3f2109542 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 04:07:50 +0000 Subject: [PATCH 4/8] release: the A4 prefill fragment as published in the Karmic beta Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-b12x-w4a16-a4-prefill.json | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json index 4a82b4504c18..b623221aa416 100644 --- a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json +++ b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json @@ -3,14 +3,22 @@ "id": "vllm-b12x-w4a16-a4-prefill", "category": "feature", "summary": "Opt-in faster prefill for W4A16 MoE: NVFP4 activations for large prefill steps, decode stays BF16", - "models": ["GLM-5.3-Flash"], + "models": [ + "GLM-5.3-Flash" + ], "compatibility": "Opt-in through B12X_W4A16_A4_PREFILL_MIN_TOKENS; without it nothing changes.", "details": [ - "With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, the b12x MoE in W4A16 mode receives the checkpoint's activation global scales, and b12x runs MoE calls of at least that many tokens with NVFP4 activations over the same packed W4A16 weights (B12X_W4A16_A4_PREFILL_TERMS=2 uses a two-term NVFP4 activation format with about 8-bit precision).", + "With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, the b12x MoE in W4A16 mode receives the checkpoint's activation global scales, and b12x runs MoE calls of at least that many tokens with NVFP4 activations over the same packed W4A16 weights. GLM-5.3-Flash TP2/DCP2 MTP3 on two RTX PRO 6000 Max-Q with a 1536-token threshold: prefill 9,528 and 9,784 tok/s at 8K and 32K tokens versus 7,490 and 7,728 with W4A16, decode unchanged.", "Decode rows always stay on the exact W4A16 path: a mixed step runs its leading decode rows through W4A16 and only its prefill rows through the faster path.", "Layers without valid activation scales (for example a draft layer without calibration) stay W4A16." ], - "authors": ["voipmonitor"], - "pull_requests": [], - "requires": [] + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 970 + ], + "requires": [ + "b12x-w4a16-a4-prefill" + ] } From a0c4fade2e28b33a6130a35151049a6e6e298037 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:07:19 +0000 Subject: [PATCH 5/8] modelopt: zero-initialize NVFP4 MoE input scales W4A16_NVFP4 experts store no input scales (GLM-5.3-Flash keeps its MTP layer's routed experts weight-only), so the registered parameters kept whatever torch.empty returned. With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, the b12x MoE takes positive finite input scales for calibrated ones, and the draft layer could run its prefill with NVFP4 activations under a garbage global scale (5.2e-32 in a GLM Spark TP2 run). Zero reads as absent: the existing guard keeps such layers on W4A16. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit d7af3ad174040404bfd82b941abe9c72f4529d5a) --- tests/quantization/test_modelopt.py | 54 +++++++++++++++++++ .../layers/quantization/modelopt.py | 7 ++- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 92cde3013ee9..0d09ae0ba6ff 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -1035,6 +1035,60 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16( assert kwargs["activation_key"] is kNvfp4Dynamic +@pytest.mark.parametrize("quant_method", ["NVFP4", "W4A16_NVFP4"]) +def test_modelopt_nvfp4_moe_input_scales_start_at_zero(dist_init, quant_method): + """Input scales the checkpoint does not store must read as absent. + + W4A16_NVFP4 experts (GLM-5.3's MTP layer) have no input scales on disk, so + the registered parameters keep their initial value. The b12x W4A16 A4 + prefill treats positive finite scales as calibrated; uninitialized memory + could enable it with a garbage global scale. + """ + from vllm.model_executor.layers.quantization import modelopt + + config = modelopt.ModelOptNvFp4Config( + quant_method=quant_method, + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + group_size=16, + ) + with ( + patch.object( + modelopt, + "select_nvfp4_moe_backend", + return_value=(MagicMock(), MagicMock()), + ), + patch.object( + modelopt, "is_global_sf_supported_for_nvfp4_backend", return_value=False + ), + ): + moe = modelopt.ModelOptNvFp4FusedMoE(config, MagicMock()) + + empty = torch.empty + + def garbage_empty(*args, **kwargs): + # Stand-in for uninitialized memory: positive and finite. + tensor = empty(*args, **kwargs) + if tensor.dtype == torch.float32: + tensor.fill_(5.2e-32) + return tensor + + layer = torch.nn.Module() + with patch.object(modelopt.torch, "empty", garbage_empty): + moe.create_weights( + layer, + num_experts=4, + hidden_size=64, + intermediate_size_per_partition=32, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + + assert torch.equal(layer.w13_input_scale, torch.zeros(4, 2)) + assert torch.equal(layer.w2_input_scale, torch.zeros(4)) + + @pytest.mark.parametrize( "per_layer_algo, expected_weight, expected_activation", [ diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 91eaf10eb280..99000450d740 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -961,8 +961,11 @@ def create_weights( global_sf_num_experts = ( global_num_experts if self.use_global_sf else num_experts ) + # Zero, not uninitialized: W4A16_NVFP4 layers (such as GLM-5.3 MTP + # experts) store no input scales, and the b12x W4A16 A4 prefill takes + # any positive finite value here for a calibrated scale. w13_input_scale = PerTensorScaleParameter( - data=torch.empty( + data=torch.zeros( global_sf_num_experts, w13_num_shards, dtype=torch.float32, @@ -972,7 +975,7 @@ def create_weights( layer.register_parameter("w13_input_scale", w13_input_scale) w2_input_scale = PerTensorScaleParameter( - data=torch.empty(global_sf_num_experts, dtype=torch.float32), + data=torch.zeros(global_sf_num_experts, dtype=torch.float32), weight_loader=weight_loader, ) layer.register_parameter("w2_input_scale", w2_input_scale) From a76c00c329404b7faed6631025e9dfbb1644ac42 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:08:01 +0000 Subject: [PATCH 6/8] release: fragment for the zero-initialized NVFP4 MoE input scales (#976) Co-Authored-By: Claude Opus 5.5 (cherry picked from commit dbece926c942b894e2967c3f6961671870fb6aee) --- .../vllm-nvfp4-moe-unloaded-input-scales.json | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-nvfp4-moe-unloaded-input-scales.json diff --git a/.lil/changes/vllm-nvfp4-moe-unloaded-input-scales.json b/.lil/changes/vllm-nvfp4-moe-unloaded-input-scales.json new file mode 100644 index 000000000000..16dae1f5ad0d --- /dev/null +++ b/.lil/changes/vllm-nvfp4-moe-unloaded-input-scales.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-nvfp4-moe-unloaded-input-scales", + "category": "fix", + "summary": "MoE layers without stored activation scales no longer enable the NVFP4-activation prefill from uninitialized memory", + "models": [ + "GLM-5.3-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "GLM-5.3-Flash checkpoints keep the MTP layer's routed experts weight-only (W4A16_NVFP4) without input scales. Their scale parameters were uninitialized, so with B12X_W4A16_A4_PREFILL_MIN_TOKENS set the draft layer could prefill with NVFP4 activations under a garbage global scale. They now start at zero, which the A4 prefill treats as absent: the draft layer stays W4A16.", + "MTP acceptance after a 13.4K-token prompt was unchanged in the measured case; the fix removes the dependence on uninitialized memory." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 976 + ], + "requires": [] +} From b94644d1327793ff962fbfe8c1bbbd1038ff30d1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:53:16 +0000 Subject: [PATCH 7/8] b12x MoE: choose the A4 prefill by row type, not by call size With B12X_W4A16_A4_PREFILL_MIN_TOKENS set, the decode rows of a step run W4A16 and its prefill rows run A4, whatever the call size (b12x bind(a4_prefill=...)). Before, prefill rows took A4 only when the call had at least the threshold's tokens, so a prompt's precision depended on how the scheduler chunked and batched it. A captured CUDA graph keeps W4A16 for all its rows, and layers without activation scales (the MTP draft layer) are not split at all. A GDN/KDA layer counts prefills on the host, so a decode-only step finds its decode rows without a device read; steps with prefill rows read query_start_loc once, as before. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit df3d9833a4761b6e4132fed8350552e844be98be) --- tests/kernels/moe/test_b12x_a4_prefill.py | 112 +++++++++++++------ vllm/model_executor/layers/fused_moe/b12x.py | 102 ++++++++++------- 2 files changed, 143 insertions(+), 71 deletions(-) diff --git a/tests/kernels/moe/test_b12x_a4_prefill.py b/tests/kernels/moe/test_b12x_a4_prefill.py index 16ad8f9f17f9..a859ce3b70d3 100644 --- a/tests/kernels/moe/test_b12x_a4_prefill.py +++ b/tests/kernels/moe/test_b12x_a4_prefill.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Host-side contract of the opt-in A4 prefill for b12x W4A16 MoE. -Covers the threshold variable, the leading decode-row count read from the -forward context, and the split that keeps decode rows on the W4A16 path. The +Covers the enabling variable, the leading decode-row count read from the +forward context, and the per-row-type choice: decode rows run W4A16, prefill +rows A4. The b12x library itself is replaced by a recorder; no GPU is needed. """ @@ -30,12 +31,19 @@ def test_min_tokens_env(monkeypatch: pytest.MonkeyPatch) -> None: assert b12x._w4a16_a4_prefill_min_tokens() == 0 -def _fake_context(monkeypatch, starts): - metadata = SimpleNamespace( - query_start_loc=torch.tensor(starts, dtype=torch.int32), - num_actual_tokens=starts[-1], - ) - context = SimpleNamespace(attn_metadata={"layer.0": metadata}) +def _fake_context(monkeypatch, starts, gdn_prefills=None): + metadata = { + "layer.0": SimpleNamespace( + query_start_loc=torch.tensor(starts, dtype=torch.int32), + num_actual_tokens=starts[-1], + ) + } + if gdn_prefills is not None: + # A GDN/KDA layer's host-side counts. + metadata["layer.1"] = SimpleNamespace( + num_prefills=gdn_prefills, num_spec_decode_tokens=starts[-1] + ) + context = SimpleNamespace(attn_metadata=metadata) monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True) monkeypatch.setattr(forward_context, "get_forward_context", lambda: context) return context @@ -52,42 +60,54 @@ def _fake_context(monkeypatch, starts): ) def test_leading_decode_rows(monkeypatch: pytest.MonkeyPatch, starts, rows) -> None: context = _fake_context(monkeypatch, starts) - assert b12x._num_leading_decode_tokens() == rows + assert b12x._num_leading_decode_tokens(starts[-1]) == rows # Cached on the forward context: one host read per step. assert context._b12x_a4_decode_rows == rows context.attn_metadata = {} - assert b12x._num_leading_decode_tokens() == rows + assert b12x._num_leading_decode_tokens(starts[-1]) == rows + + +def test_a_decode_only_step_needs_no_device_read(monkeypatch: pytest.MonkeyPatch): + context = _fake_context(monkeypatch, [0, 4, 8], gdn_prefills=0) + context.attn_metadata["layer.0"].query_start_loc = None # would need a sync + assert b12x._num_leading_decode_tokens(8) == 8 def test_leading_decode_rows_without_context(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: False) - assert b12x._num_leading_decode_tokens() == 0 + assert b12x._num_leading_decode_tokens(8) == 0 class _Recorder: def __init__(self) -> None: - self.calls: list[tuple[int, int]] = [] + self.calls: list[tuple[int, int, bool | None]] = [] def bind( self, plan, *, scratch, a, experts, topk_weights, topk_ids, output, **kwargs ): - del plan, scratch, experts, kwargs + del plan, scratch, experts assert ( a.shape[0] == topk_weights.shape[0] == topk_ids.shape[0] == output.shape[0] ) - return SimpleNamespace(rows=int(a.shape[0]), first=int(a.data_ptr())) + return SimpleNamespace( + rows=int(a.shape[0]), + first=int(a.data_ptr()), + a4=kwargs.get("a4_prefill"), + ) def run(self, *, binding) -> None: - self.calls.append((binding.rows, binding.first)) + self.calls.append((binding.rows, binding.first, binding.a4)) -def _experts(monkeypatch, quant_mode: str): +def _experts(monkeypatch, quant_mode: str, a4_scales: bool = True): recorder = _Recorder() monkeypatch.setattr(b12x, "_require_b12x_fused_moe", lambda: recorder) + monkeypatch.setattr(b12x, "_is_current_stream_capturing", lambda: False) experts = object.__new__(B12xExperts) experts._apply_router_weight_on_input = False experts._quant_mode = quant_mode - experts._prepared = lambda: object() + prepared = SimpleNamespace(_impl=SimpleNamespace(a4_prefill_scales=a4_scales)) + experts._prepared = lambda: prepared experts._plan_for_tokens = lambda tokens, **kwargs: object() return experts, recorder @@ -114,31 +134,59 @@ def _apply(experts, tokens): return hidden -def test_mixed_step_splits_decode_rows(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("tokens", [2312, 64]) +def test_decode_rows_stay_w4a16_and_prefill_rows_run_a4( + monkeypatch: pytest.MonkeyPatch, tokens +) -> None: + """Small and large prefills alike: the row type decides, not the size.""" monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") - _fake_context(monkeypatch, [0, 4, 8, 2312]) + _fake_context(monkeypatch, [0, 4, 8, tokens]) experts, recorder = _experts(monkeypatch, "w4a16") - hidden = _apply(experts, 2312) + hidden = _apply(experts, tokens) row = hidden.element_size() * hidden.shape[1] assert recorder.calls == [ - (8, hidden.data_ptr()), - (2304, hidden.data_ptr() + 8 * row), + (8, hidden.data_ptr(), False), + (tokens - 8, hidden.data_ptr() + 8 * row, True), ] +def test_a_prefill_only_step_runs_a4(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + _fake_context(monkeypatch, [0, 300]) + experts, recorder = _experts(monkeypatch, "w4a16") + _apply(experts, 300) + assert [(rows, a4) for rows, _, a4 in recorder.calls] == [(300, True)] + + +def test_a_decode_only_step_stays_w4a16(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + _fake_context(monkeypatch, [0, 4, 8, 12], gdn_prefills=0) + experts, recorder = _experts(monkeypatch, "w4a16") + _apply(experts, 12) + assert [(rows, a4) for rows, _, a4 in recorder.calls] == [(12, False)] + + +def test_a_captured_graph_stays_w4a16(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + experts, recorder = _experts(monkeypatch, "w4a16") + monkeypatch.setattr(b12x, "_is_current_stream_capturing", lambda: True) + _apply(experts, 32) + assert [(rows, a4) for rows, _, a4 in recorder.calls] == [(32, False)] + + @pytest.mark.parametrize( - "quant_mode, threshold, tokens", + "quant_mode, threshold, a4_scales", [ - ("w4a16", "1536", 64), # below the threshold: one call - ("w4a16", "0", 2312), # option off - ("nvfp4", "1536", 2312), # not a W4A16 plan + ("w4a16", "0", True), # option off + ("nvfp4", "1536", True), # not a W4A16 plan + ("w4a16", "1536", False), # no activation scales (the MTP draft layer) ], ) -def test_no_split( - monkeypatch: pytest.MonkeyPatch, quant_mode, threshold, tokens +def test_no_choice_without_an_a4_path( + monkeypatch: pytest.MonkeyPatch, quant_mode, threshold, a4_scales ) -> None: monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", threshold) - _fake_context(monkeypatch, [0, 4, tokens]) - experts, recorder = _experts(monkeypatch, quant_mode) - _apply(experts, tokens) - assert [rows for rows, _ in recorder.calls] == [tokens] + _fake_context(monkeypatch, [0, 4, 2312]) + experts, recorder = _experts(monkeypatch, quant_mode, a4_scales) + _apply(experts, 2312) + assert [(rows, a4) for rows, _, a4 in recorder.calls] == [(2312, None)] diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 1dfd5375e8fa..c77e33fcfa42 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -233,21 +233,23 @@ def _w4a16_a4_prefill_min_tokens() -> int: def _w4a16_a4_prefill_enabled() -> bool: - """b12x reads B12X_W4A16_A4_PREFILL_MIN_TOKENS at plan time; it needs the - activation global scales the A16 path otherwise ignores.""" + """b12x compiles the A4 prefill path when B12X_W4A16_A4_PREFILL_MIN_TOKENS + is set; it needs the activation global scales the A16 path otherwise + ignores.""" return _w4a16_a4_prefill_min_tokens() > 0 _A4_SPLIT_LOGGED = False -def _num_leading_decode_tokens() -> int: - """Leading decode-row count of the current decodes-first batch (0 if unknown). +def _num_leading_decode_tokens(tokens: int) -> int: + """Leading decode-row count of the current decodes-first batch. - Only called for steps at or above the A4 threshold, which never run inside a - captured CUDA graph: one host read of query_start_loc per step, cached on the - forward context. Requests with query length <= 4 (1 + 3 MTP tokens) count as - decode rows; KDA's reorder keeps them first. + Requests with query length <= 4 (1 + 3 MTP tokens) count as decode rows; + KDA's reorder keeps them first. A GDN/KDA layer counts prefills on the + host, so a decode-only step needs no device read. Otherwise one host read + of query_start_loc per step, cached on the forward context. Never called + inside a captured CUDA graph. 0 without a forward context. """ from vllm.forward_context import get_forward_context, is_forward_context_available @@ -260,15 +262,25 @@ def _num_leading_decode_tokens() -> int: rows = 0 metadata = context.attn_metadata if isinstance(metadata, dict): - for value in metadata.values(): - starts = getattr(value, "query_start_loc", None) - if isinstance(starts, torch.Tensor) and hasattr(value, "num_actual_tokens"): - bounds = starts.tolist() - for lo, hi in zip(bounds, bounds[1:]): - if not 0 < hi - lo <= 4: - break - rows = hi - break + prefills = [ + value.num_prefills + for value in metadata.values() + if hasattr(value, "num_spec_decode_tokens") + ] + if prefills and not any(prefills): + rows = tokens + else: + for value in metadata.values(): + starts = getattr(value, "query_start_loc", None) + if isinstance(starts, torch.Tensor) and hasattr( + value, "num_actual_tokens" + ): + bounds = starts.tolist() + for lo, hi in zip(bounds, bounds[1:]): + if not 0 < hi - lo <= 4: + break + rows = hi + break with contextlib.suppress(AttributeError): context._b12x_a4_decode_rows = rows return rows @@ -963,34 +975,46 @@ def apply( raise ValueError("b12x MoE requires contiguous caller-owned workspace2") scratch = workspace2.view(-1).view(torch.uint8) tokens = int(hidden_states.shape[0]) - a4_min = _w4a16_a4_prefill_min_tokens() - if a4_min and tokens >= a4_min and self._quant_mode == "w4a16": - # Decode always stays A16: a mixed step runs its leading decode rows - # through W4A16 and only the prefill rows through the A4 path. - decode_rows = _num_leading_decode_tokens() - if 0 < decode_rows < tokens: + if ( + self._quant_mode == "w4a16" + and _w4a16_a4_prefill_enabled() + and getattr(getattr(prepared, "_impl", None), "a4_prefill_scales", False) + ): + # The row type decides the A4 prefill: decode rows always run + # W4A16 and prefill rows always A4, whatever the call size. A + # captured CUDA graph (decode steps and the smallest mixed steps) + # keeps W4A16 for all its rows. Layers without activation scales + # (the MTP draft layer) have no A4 path to choose. + if _is_current_stream_capturing(): + parts = ((0, tokens, False),) + else: + decode_rows = min(_num_leading_decode_tokens(tokens), tokens) + parts = ((0, decode_rows, False), (decode_rows, tokens, True)) global _A4_SPLIT_LOGGED - if not _A4_SPLIT_LOGGED: + if not _A4_SPLIT_LOGGED and 0 < decode_rows < tokens: _A4_SPLIT_LOGGED = True logger.info( - "b12x W4A16 A4 prefill split: %d decode rows stay A16, " - "%d prefill rows", + "b12x W4A16 A4 prefill by row type: %d decode rows stay " + "W4A16, %d prefill rows run A4", decode_rows, tokens - decode_rows, ) - for lo, hi in ((0, decode_rows), (decode_rows, tokens)): - binding = _require_b12x_fused_moe().bind( - plan, - scratch=scratch, - a=hidden_states[lo:hi], - experts=prepared, - topk_weights=topk_weights[lo:hi], - topk_ids=topk_ids[lo:hi], - output=output[lo:hi], - input_scales_static=True, - ) - _require_b12x_fused_moe().run(binding=binding) - return + for lo, hi, a4_prefill in parts: + if hi == lo: + continue + binding = _require_b12x_fused_moe().bind( + plan, + scratch=scratch, + a=hidden_states[lo:hi], + experts=prepared, + topk_weights=topk_weights[lo:hi], + topk_ids=topk_ids[lo:hi], + output=output[lo:hi], + input_scales_static=True, + a4_prefill=a4_prefill, + ) + _require_b12x_fused_moe().run(binding=binding) + return binding = _require_b12x_fused_moe().bind( plan, scratch=scratch, From 4dc4de61f5b895738d320397c86c2a9ac5d49dc4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 14:43:54 +0000 Subject: [PATCH 8/8] release: fragment for the row-type A4 prefill (#977) Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 109d044369f4c680deda36ceefd830169efe08a5) --- .../vllm-b12x-a4-prefill-row-type.json | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 .lil/changes/vllm-b12x-a4-prefill-row-type.json diff --git a/.lil/changes/vllm-b12x-a4-prefill-row-type.json b/.lil/changes/vllm-b12x-a4-prefill-row-type.json new file mode 100644 index 000000000000..070b60665706 --- /dev/null +++ b/.lil/changes/vllm-b12x-a4-prefill-row-type.json @@ -0,0 +1,24 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-b12x-a4-prefill-row-type", + "category": "performance", + "summary": "The A4 prefill now covers every prefill row, not only steps of 1,536 tokens or more: short prompts prefill up to 11% faster", + "models": [ + "GLM-5.3-Flash" + ], + "compatibility": "Only with B12X_W4A16_A4_PREFILL_MIN_TOKENS set (prefill-activations a4); decode stays W4A16.", + "details": [ + "The b12x MoE now chooses the precision by row type: decode and MTP verification rows run W4A16 and prefill rows run A4, whatever the step size, so a prompt's precision no longer depends on how the scheduler batches it. Steps replayed from a CUDA graph keep W4A16.", + "GLM-5.3-Flash glm53-tp2 on two RTX PRO 6000 Max-Q: a 1,154-token prompt prefilled at 6,498 instead of 5,868 tok/s; 8K/32K prefill and decode unchanged; needle-checksum 492, 492, 490 of 500 against 490, 492, 497." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 977 + ], + "requires": [ + "vllm-b12x-w4a16-a4-prefill", + "b12x-w4a16-a4-prefill-per-call" + ] +}