[KK] b12x MoE: opt-in NVFP4-activation prefill for W4A16 experts, decode rows stay W4A16

.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"
+  ]
+}

.lil/changes/vllm-b12x-w4a16-a4-prefill.json

@@ -0,0 +1,24 @@
+{
+  "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. 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": [
+    970
+  ],
+  "requires": [
+    "b12x-w4a16-a4-prefill"
+  ]
+}

.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": []
+}

tests/kernels/moe/test_b12x_a4_prefill.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Host-side contract of the opt-in A4 prefill for b12x W4A16 MoE.
+
+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.
+"""
+
+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, 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
+
+
+@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(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(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(8) == 0
+
+
+class _Recorder:
+    def __init__(self) -> None:
+        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
+        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()),
+            a4=kwargs.get("a4_prefill"),
+        )
+
+    def run(self, *, binding) -> None:
+        self.calls.append((binding.rows, binding.first, binding.a4))
+
+
+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
+    prepared = SimpleNamespace(_impl=SimpleNamespace(a4_prefill_scales=a4_scales))
+    experts._prepared = lambda: prepared
+    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
+
+
+@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, tokens])
+    experts, recorder = _experts(monkeypatch, "w4a16")
+    hidden = _apply(experts, tokens)
+    row = hidden.element_size() * hidden.shape[1]
+    assert recorder.calls == [
+        (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, a4_scales",
+    [
+        ("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_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, 2312])
+    experts, recorder = _experts(monkeypatch, quant_mode, a4_scales)
+    _apply(experts, 2312)
+    assert [(rows, a4) for rows, _, a4 in recorder.calls] == [(2312, None)]

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",
     [