From e0ad58ef5873fbd0cf4ed3970c0c8008baf65e03 Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Fri, 18 Sep 2026 12:45:22 +0100 Subject: [PATCH 001/365] fix(qwen): propagate MTP positional overrides Normalize internally constructed MTP drafts before applying target positional overrides, preserving their MTP discriminator for both explicit and automatically detected drafts. Co-authored-by: OMP Agent Signed-off-by: Naadir Jeewa --- .../test_speculative_draft_hf_overrides.py | 37 +++++++++++ tests/models/qwen4_exp/test_b12x_qsa.py | 22 +++++++ vllm/config/speculative.py | 61 ++++++++++++++++--- 3 files changed, 111 insertions(+), 9 deletions(-) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index bc5300e53fce..b81f1f4e292f 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -22,6 +22,8 @@ from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig +pytestmark = pytest.mark.skip_global_cleanup + def _make_hf_config(**kwargs) -> PretrainedConfig: defaults = dict( @@ -43,6 +45,41 @@ def test_dict_overrides_are_not_forwarded_to_draft(): assert composed is SpeculativeConfig.hf_config_override +@pytest.mark.cpu_test +@pytest.mark.parametrize("method", ["mtp", "draft_model"]) +def test_mtp_draft_receives_target_dict_overrides(method: str): + """An in-model MTP draft shares the target's positional geometry.""" + target_hf_overrides = { + "architectures": ["Qwen3_8FlashNextForCausalLM"], + "model_type": "qwen3_8_flash_next", + "text_config": { + "max_position_embeddings": 1048576, + "rope_parameters": {"rope_type": "yarn", "factor": 4.0}, + }, + } + override = SpeculativeConfig.get_draft_hf_overrides(method, target_hf_overrides) + assert callable(override) + text_config = _make_hf_config( + max_position_embeddings=262144, + rope_parameters={"rope_type": "default"}, + hc_count=4, + mtp_num_hidden_layers=2, + num_attention_heads=4, + ) + source = _make_hf_config( + architectures=["Qwen3_8FlashNextForCausalLM"], + model_type="qwen3_8_flash_next", + text_config=text_config, + ) + + out = override(source) + + assert out.model_type == "qwen4_exp_mtp" + assert out.architectures == ["Qwen4ExpMTP"] + assert out.text_config.max_position_embeddings == 1048576 + assert out.text_config.rope_parameters == {"rope_type": "yarn", "factor": 4.0} + + @pytest.mark.cpu_test def test_none_overrides_fall_back_to_arch_mapping(): composed = SpeculativeConfig.compose_draft_hf_overrides(None) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 48540e675c36..36452ced19ea 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/tests/models/qwen4_exp/test_b12x_qsa.py @@ -490,6 +490,7 @@ def test_qsa_selector_fork_requires_one_full_graph( def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None: + """Prefill plan capacities cover each requested context ceiling.""" assert qsa_module._qsa_prefill_context_capacities(262144, 4096) == ( 4096, 8192, @@ -513,6 +514,27 @@ def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None: ) +def test_qsa_prefill_binding_accepts_pass_one_at_one_million_tokens() -> None: + """Pass one uses a prefill plan at the target context ceiling.""" + owner = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + owner.max_decode_rows = 16 + owner.max_seq_len = 1048576 + prefill_plan = object() + owner._qsa_prefill_bindings = ( + qsa_module._QSAContextPlan( + max_seq_len=1048576, + caps=None, + plan=prefill_plan, + ), + ) + owner._qsa_decode_context = None + + assert ( + owner._qsa_binding_for_workload(rows=17, max_seq_len=1048576).plan + is prefill_plan + ) + + def test_qsa_run_consumes_projection_views_and_writes_live_output(monkeypatch) -> None: rows = 2 query = torch.randn(rows, 6, 256, dtype=torch.bfloat16) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 42df7dbd911d..1384d7897c45 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1088,9 +1088,36 @@ def _apply_composed_hf_override( target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig], hf_config: PretrainedConfig, ) -> PretrainedConfig: + """Apply the draft normalization before a target config transform.""" hf_config = SpeculativeConfig.hf_config_override(hf_config) return target_hf_overrides(hf_config) + @staticmethod + def _apply_mtp_dict_hf_override( + target_hf_overrides: dict[str, Any], + hf_config: PretrainedConfig, + *, + require_mtp: bool = False, + ) -> PretrainedConfig: + """Apply target patches without replacing the MTP discriminator.""" + hf_config = SpeculativeConfig.hf_config_override(hf_config) + if require_mtp and hf_config.model_type not in get_args(MTPModelTypes): + return hf_config + for key, value in target_hf_overrides.items(): + if key in {"architectures", "model_type"}: + continue + target = getattr(hf_config, key, None) + if isinstance(value, dict) and target is not None: + if isinstance(target, dict): + target.update(value) + continue + if hasattr(target, "__dict__"): + for nested_key, nested_value in value.items(): + setattr(target, nested_key, nested_value) + continue + setattr(hf_config, key, value) + return hf_config + @staticmethod def compose_draft_hf_overrides( target_hf_overrides: HfOverrides | None, @@ -1117,6 +1144,27 @@ def compose_draft_hf_overrides( SpeculativeConfig._apply_composed_hf_override, target_hf_overrides ) + @staticmethod + def get_draft_hf_overrides( + method: str, + target_hf_overrides: HfOverrides | None, + ) -> HfOverrides: + """Return overrides appropriate for the selected draft method.""" + if method == "medusa": + return {"model_type": "medusa"} + if method == "mtp" and isinstance(target_hf_overrides, dict): + return functools.partial( + SpeculativeConfig._apply_mtp_dict_hf_override, + target_hf_overrides, + ) + if method == "draft_model" and isinstance(target_hf_overrides, dict): + return functools.partial( + SpeculativeConfig._apply_mtp_dict_hf_override, + target_hf_overrides, + require_mtp=True, + ) + return SpeculativeConfig.compose_draft_hf_overrides(target_hf_overrides) + @staticmethod def _is_custom_proposer_path(model: str | None) -> bool: """True if ``model`` is a dotted import path (e.g. ``pkg.MyProposer``).""" @@ -1130,6 +1178,7 @@ def _is_custom_proposer_path(model: str | None) -> bool: return len(parts) >= 2 and all(part.isidentifier() for part in parts) def __post_init__(self): + """Validate speculative settings and construct the draft configuration.""" # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter # will be used to set the draft model, eagle head, or additional weight @@ -1318,15 +1367,9 @@ def __post_init__(self): # detect them. When the method is explicitly "medusa", inject # model_type so MedusaConfig.from_pretrained is used instead. draft_hf_overrides: HfOverrides - if self.method == "medusa": - draft_hf_overrides = {"model_type": "medusa"} - else: - # Compose any callable hf_overrides set on the target so the - # draft config receives the same transform (e.g. the test - # shrink). Dict overrides stay target-only. - draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides( - self.target_model_config.hf_overrides - ) + draft_hf_overrides = SpeculativeConfig.get_draft_hf_overrides( + self.method, self.target_model_config.hf_overrides + ) self.draft_model_config = ModelConfig( model=self.model, runner="draft", From 1485b050552c9f669de848f184f742055d620bec Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 20:40:28 +0000 Subject: [PATCH 002/365] fix(ds4.1): release transient MHC and WO graph activations Retain caller-owned scratch without extending activation lifetimes through B12X bindings. Keep prepared plans, fresh bindings, precision and CUDA graph coverage unchanged. Native GPU tests reproduce retained operands in both MHC modes and both WO modes. All four corrected cases and the block-linear regression pass exact changed-input graph replay. TP4 DSpark K7 serving restores logical KV capacity from 1,920,034 to 4,466,215 tokens and passes bounded text, prefix and vision checks. C1 is -0.65%, C8 -2.32% and 32k prefill +0.44% in sequential samples; regression-free throughput is not claimed. Ruff findings and eleven source mypy findings reproduce on the untouched base. Tests pass mypy. Skip those pre-existing hook failures and whole-file formatting for this isolated local commit; other commit hooks remain enabled. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit a28be967e5779629dbc212a312a3ad8924c0623f) --- tests/v1/attention/test_b12x_v41_workspace.py | 241 ++++++++++++++++++ vllm/models/deepseek_v4_1/attention.py | 7 +- vllm/models/deepseek_v4_1/b12x_layers.py | 4 +- 3 files changed, 249 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index e6bd78e6719d..1a87261c3e7a 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -422,6 +422,144 @@ def run(rows): session.close() +@pytest.mark.parametrize("operation", ["pre", "post_pre"]) +def test_mhc_capture_retains_scratch_not_caller_activations( + native_workspace, monkeypatch, operation +): + """Intermediate MHC outputs must be reusable across shared-pool graphs.""" + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4_1 import b12x_layers + from vllm.utils.b12x import B12xWorkload, register_b12x_layer + + _, manager, workspace = native_workspace + monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (8, 16)) + device = torch.device("cuda") + hidden = 5120 + torch.manual_seed(146) + layer = torch.nn.Module() + for kind in ("attn", "ffn"): + setattr(layer, f"hc_{kind}_fn", torch.randn(24, hidden * 4, device=device) / 64) + setattr(layer, f"hc_{kind}_scale", torch.ones(3, device=device)) + setattr(layer, f"hc_{kind}_base", torch.zeros(24, device=device)) + setattr( + layer, + f"{kind}_norm", + SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device) + ), + ) + layer.hc_attn_fn_broadcast = None + module = layer._b12x_mhc = b12x_layers.B12xMHC( + SimpleNamespace( + hidden_size=hidden, + rms_norm_eps=1e-20, + hc_eps=1e-6, + hc_sinkhorn_iters=20, + hc_mult=4, + ) + ) + name = f"test.mhc.capture.{operation}" + module.bind_layer_name(name) + register_b12x_layer(name, layer) + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=16, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare( + tuple( + request + for unit in module.get_b12x_preparation_units(layer, workload) + for request in unit.requests + ) + ) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in module._plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + residual = torch.randn(16, 4, hidden, dtype=torch.bfloat16, device=device) + previous = torch.randn(16, hidden, dtype=torch.bfloat16, device=device) + pre = torch.full((16, 4), 0.25, device=device) + previous_post = pre.clone() + comb = torch.eye(4, device=device).expand(16, -1, -1).contiguous() + destinations = [ + torch.empty_like(residual), + torch.empty_like(pre), + torch.empty_like(comb), + torch.empty_like(previous), + torch.empty_like(pre), + ] + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + kwargs = ( + {} + if operation == "pre" + else dict( + previous_output=previous[:rows], + previous_post=previous_post[:rows], + previous_comb=comb[:rows], + ) + ) + values = module.pre( + residual[:rows], + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + layer.attn_norm.weight, + pre[:rows], + **kwargs, + ) + for destination, value in zip(destinations, values, strict=True): + references.append(weakref.ref(value)) + destination[:rows].copy_(value) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + residual.normal_() + previous.normal_() + run(rows) + expected = [output[:rows].clone() for output in destinations] + for output in destinations: + output.fill_(float("nan")) + graphs[rows].replay() + for output, reference in zip(destinations, expected, strict=True): + torch.testing.assert_close(output[:rows], reference, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + @pytest.mark.parametrize("draft_tokens", [5, 7]) def test_parallel_draft_reservation_keeps_decode_split_parallelism( native_workspace, draft_tokens @@ -824,6 +962,109 @@ def projection(n, k, value): del resources +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_output_projection_capture_releases_caller_activations( + native_workspace, is_prefill +): + """WO bindings may borrow inputs and results without pinning them per graph.""" + from b12x.preparation import PreparationSession + + from vllm.utils.b12x import B12xWorkload + + attention, manager, workspace = native_workspace + device = torch.device("cuda") + torch.manual_seed(147) + layer = _layer(attention) + layer.capacity = 16 + layer.n_local_groups = 2 + layer.n_local_heads = 16 + layer.head_dim = 512 + layer.rope_head_dim = 64 + layer.o_lora_rank = 1024 + layer.hidden_size = 5120 + layer._wo_plans = {} + angles = torch.randn(32, 32, device=device) + layer.rotary_emb = SimpleNamespace( + cos_sin_cache=torch.cat((angles.cos(), angles.sin()), dim=-1) + ) + for name, shape in (("wo_a", (2048, 4096)), ("wo_b", (5120, 2048))): + setattr( + layer, + name, + SimpleNamespace( + weight=(torch.randn(shape, device=device) / 32).to(torch.float8_e4m3fn), + weight_scale_inv=torch.ones( + shape[0] // 32, shape[1] // 32, device=device + ).to(torch.float8_e8m0fnu), + ), + ) + layer.setup_wo_projection() + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=32, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare(layer._wo_preparation_unit(workload).requests) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in layer._wo_plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + inputs = torch.randn(16, 16, 512, dtype=torch.bfloat16, device=device) + positions = torch.arange(16, device=device) + output = torch.empty(16, 5120, dtype=torch.bfloat16, device=device) + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + source = inputs[:rows] + 1 + result = layer._o_proj(source, positions[:rows], is_prefill=is_prefill) + references.extend((weakref.ref(source), weakref.ref(result))) + if result._base is not None: + references.append(weakref.ref(result._base)) + output[:rows].copy_(result) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + inputs.normal_() + positions.copy_(torch.randperm(16, device=device)) + run(rows) + expected = output[:rows].clone() + output.fill_(float("nan")) + graphs[rows].replay() + torch.testing.assert_close(output[:rows], expected, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + def test_output_projection_uses_fused_block32_path(native_workspace, monkeypatch): from dataclasses import dataclass diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 661e1f862bee..06c590ce5bc8 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1706,9 +1706,10 @@ def _o_proj(self, o, positions, *, is_prefill=False): ) require_prepared(plan, "gemm.wo_projection", o.device) weights = self._wo_projection_weights + scratch = _scratch(plan) binding = wo_projection.bind_inv_rope( plan, - scratch=_scratch(plan), + scratch=scratch, o=o, positions=positions, cos_sin_cache=self.rotary_emb.cos_sin_cache, @@ -1725,7 +1726,9 @@ def _o_proj(self, o, positions, *, is_prefill=False): device=o.device, ) binding = replace(binding, output=output) - retain_cuda_graph_capture_resource(binding) + # The model owns the plan and weights; graph pools manage activation + # lifetimes. Keep only the borrowed workspace, not layer activations. + retain_cuda_graph_capture_resource(scratch) local = wo_projection.run_inv_rope( binding=binding, plan=plan, diff --git a/vllm/models/deepseek_v4_1/b12x_layers.py b/vllm/models/deepseek_v4_1/b12x_layers.py index de605dfd11c0..57292c9db2e1 100644 --- a/vllm/models/deepseek_v4_1/b12x_layers.py +++ b/vllm/models/deepseek_v4_1/b12x_layers.py @@ -586,7 +586,9 @@ def _mhc_pre( comb=comb, pre_out=pre_out, ) - retain_cuda_graph_capture_resource(binding) + # Caller-owned outputs follow PyTorch's graph-pool lifetimes. Retaining + # them through the binding prevents reuse across layers and graph shapes. + retain_cuda_graph_capture_resource(scratch) if previous_output is None: mhc.run_pre( residual, From 09182242a270d259d2855b701098698c99c3c8e4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 21:16:53 +0000 Subject: [PATCH 003/365] fix(ds4.1): release captured query and context activations Retain caller-owned scratch and attention metadata without retaining transient query/index operands or DSpark context-projection results. Prepared plans, per-call bindings, numerical kernels, cache geometry and graph coverage remain unchanged. Real SM120 tests reproduce the immutable-image ownership defects, including detached query storage. Twenty-four native lifetime and exact changed-input replay cases pass across decode/prefill and three page geometries. TP4 DSpark K7 memory snapshots remove 858657792 additional live graph bytes per rank and pass text, prefix and vision checks. Independent uninstrumented throughput qualification is recorded separately. The test file passes mypy; existing source typing findings remain. Ruff diagnostics match the parent revision. Skip the reproducing Ruff/mypy findings and unrelated whole-file formatting while retaining the remaining commit hooks. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 42627f07f551a5a38575eb1baca9693d7856c35e) --- tests/v1/attention/test_b12x_v41_workspace.py | 86 +++++++++++++------ vllm/models/deepseek_v4_1/attention.py | 22 ++--- vllm/models/deepseek_v4_1/nvidia/dspark.py | 2 +- 3 files changed, 72 insertions(+), 38 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 1a87261c3e7a..549ae1e88fec 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -4,10 +4,12 @@ import gc import weakref +from functools import partial from types import SimpleNamespace import pytest import torch +from torch.multiprocessing.reductions import StorageWeakRef from vllm.config import CUDAGraphMode @@ -336,9 +338,11 @@ def test_index_preparation_reuses_reserved_workspace(native_workspace): torch.testing.assert_close(cache, torch.zeros_like(cache), rtol=0, atol=0) +@pytest.mark.parametrize("projection_kind", ["linear", "dspark_context"]) def test_block_linear_capture_retains_scratch_not_caller_activations( native_workspace, monkeypatch, + projection_kind, ): from vllm.models.deepseek_v4_1 import b12x_layers @@ -353,10 +357,22 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( layer.weight = torch.randn(256, 256, device=device).to(torch.float8_e4m3fn) layer.weight_scale_inv = torch.ones(8, 8, device=device).to(torch.float8_e8m0fnu) monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (16,)) - method = b12x_layers.B12xFP8LinearMethod( - SimpleNamespace(weight_block_size=[32, 32]) - ) - method.process_weights_after_loading(layer) + if projection_kind == "linear": + method = b12x_layers.B12xFP8LinearMethod( + SimpleNamespace(weight_block_size=[32, 32]) + ) + method.process_weights_after_loading(layer) + project = partial(method.apply, layer) + plans = layer.b12x_plans + else: + from vllm.models.deepseek_v4_1.nvidia import dspark + + monkeypatch.setattr(dspark, "_execution_capacities", lambda: (16,)) + method = dspark._ContextKVProjection( + SimpleNamespace(fused_wqa_wkv=layer, q_lora_rank=0), 16 + ) + project = method + plans = method.plans workload = B12xWorkload( stage="weights", token_counts=(8, 16), @@ -376,7 +392,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( ) session.freeze() manager.reserve_all( - *((spec.shape, spec.dtype) for spec in layer.b12x_plans[0].scratch_specs()) + *((spec.shape, spec.dtype) for spec in plans[0].scratch_specs()) ) manager.lock() inputs = torch.randn(16, 256, dtype=torch.bfloat16, device=device) @@ -385,7 +401,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( def run(rows): source = inputs[:rows] + 1 - output = method.apply(layer, source) + output = project(source) references.extend((weakref.ref(source), weakref.ref(output))) outputs[:rows].copy_(output) @@ -695,6 +711,7 @@ def test_attention_shared_scratch_graph_replay( device = torch.device("cuda", torch.accelerator.current_device_index()) torch.manual_seed(142) layer = _layer(attention, swa_page=swa_page) + layer.rotary_emb = SimpleNamespace(cos_sin_cache=torch.empty(0, device=device)) layer.config.cache_config.block_size = main_page layer.max_model_len = 4096 if rows == 36: @@ -755,13 +772,12 @@ def metadata(page): device=device, dtype=torch.uint8, ) - _write_cache( - kv, - cache, - torch.arange(length, device=device) + page, - page_size=page, - cache_kind=kind, - cache_format="deepseek_v41", + from b12x.attention._shared.mla.compressed_reference import ( + pack_deepseek_v41_cache_reference, + ) + + cache[1:].copy_( + pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -782,12 +798,6 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 - _write_index_keys( - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - page_size=main_page, - ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -811,17 +821,35 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) + attention.dsa_indexer.quantize_write_index_k_mxfp4( + layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + ) + activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): - attention.dsa_indexer.quantize_q_mxfp4( - layer._index_plan("decode", min(rows, layer.DECODE_CHUNK)), - iq, - q_mxfp4=packed, - q_scales=scales, - ) - layer.forward_mqa( - q, None, positions, out, index_query=(packed, scales, weights) + query = q.clone() + index_query = packed.clone(), scales.clone(), weights.clone() + activation_refs.extend( + (name, StorageWeakRef(tensor.untyped_storage())) + for name, tensor in zip( + ("query", "index_query", "index_scales", "index_weights"), + (query, *index_query), + ) ) + mode = "decode" if is_decode else "prefill" + chunk = layer.DECODE_CHUNK if is_decode else layer.INDEX_CHUNK + for offset in range(0, rows, chunk): + end = min(offset + chunk, rows) + attention.dsa_indexer.quantize_q_mxfp4( + layer._index_plan(mode, end - offset), + iq[offset:end], + q_mxfp4=index_query[0][offset:end], + q_scales=index_query[1][offset:end], + ) + layer.forward_mqa(query, None, positions, out, index_query=index_query) run() batched_output = out.clone() @@ -845,6 +873,10 @@ def run(): torch.cuda.graph(graph, stream=stream), ): run() + # Graph resources own scratch, not the per-layer query activations. + gc.collect() + retained = [name for name, reference in activation_refs if not reference.expired()] + assert not retained, f"Capture resources retained caller activations: {retained}" # Changed queries exercise both selection and attention on replay, after # all plan storage has been reused by another operation. for step, seed in enumerate((143, 144)): diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 06c590ce5bc8..4dd404ee4853 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1454,16 +1454,17 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): candidate_lengths=source._candidate_lens[offset:end], ) index_plan = self._index_plan(index_mode, count) + scratch = current_workspace_manager().get_simultaneous( + *( + (spec.shape, spec.dtype) + for spec in dsa_indexer.scratch_specs( + index_plan, device=iq_data.device + ) + ) + ) binding = dsa_indexer.bind( index_plan, - scratch=current_workspace_manager().get_simultaneous( - *( - (spec.shape, spec.dtype) - for spec in dsa_indexer.scratch_specs( - index_plan, device=iq_data.device - ) - ) - ), + scratch=scratch, q_mxfp4=iq_data[offset:end], q_scales=iq_scale[offset:end], query_weights=iw[offset:end], @@ -1475,7 +1476,7 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): output_indices=self.topk_indices_buffer[offset:end], **candidate_args, ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) dsa_indexer.score(binding) dsa_indexer.select(binding) @@ -1534,7 +1535,8 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): swa_lengths=swa_lengths[:rows], **kwargs, ) - retain_cuda_graph_capture_resource(binding) + # Keep scratch and metadata, not aliases of the transient query storage. + retain_cuda_graph_capture_resource(buffers) mla.run( binding=binding, swa_k_cache=self.swa_cache_layer.kv_cache, diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v4_1/nvidia/dspark.py index 5251cc82d9e4..61ba555c05d3 100644 --- a/vllm/models/deepseek_v4_1/nvidia/dspark.py +++ b/vllm/models/deepseek_v4_1/nvidia/dspark.py @@ -186,7 +186,7 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor: packed_weight=self.weight, output=out.view(rows, self.weight.out_features, 1), ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) block_fp8_linear.run(binding=binding) return out From 353ff02af20b0f7dc67d927b48df05569d8aaeec Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 12:25:32 +0000 Subject: [PATCH 004/365] test(ds41): verify MHC ownership across breakable graph segments Exercise pre and lagged post-pre outputs across an eager segment boundary and changed-input replay on a shared graph pool. Require transient Python owners to release while caller scratch remains retained. Native graph semantics and model code are unchanged. Applicable hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 3ead61584ac374ca5825ffb8aa21ad8a1a05beca) --- tests/v1/attention/test_b12x_v41_workspace.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 549ae1e88fec..260f386aa122 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -439,12 +439,14 @@ def run(rows): @pytest.mark.parametrize("operation", ["pre", "post_pre"]) +@pytest.mark.parametrize("capture_kind", ["full", "breakable"]) def test_mhc_capture_retains_scratch_not_caller_activations( - native_workspace, monkeypatch, operation + native_workspace, monkeypatch, operation, capture_kind ): """Intermediate MHC outputs must be reusable across shared-pool graphs.""" from b12x.preparation import PreparationSession + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture from vllm.models.deepseek_v4_1 import b12x_layers from vllm.utils.b12x import B12xWorkload, register_b12x_layer @@ -537,6 +539,10 @@ def run(rows): pre[:rows], **kwargs, ) + capture = BreakableCUDAGraphCapture.current() + if capture is not None: + # The lagged outputs cross a segment boundary before consumption. + capture.add_eager(lambda: None) for destination, value in zip(destinations, values, strict=True): references.append(weakref.ref(value)) destination[:rows].copy_(value) @@ -547,14 +553,23 @@ def run(rows): for rows in (16, 8): run(rows) torch.accelerator.synchronize() - graph = torch.cuda.CUDAGraph() + graph = ( + torch.cuda.CUDAGraph() + if capture_kind == "full" + else BreakableCUDAGraphCapture(pool=pool) + ) graphs[rows] = graph + context = ( + torch.cuda.graph(graph, pool=pool) if capture_kind == "full" else graph + ) with ( session.capture(), workspace.collect_cuda_graph_capture_resources() as resources, - torch.cuda.graph(graph, pool=pool), + torch.cuda.stream(torch.cuda.Stream()), + context, ): run(rows) + torch.accelerator.synchronize() owners.append(resources) gc.collect() assert all(reference() is None for reference in references) From 64752179cd97691b72df2f00aee04a7ee94ffd21 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:17:17 +0000 Subject: [PATCH 005/365] fix(ds4): reuse MHC autotuning output buffers Allocate one output set per preparation candidate so repeated CUDA graph timing samples do not accumulate functional-output copies beyond the primed memory budget. Serving kernels, operation queries, candidate search and caller-owned scratch are unchanged. Validation: 13 native SM120 tests pass, including exact functional parity and changed-input CUDA graph replay at 3 and 4096 tokens. The baseline fails output-pointer reuse. DS4 TP2 completes weights preparation; end-to-end serving remains gated by KV admission capacity. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 2b2fa507b89d9652c5d87db61c44ee89b5fe864f) --- .../kernels/test_b12x_linear.py | 112 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 22 +++- 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 72c4f6d130ff..6a1a540400e3 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -7,6 +7,7 @@ import sys import types from dataclasses import dataclass +from typing import Any import pytest import torch @@ -1874,6 +1875,117 @@ def test_b12x_mhc_declares_the_operations_each_layer_runs( mhc._plan_for("pre", 8) +@torch.inference_mode() +@pytest.mark.parametrize("operation", ["pre", "post_pre", "post_pre_bf16", "post"]) +@pytest.mark.parametrize("tokens", [3, 4096]) +def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( + monkeypatch, operation, tokens +): + """Tuning fixtures retain one output set per candidate, not per sample.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("native b12x mHC requires SM12x") + from b12x.norm.mhc import _impl + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x + + device = torch.device("cuda", torch.accelerator.current_device_index()) + hidden, mult = 4096, 4 + fn = torch.randn(24, mult * hidden, device=device) * 0.001 + norm = types.SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device), + variance_epsilon=1e-6, + ) + layer = types.SimpleNamespace( + hc_attn_fn=fn, + hc_ffn_fn=fn, + hc_attn_scale=torch.ones(3, device=device), + hc_ffn_scale=torch.ones(3, device=device), + hc_attn_base=torch.zeros(24, device=device), + hc_ffn_base=torch.zeros(24, device=device), + hc_attn_fn_broadcast=fn.view(24, mult, hidden).sum(1), + attn_norm=norm, + ffn_norm=norm, + hc_ffn_fn_bf16=fn.to(torch.bfloat16), + ) + mhc = dsv4_b12x.B12xMHCResidual( + hidden_size=hidden, + hc_mult=mult, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=20, + ) + workload = B12xWorkload( + stage="weights", + token_counts=(tokens,), + fixed_token_counts=(), + output_dtype=torch.bfloat16, + max_tokens=tokens, + max_seqs=1, + max_model_len=tokens, + ) + (unit,) = mhc.get_b12x_preparation_units(layer, workload) + request = next( + item for item in unit.requests if item.name.endswith(f".{operation}.m{tokens}") + ) + native_operation = "post_pre" if operation == "post_pre_bf16" else operation + name = f"_b12x_mhc_{native_operation}_impl" + native = getattr(_impl, name) + captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + + def observe(*args, **kwargs): + captured_call[:] = [(args, kwargs)] + return native(*args, **kwargs) + + def outputs(value): + return value if isinstance(value, tuple) else (value,) + + with PreparationSession( + device=device, autotune=False, compile_workers=0 + ) as session: + session.prepare((request,)) + call = request.prepare_call(request.plan.prepared.state) + with monkeypatch.context() as hooks: + hooks.setattr(_impl, name, observe) + call.produce() + first = outputs(call.invoke()) + args, kwargs = captured_call.pop() + reference_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} + } + pointers = [tensor.data_ptr() for tensor in first] + second = outputs(call.invoke()) + assert [tensor.data_ptr() for tensor in second] == pointers + for actual, expected in zip( + second, outputs(native(*args, **reference_kwargs)), strict=True + ): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + try: + with session.capture(), torch.cuda.graph(graph): + call.invoke() + call.invoke() + for _ in range(2): + call.produce() + for tensor in first: + tensor.fill_(float("nan")) + allocated = torch.accelerator.memory_allocated(device) + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated(device) == allocated + assert [ + tensor.data_ptr() for tensor in outputs(call.output) + ] == pointers + expected = outputs(native(*args, **reference_kwargs)) + for actual, wanted in zip(first, expected, strict=True): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + finally: + graph.reset() + + def test_b12x_mhc_keeps_plans_across_workloads(monkeypatch) -> None: from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 98c210d79767..232d84b2dc94 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -435,6 +435,24 @@ def prepare(state): dtype=torch.float32, device=device, ) + # Each candidate owns one output set. Functional allocations inside + # repeated capture samples would reserve additional graph-pool copies + # that are absent from the tuner's primed-residency measurement. + residual_out = torch.empty( + (tokens, self.hc_mult, self.hidden_size), + dtype=torch.bfloat16, + device=device, + ) + outputs = ( + {} + if operation == "post" + else { + "residual_out": residual_out, + "y_out": torch.empty_like(x), + "post_out": torch.empty_like(post), + "comb_out": torch.empty_like(comb), + } + ) def produce(): residual.normal_() @@ -455,6 +473,7 @@ def produce(): norm_weight=attn_norm.weight, norm_eps=float(attn_norm.variance_epsilon), _state=state, + **outputs, ) elif operation in ("post_pre", "post_pre_bf16"): if operation == "post_pre": @@ -488,10 +507,11 @@ def produce(): norm_weight=norm.weight, norm_eps=float(norm.variance_epsilon), _state=state, + **outputs, ) else: run = lambda: _impl._b12x_mhc_post_impl( - x, residual, post, comb, _state=state + x, residual, post, comb, out=residual_out, _state=state ) def execute(): From 15f47ba914fe4a388e97f0fd8bbb61b6099288d9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:54:03 +0000 Subject: [PATCH 006/365] fix(workspace): reserve B12X scratch by model owner Keep target and DSpark draft scratch capacities independent during profile hooks and B12X preparation. Reserve shared modules in both lanes and propagate each lane capacity across microbatches before graph capture. Preserve the complete attention shape envelope and unscoped coordinator compatibility. Validation: 83 focused workspace, preparation and graph-profile tests pass. Stock DS4 TP2/DSpark K5 with a 4096-token budget admits 7.56 GiB KV instead of 6.12 GiB, retaining the 1M-token model limit. Serving throughput qualification is pending. The unchanged base reproduces eight coordinator mypy errors and the existing pickle-import lint finding; remaining hooks pass. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit ea358a87e72e6f44baf14c436d278b5f4d52b92c) --- tests/model_executor/test_b12x_warmup.py | 25 +++++ tests/v1/worker/test_workspace.py | 117 +++++++++++++++++++++ vllm/model_executor/warmup/b12x_prepare.py | 48 +++++++-- vllm/utils/b12x.py | 7 +- vllm/v1/worker/b12x_startup.py | 20 +++- vllm/v1/worker/gpu/model_runner.py | 29 +++-- vllm/v1/worker/workspace.py | 40 +++++-- 7 files changed, 259 insertions(+), 27 deletions(-) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index c54ce7c3856d..41690faa4fef 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -348,6 +348,31 @@ def test_scope_unit_calls_wraps_mapping_factories() -> None: request.prepare_call = {2: lambda state: _Call(run=lambda: _workspace_lane.get())} scoped = scope_b12x_unit_calls(_unit("u", request), 1) assert scoped.requests[0].prepare_call[2](object()).run() == 1 + assert scoped.workspace_lanes == (1,) + + +@pytest.mark.parametrize("draft_lane", [0, 1]) +def test_shared_preparation_scratch_covers_both_model_owners(draft_lane): + target, draft, shared = (torch.nn.Module() for _ in range(3)) + for module in (target, draft, shared): + module.b12x_preparation_provider = _Provider() + target.shared = shared + draft.shared = shared + units = b12x_prepare.collect_b12x_units( + _worker(target, draft=draft, draft_lane=draft_lane), _workload() + ) + ownership = b12x_prepare._request_workspace_lanes(units) + assert ownership[f"{id(target):x}"] == (0,) + assert ownership[f"{id(draft):x}"] == (draft_lane,) + assert ownership[f"{id(shared):x}"] == tuple(sorted({0, draft_lane})) + assert ( + sum( + f"{id(shared):x}" == request.name + for unit in units + for request in unit.requests + ) + == 1 + ) def test_blockscaled_holder_declares_provided_workspace_under_the_cap( diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 0eba7b71d592..1699fe90ea7f 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from typing import cast import pytest @@ -167,6 +168,122 @@ def test_workspace_lane_validation(monkeypatch) -> None: workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) +def test_profile_reserves_model_lanes_without_replicating_target_capacity( + monkeypatch, +) -> None: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + monkeypatch.setattr(workspace, "_manager", manager) + calls: list[tuple[int, int]] = [] + + class ScratchOwner(torch.nn.Module): + def __init__(self, size): + super().__init__() + self.size = size + + def reserve_profile_scratch(self): + calls.append((self.size, workspace._workspace_lane.get())) + manager.get_simultaneous(((self.size,), torch.uint8)) + + target = ScratchOwner(4096) + draft = ScratchOwner(512) + shared = ScratchOwner(256) + target.add_module("shared", shared) + draft.add_module("shared", shared) + runner = SimpleNamespace( + get_model=lambda: target, + get_draft_model=lambda: draft, + _draft_workspace_lane=1, + compilation_config=SimpleNamespace( + static_forward_context={ + "target": target, + "alias": target, + "draft": draft, + "shared": shared, + } + ), + ) + GPUModelRunner._reserve_profile_scratch(runner) + + assert calls == [(4096, 0), (512, 1), (256, 0), (256, 1)] + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + assert [buffer.numel() for buffer in buffers] == [4096, 512, 4096, 512] + pointers = [buffer.data_ptr() for buffer in buffers] + assert len(set(pointers)) == 4 + manager.lock() + manager.reserve_by_lane() + for ubatch in range(2): + active_ubatch[0] = ubatch + for lane, size in enumerate((4096, 512)): + with workspace.use_workspace_lane(lane): + (view,) = manager.get_simultaneous(((size,), torch.uint8)) + assert view.data_ptr() == pointers[ubatch * 2 + lane] + + +def test_lane_reservation_rejects_unreserved_microbatch_when_locked(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((256,), torch.uint8)) + manager.lock() + with pytest.raises(AssertionError, match="reserve_by_lane"): + manager.reserve_by_lane() + + +@pytest.mark.parametrize("scoped", [False, True]) +def test_preparation_reservation_preserves_per_lane_profile_capacity( + monkeypatch, scoped +): + from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator + + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((4096,), torch.uint8)) + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((256,), torch.uint8)) + requests = [] + for name, size in (("target", 1024), ("draft", 512), ("shared", 768)): + spec = SimpleNamespace(shape=(size,), dtype=torch.uint8) + requests.append( + SimpleNamespace( + name=name, + plan=SimpleNamespace(scratch_specs=lambda spec=spec: (spec,)), + ) + ) + job = SimpleNamespace( + advance=lambda **kwargs: SimpleNamespace(done=True, pending_compilation=False), + result=lambda: SimpleNamespace(close=lambda: None), + close=lambda: None, + ) + coordinator = B12xPreparationCoordinator( + SimpleNamespace(begin=lambda *args, **kwargs: job), + [(tuple(requests), False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=manager, + request_workspace_lanes=( + {"target": (0,), "draft": (1,), "shared": (0, 1)} if scoped else None + ), + ) + result = coordinator.advance() + assert result["done"] and result["error"] is None + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + expected = [4096, 768, 4096, 768] if scoped else [4096] * 4 + assert [buffer.numel() for buffer in buffers] == expected + assert len({buffer.data_ptr() for buffer in buffers}) == 4 + manager.lock() + manager.reserve_by_lane() + + def test_cuda_graph_capture_resources_are_scoped_to_collector() -> None: outside = object() first = object() diff --git a/vllm/model_executor/warmup/b12x_prepare.py b/vllm/model_executor/warmup/b12x_prepare.py index dafff99c8e49..6ec0d9768778 100644 --- a/vllm/model_executor/warmup/b12x_prepare.py +++ b/vllm/model_executor/warmup/b12x_prepare.py @@ -228,6 +228,7 @@ def _units_from_modules( workload: B12xWorkload, *, seen: set[int] | None = None, + module_lanes: dict[int, tuple[int, ...]] | None = None, ) -> Iterable[B12xPreparationUnit]: if seen is None: seen = set() @@ -254,7 +255,12 @@ def _units_from_modules( f"{type(provider).__qualname__} returned a non-unit " "preparation value" ) - yield unit + lanes = ( + (workload.lane,) + if module_lanes is None + else module_lanes.get(id(module), (workload.lane,)) + ) + yield replace(unit, workspace_lanes=lanes) def mark_b12x_eager_shapes(worker: Worker) -> None: @@ -328,26 +334,43 @@ def collect_b12x_units( """Collect every unit of one stage from the model, the draft, and comms.""" mark_b12x_eager_shapes(worker) units: list[B12xPreparationUnit] = [] + model = worker.get_model() draft = worker.get_draft_model() + lane = _draft_lane(worker) if draft is not None else 0 + lanes_by_module: dict[int, set[int]] = {} + for owner, owner_lane in ((model, 0), (draft, lane)): + if owner is not None: + for module in owner.modules(): + lanes_by_module.setdefault(id(module), set()).add(owner_lane) + module_lanes = {key: tuple(sorted(lanes)) for key, lanes in lanes_by_module.items()} if draft is not None: - lane = _draft_lane(worker) workload = _draft_workload(worker, workload, lane=0) # Target and draft can share the same embedding and output-head modules. seen: set[int] = set() - units.extend(_units_from_modules(worker.get_model(), workload, seen=seen)) + units.extend( + _units_from_modules(model, workload, seen=seen, module_lanes=module_lanes) + ) if draft is not None: draft_workload = replace(workload, lane=lane) from vllm.v1.worker.workspace import use_workspace_lane with use_workspace_lane(lane): - draft_units = list(_units_from_modules(draft, draft_workload, seen=seen)) + draft_units = list( + _units_from_modules( + draft, draft_workload, seen=seen, module_lanes=module_lanes + ) + ) if lane: draft_units = [scope_b12x_unit_calls(unit, lane) for unit in draft_units] units.extend(draft_units) for provider in b12x_unit_providers(): hook = getattr(provider, "get_b12x_preparation_units", None) if callable(hook): - units.extend(hook(provider, workload)) + # Communicators serve target and draft executions independently. + units.extend( + replace(unit, workspace_lanes=tuple(sorted({0, lane}))) + for unit in hook(provider, workload) + ) # The state stage also carries the weights-stage units: their prepared # plans are skipped by the session, and any plan a layer declared for a # count that only the state stage exposes is prepared before capture. @@ -447,6 +470,14 @@ def b12x_batches(units: Iterable[B12xPreparationUnit], *, autotune: bool = True) return batches +def _request_workspace_lanes(units: Iterable[B12xPreparationUnit]): + return { + request.name: unit.workspace_lanes + for unit in units + for request in unit.requests + } + + def get_b12x_session(worker: Worker): """Return the worker's preparation session, creating it on first use.""" session = getattr(worker, "_b12x_session", None) @@ -492,10 +523,12 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): from vllm.v1.worker.workspace import current_workspace_manager batches = [] + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage="state" if stage == "bind" else stage) + units = collect_b12x_units(worker, workload) batches = b12x_batches( - collect_b12x_units(worker, workload), + units, autotune=( bool(worker.vllm_config.kernel_config.enable_b12x_autotune) and os.environ.get("B12X_AUTOTUNE", "1") != "0" @@ -531,6 +564,7 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if batches else None, + request_workspace_lanes=_request_workspace_lanes(units), ) @@ -608,6 +642,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: from vllm.v1.worker.workspace import current_workspace_manager requests: tuple[PreparationRequest, ...] = () + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage=stage) units = collect_b12x_units(worker, workload) @@ -621,6 +656,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if requests else None, + request_workspace_lanes=_request_workspace_lanes(units), ) outcome = coordinator.status() while not outcome["done"]: diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 11c16cfe4227..25e6ea3d2618 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -94,6 +94,7 @@ class B12xPreparationUnit: requests: tuple[PreparationRequest, ...] stage: Literal["weights", "state"] autotune: bool = True + workspace_lanes: tuple[int, ...] = (0,) def __post_init__(self): if not isinstance(self.name, str) or not self.name: @@ -103,6 +104,10 @@ def __post_init__(self): raise ValueError("invalid native preparation stage") if type(self.autotune) is not bool: raise TypeError("autotune must be boolean") + lanes = tuple(self.workspace_lanes) + if not lanes or any(type(lane) is not int or lane < 0 for lane in lanes): + raise ValueError("workspace lanes must be nonnegative integers") + object.__setattr__(self, "workspace_lanes", tuple(sorted(set(lanes)))) requests = tuple(self.requests) names = [request.name for request in requests] if len(names) != len(set(names)): @@ -210,7 +215,7 @@ def wrap_calls(calls): ) for request in unit.requests ) - return replace(unit, requests=requests) + return replace(unit, requests=requests, workspace_lanes=(lane,)) def b12x_preparation_token_counts( diff --git a/vllm/v1/worker/b12x_startup.py b/vllm/v1/worker/b12x_startup.py index 275bb0bf79fd..49fd5f51b5d7 100644 --- a/vllm/v1/worker/b12x_startup.py +++ b/vllm/v1/worker/b12x_startup.py @@ -38,11 +38,13 @@ def __init__( world_group, process_local_only: bool = False, workspace=None, + request_workspace_lanes: dict[str, tuple[int, ...]] | None = None, ) -> None: if type(global_rank) is not int or global_rank < 0: raise ValueError("global_rank must be a nonnegative integer") self.session = session self._workspace = workspace + self._request_workspace_lanes = request_workspace_lanes self._active_requests = () self.global_rank = global_rank self.world_group = world_group @@ -275,13 +277,25 @@ def _advance_local(self) -> None: job.result().close() if self._workspace is not None: + from vllm.v1.worker.workspace import use_workspace_lane + for request in self._active_requests: specs = tuple(request.plan.scratch_specs()) if specs: - self._workspace.get_simultaneous( - *((spec.shape, spec.dtype) for spec in specs) + lanes = ( + (0,) + if self._request_workspace_lanes is None + else self._request_workspace_lanes[request.name] ) - self._workspace.reserve_all() + for lane in lanes: + with use_workspace_lane(lane): + self._workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in specs) + ) + if self._request_workspace_lanes is None: + self._workspace.reserve_all() + else: + self._workspace.reserve_by_lane() self._active_requests = () self._job = None if self._batches: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43da8b0e4646..aba519befaff 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1004,15 +1004,26 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: self.pooling_runner.dummy_pooler_run(hidden_states) def _reserve_profile_scratch(self) -> None: - seen: set[int] = set() + model_lanes: dict[int, set[int]] = {} + for model, lane in ( + (self.get_model(), 0), + (self.get_draft_model(), self._draft_workspace_lane), + ): + if model is not None: + for module in model.modules(): + model_lanes.setdefault(id(module), set()).add(lane) + seen: set[tuple[int, int]] = set() for module in self.compilation_config.static_forward_context.values(): - if id(module) in seen: - continue - seen.add(id(module)) - reserve = getattr(module, "reserve_profile_scratch", None) - if reserve is not None: - reserve() - current_workspace_manager().reserve_all() + for lane in sorted(model_lanes.get(id(module), {0})): + key = (id(module), lane) + if key in seen: + continue + seen.add(key) + reserve = getattr(module, "reserve_profile_scratch", None) + if reserve is not None: + with use_workspace_lane(lane): + reserve() + current_workspace_manager().reserve_by_lane() @torch.inference_mode() def profile_run( @@ -1054,7 +1065,7 @@ def profile_run( torch.accelerator.synchronize() del hidden_states, sample_hidden_states self._profile_deepseek_v4_attention(prepare_profile_state) - current_workspace_manager().reserve_all() + current_workspace_manager().reserve_by_lane() self.reset_encoder_cache() gc.collect() diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 81adeee47065..ddde9f8b1df5 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -229,15 +229,40 @@ def reserve_all( required_bytes, max(map(self._workspace_size_bytes, self._current_workspaces), default=0), ) + self._reserve_slots( + [required_bytes] * len(self._current_workspaces), operation="reserve_all" + ) + + def reserve_by_lane(self) -> None: + """Cover every microbatch without copying capacity between model lanes. + + A target and its drafter have separate scratch contracts. A microbatch + can execute either model, but each model always uses its assigned lane. + """ + lane_bytes = [ + max( + map( + self._workspace_size_bytes, + self._current_workspaces[lane :: self._num_lanes], + ), + default=0, + ) + for lane in range(self._num_lanes) + ] + self._reserve_slots( + lane_bytes * self._num_ubatches, operation="reserve_by_lane" + ) + + def _reserve_slots(self, sizes: list[int], *, operation: str) -> None: undersized = [ workspace_id for workspace_id, workspace in enumerate(self._current_workspaces) - if self._workspace_size_bytes(workspace) < required_bytes + if self._workspace_size_bytes(workspace) < sizes[workspace_id] ] if self._locked and undersized: raise AssertionError( - "Workspace is locked but reserve_all requires " - f"{required_bytes / _MB:.2f} MB in slot(s) {undersized}." + f"Workspace is locked but {operation} requires " + f"larger capacity in slot(s) {undersized}." ) for workspace_id in undersized: @@ -246,14 +271,13 @@ def reserve_all( del current_workspace torch.accelerator.empty_cache() self._current_workspaces[workspace_id] = torch.empty( - (required_bytes,), dtype=torch.uint8, device=self._device + (sizes[workspace_id],), dtype=torch.uint8, device=self._device ) if envs.VLLM_DEBUG_WORKSPACE and undersized: logger.info( - "[WORKSPACE DEBUG] Reserved %.2f MB in execution slots %s", - required_bytes / _MB, - undersized, + "[WORKSPACE DEBUG] Reserved execution slots (slot, MB): %s", + [(slot, sizes[slot] / _MB) for slot in undersized], ) def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: @@ -327,7 +351,7 @@ def get_caller_info() -> str: # stream-ordered, so wait for the device first. Growth happens only # before the workspace is locked, never in steady-state serving. if self._device.type == "cuda" and current_workspace is not None: - torch.cuda.synchronize(self._device) + torch.accelerator.synchronize(self._device) self._current_workspaces[workspace_id] = None del current_workspace # Release the freed segment back to CUDA so the caching From e4f9b816f59999d37a32704fed5d393b363f5ba0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 03:29:17 +0000 Subject: [PATCH 007/365] fix(gdn): bind caller-owned prefill activations directly Keep only sequence metadata on each GDN layer. Bind projected inputs and invocation-owned outputs through the native variable-row contract instead of retaining scheduler-capacity copies per layer. Priming uses temporary activation storage; recurrent-state ownership and kernel arithmetic are unchanged. Validation coverage includes token-capacity-independent metadata storage and CUDA graph replay with 49 or 64 rows, changing worklists, immutable inputs, and allocation counters. GPU qualification is recorded separately. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 983c38e2ef0b8f3acaaa029f952ba1c2f26f8033) --- tests/kernels/mamba/test_gdn_prefill_b12x.py | 33 ++++++++- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 72 ++++++++++--------- .../layers/mamba/ops/b12x_gdn_prefill.py | 58 +++++---------- 3 files changed, 87 insertions(+), 76 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_prefill_b12x.py b/tests/kernels/mamba/test_gdn_prefill_b12x.py index 50d66d32764f..153722421c08 100644 --- a/tests/kernels/mamba/test_gdn_prefill_b12x.py +++ b/tests/kernels/mamba/test_gdn_prefill_b12x.py @@ -10,6 +10,7 @@ from vllm.model_executor.layers.mamba.ops.b12x_gdn_prefill import ( B12xGdnPrefill, + GdnPrefillStaging, prefill_capacities, ) @@ -27,6 +28,25 @@ def test_prefill_capacity_family_covers_exact_scheduler_limit(capacity, expected assert prefill_capacities(capacity) == expected +@pytest.mark.parametrize("max_tokens", (16, 6019, 32768)) +def test_prefill_staging_memory_depends_on_sequences_not_token_capacity(max_tokens): + staging = GdnPrefillStaging.allocate( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + assert staging.nbytes == ((16 + 1) + 4 * 16 + 2) * 4 + assert staging.is_compatible( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_prepared_gdn_decode_plan_is_scoped_to_the_bound_recurrent_pool( default_vllm_config, @@ -324,8 +344,10 @@ def test_prefill_pooled_state_graph_replays_changed_lengths_and_slots(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", (49, 64)) def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists( default_vllm_config, + rows, ): from b12x._lib.runtime_control import kernel_resolution_guard from b12x.preparation import PreparationSession @@ -353,7 +375,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pytest.skip("SM12x is required") torch.manual_seed(23) device = torch.device("cuda") - rows, seqs, key_heads, value_heads = 64, 3, 2, 6 + seqs, key_heads, value_heads = 3, 2, 6 width = (2 * key_heads + value_heads) * 128 layer = QwenGatedDeltaNetAttention.__new__(QwenGatedDeltaNetAttention) torch.nn.Module.__init__(layer) @@ -383,7 +405,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pool, ) layer._b12x_preparation_prefix = "test.mixed-gdn" - layer._b12x_prefill_max_tokens = rows + layer._b12x_prefill_max_tokens = 64 layer._b12x_prefill_max_seqs = seqs layer._b12x_prefill_staging = None layer._b12x_decode_plan = None @@ -446,6 +468,11 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists ), attn_metadata=attention, ) + immutable_inputs = { + name: value.clone() + for name, value in inputs.items() + if isinstance(value, torch.Tensor) and name != "core_attn_out" + } saved_conv, saved_pool = conv.clone(), pool.clone() def stage(lengths, computed, drafts, swap): @@ -593,6 +620,8 @@ def oracle(): ) assert_close("mixed state", pool, expected_pool, ratio=5e-3) torch.testing.assert_close(conv, expected_conv, rtol=0, atol=0) + for name, value in immutable_inputs.items(): + torch.testing.assert_close(inputs[name], value, rtol=0, atol=0) finally: guard.__exit__(None, None, None) finally: diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1169d9304d43..f75bc330347b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -959,16 +959,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): if api is None: raise RuntimeError("b12x GDN prefill was not initialized") recurrent_state = self.kv_cache[1] - staging = self._b12x_prefill_staging - resident_nbytes = 0 - if staging is not None and staging.is_compatible( - max_tokens=self._b12x_prefill_max_tokens, - max_seqs=self._b12x_prefill_max_seqs, - key_heads=self._b12x_local_key_heads, - value_heads=self._b12x_local_value_heads, - device=recurrent_state.device, - ): - resident_nbytes = staging.nbytes caps = api.Caps( device=current_platform.current_device(), max_tokens=capacity, @@ -979,8 +969,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): state_dtype=recurrent_state.dtype, checkpoint_export=True, null_state_index=0, - staging_key=(self._b12x_preparation_prefix, "gdn-prefill-staging"), - staging_resident_nbytes=resident_nbytes, ) return api.plan( caps, @@ -1148,7 +1136,6 @@ def _benchmark_b12x_gdn_prefill(self, state, capacity: int): def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): from b12x.preparation import PreparedCall - owners: tuple[torch.Tensor, ...] specs = tuple(state.layout.scratch_specs()) if len(specs) != 1: raise RuntimeError("b12x GDN prefill requires one scratch buffer") @@ -1159,11 +1146,30 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): # Trial and prepare factories own their scratch; the runtime path in # B12xGdnPrefill.run draws from the workspace manager instead. scratch = torch.empty(spec.shape, dtype=spec.dtype, device=device) + # Priming and candidate races use temporary inputs. Serving binds the + # caller's projection and output tensors directly, without retaining + # scheduler-capacity activation storage on every layer. + mixed_qkv = torch.empty( + ( + capacity, + (2 * self._b12x_local_key_heads + self._b12x_local_value_heads) * 128, + ), + dtype=torch.bfloat16, + device=device, + ) + a = torch.empty( + (capacity, self._b12x_local_value_heads), + dtype=torch.bfloat16, + device=device, + ) + b = torch.empty_like(a) + output = torch.empty( + (capacity, self._b12x_local_value_heads, 128), + dtype=torch.bfloat16, + device=device, + ) + owners: tuple[torch.Tensor, ...] if benchmark: - mixed_qkv = torch.empty_like(staging.mixed_qkv[:capacity]) - a = torch.empty_like(staging.a[:capacity]) - b = torch.empty_like(staging.b[:capacity]) - output = torch.empty_like(staging.output[:capacity]) cu_seqlens = torch.empty_like(staging.query_start_loc) indices = torch.full_like(staging.initial_indices, slot) final_indices = torch.full_like(staging.final_indices, slot) @@ -1187,12 +1193,6 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): num_tokens, ) else: - mixed_qkv = staging.mixed_qkv[:capacity] - a, b, output = ( - staging.a[:capacity], - staging.b[:capacity], - staging.output[:capacity], - ) cu_seqlens = staging.query_start_loc indices, final_indices = staging.initial_indices, staging.final_indices checkpoint_indices, offsets = ( @@ -2346,6 +2346,15 @@ def _forward_core( raise RuntimeError( "b12x GDN prefill cache and runtime metadata must be bound" ) + prefill_output = torch.empty( + ( + conv_output_prefill.shape[0], + self._b12x_local_value_heads, + self.head_v_dim, + ), + device=conv_output_prefill.device, + dtype=conv_output_prefill.dtype, + ) runner.run( mixed_qkv=conv_output_prefill, a=a_prefill, @@ -2355,12 +2364,10 @@ def _forward_core( has_initial_state=prefill_has_initial_state, live_counts=attn_metadata.b12x_prefill_live_counts, checkpoint=attn_metadata.prefill_checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) - core_attn_out_non_spec = runner.output[ - : conv_output_prefill.shape[0] - ].unsqueeze(0) + core_attn_out_non_spec = prefill_output.unsqueeze(0) else: initial_state = ssm_state[prefill_state_indices] initial_state[~prefill_has_initial_state, ...] = 0 @@ -2989,6 +2996,7 @@ def _forward_core_b12x_mixed( query_start_loc=metadata.query_start_loc, metadata=metadata.convolution_metadata(rows), ).transpose(0, 1) + prefill_output = torch.empty_like(core_attn_out) runner.run( mixed_qkv=convolved, a=a.index_select(0, non_spec_indices), @@ -2998,22 +3006,22 @@ def _forward_core_b12x_mixed( has_initial_state=metadata.has_initial_state, live_counts=metadata.live_counts, checkpoint=metadata.checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) self._rms_norm_gated_cuda( - runner.output[:rows], + prefill_output, output_gate.index_select(0, non_spec_indices), - runner.output[:rows], + prefill_output, ) core_attn_out.zero_() width = core_attn_out.shape[-2] * core_attn_out.shape[-1] _scatter_b12x_gdn_output[(rows, triton.cdiv(width, 256))]( - runner.output, + prefill_output, non_spec_indices, metadata.live_counts, core_attn_out, - SOURCE_STRIDE=runner.output.stride(0), + SOURCE_STRIDE=prefill_output.stride(0), OUTPUT_STRIDE=core_attn_out.stride(0), WIDTH=width, BLOCK=256, diff --git a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py index b13ad1904238..67fd3fd419bd 100644 --- a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py +++ b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py @@ -69,16 +69,16 @@ def prefill_capacities(max_tokens: int) -> tuple[int, ...]: @dataclass(frozen=True) class GdnPrefillStaging: - """Reusable capacity buffers independent of a recurrent-state generation.""" + """Sequence metadata independent of a recurrent-state generation. + + Activations and outputs remain caller-owned. Native GDN bindings accept + fewer rows than a plan's capacity, so no padded activation copies are needed. + """ max_tokens: int max_seqs: int key_heads: int value_heads: int - mixed_qkv: torch.Tensor - a: torch.Tensor - b: torch.Tensor - output: torch.Tensor query_start_loc: torch.Tensor initial_indices: torch.Tensor final_indices: torch.Tensor @@ -97,26 +97,12 @@ def allocate( value_heads: int, device: torch.device, ) -> "GdnPrefillStaging": - mixed_qkv = torch.empty( - (max_tokens, (2 * key_heads + value_heads) * 128), - dtype=torch.bfloat16, - device=device, - ) - a = torch.empty((max_tokens, value_heads), dtype=torch.bfloat16, device=device) initial_indices = torch.zeros(max_seqs, dtype=torch.int32, device=device) return cls( max_tokens=max_tokens, max_seqs=max_seqs, key_heads=key_heads, value_heads=value_heads, - mixed_qkv=mixed_qkv, - a=a, - b=torch.empty_like(a), - output=torch.empty( - (max_tokens, value_heads, 128), - dtype=torch.bfloat16, - device=device, - ), query_start_loc=torch.zeros(max_seqs + 1, dtype=torch.int32, device=device), initial_indices=initial_indices, final_indices=torch.zeros_like(initial_indices), @@ -140,7 +126,7 @@ def is_compatible( and self.max_seqs >= max_seqs and self.key_heads == key_heads and self.value_heads == value_heads - and self.mixed_qkv.device == device + and self.query_start_loc.device == device ) @property @@ -148,10 +134,6 @@ def nbytes(self) -> int: return sum( tensor.numel() * tensor.element_size() for tensor in ( - self.mixed_qkv, - self.a, - self.b, - self.output, self.query_start_loc, self.initial_indices, self.final_indices, @@ -164,10 +146,10 @@ def nbytes(self) -> int: class B12xGdnPrefill: - """Layer-held GDN capacity family with reusable staging buffers. + """Layer-held GDN capacity family with reusable sequence metadata. The layer supplies one declared ``Plan`` for every admitted capacity. - This helper owns only capacity staging buffers; the recurrent-state pool + This helper owns only metadata staging buffers; the recurrent-state pool remains a binding-time caller resource and can be invalidated without discarding the plans or these buffers. """ @@ -217,10 +199,6 @@ def __init__( ): raise ValueError("GDN staging buffers do not cover the prepared capacity") self.staging = staging - self.mixed_qkv = staging.mixed_qkv - self.a = staging.a - self.b = staging.b - self.output = staging.output self.query_start_loc = staging.query_start_loc self.initial_indices = staging.initial_indices self.final_indices = staging.final_indices @@ -259,9 +237,6 @@ def run( capacity = next(capacity for capacity in self.capacities if rows <= capacity) plan = self.plans[capacity] (scratch,) = get_b12x_scratch_buffers(plan) - self.mixed_qkv[:rows].copy_(mixed_qkv) - self.a[:rows].copy_(a) - self.b[:rows].copy_(b) _stage_metadata[(1,)]( query_start_loc, state_indices, @@ -281,18 +256,18 @@ def run( MAX_SEQS=self.max_seqs, BLOCK=triton.next_power_of_2(self.max_seqs + 1), ) - q, k, v = self.mixed_qkv[:capacity].split( + q, k, v = mixed_qkv.split( (self.key_heads * 128, self.key_heads * 128, self.value_heads * 128), dim=-1, ) binding = self.api.bind( plan, scratch=scratch, - q=q.view(capacity, self.key_heads, 128), - k=k.view(capacity, self.key_heads, 128), - v=v.view(capacity, self.value_heads, 128), - a=self.a[:capacity], - b=self.b[:capacity], + q=q.view(rows, self.key_heads, 128), + k=k.view(rows, self.key_heads, 128), + v=v.view(rows, self.value_heads, 128), + a=a, + b=b, A_log=self.A_log, dt_bias=self.dt_bias, recurrent_state=self.recurrent_state, @@ -303,13 +278,12 @@ def run( checkpoint_offsets=self.checkpoint_offsets, num_seqs=self.num_seqs, num_tokens=self.num_tokens, - output=self.output[:capacity], + output=output[:rows], ) self.api.run( binding, scale=scale, eps=eps, - max_live_tokens=capacity, + max_live_tokens=rows, max_live_seqs=self.max_seqs, ) - output[:rows].copy_(self.output[:rows]) From 465bcece5afeb12504aaee7d8e819d9da0dbcb68 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:15:40 +0000 Subject: [PATCH 008/365] fix(ds4): reserve attention scratch from prepared declarations State preparation reserves the actual decode and single-pass extend plans before KV admission. Remove the independent split-envelope reservation from metadata-free profiling while retaining model-lane ownership and locked workspace bounds. Validation: metadata plans require at most 282 MiB for TP2 K5 attention with 4096 prefill rows and 48 decode rows, versus a 1633 MiB independent split-envelope estimate. Added metadata-free profile regression coverage; native and complete-model qualification pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit bd00a4acf031745194b43401e30776e2ed37fa7e) --- tests/v1/worker/test_workspace.py | 45 ++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 109 +------------------------ 2 files changed, 47 insertions(+), 107 deletions(-) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 1699fe90ea7f..d31428e706e7 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -313,3 +313,48 @@ def test_suspended_graph_resources_restore_collector_after_failure() -> None: raise ValueError("eager operation failed") assert workspace.retain_cuda_graph_capture_resource(owner) assert resources == [owner] + + +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( + monkeypatch, image_tokens +) -> None: + """Prepared state declarations own attention scratch before KV admission.""" + from b12x.attention import compressed_sparse_mla + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4 + + requested = [] + monkeypatch.setattr( + workspace, + "_manager", + SimpleNamespace(get_simultaneous=lambda *specs: requested.append(specs)), + ) + monkeypatch.setattr( + dsv4, "_require_b12x_compressed_sparse_mla", lambda: compressed_sparse_mla + ) + monkeypatch.setattr( + dsv4, "get_forward_context", lambda: SimpleNamespace(attn_metadata=None) + ) + layer = dsv4.DeepseekV4B12xAttention.__new__(dsv4.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.compress_ratio = 128 + layer.max_model_len = 1048576 + layer.max_num_batched_tokens = 4096 + layer.window_size = 128 + layer.max_image_tokens = image_tokens + layer.swa_cache_layer = SimpleNamespace(block_size=256) + layer.vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + use_dspark=lambda: True, num_speculative_tokens=5 + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4096, max_num_seqs=8), + ) + q = torch.randn((3, 32, 512), dtype=torch.bfloat16) + original_q = q.clone() + output = torch.full_like(q, float("nan")) + layer.forward_mqa(q, torch.empty(0), torch.arange(3), output) + + torch.testing.assert_close(q, original_q, rtol=0, atol=0) + assert torch.count_nonzero(output) == 0 + assert requested == [] diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 232d84b2dc94..14512c22cde3 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -4,9 +4,7 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass, replace -from functools import cache from typing import TYPE_CHECKING, Any, ClassVar, cast import torch @@ -60,7 +58,6 @@ DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.worker.workspace import ( - current_workspace_manager, retain_cuda_graph_capture_resource, ) @@ -558,24 +555,6 @@ def _c128a_profile_widths(max_width: int) -> tuple[int, ...]: return tuple(sorted(widths)) -@cache -def _max_q_chunks( - max_rows: int, - width: int, - split_chunks_for_contract: Callable[..., int], - decode_row_capacity: int | None, -) -> int: - return max( - rows - * split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - for rows in range(1, max(int(max_rows), 1) + 1) - ) - - def _cache_page_view( cache: torch.Tensor, page_size: int, @@ -1289,16 +1268,6 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: out = tensor_model_parallel_all_reduce(out) return out - def reserve_profile_scratch(self) -> None: - super().reserve_profile_scratch() - device = self.q_norm.weight.device - if device.type == "cuda": - # Reserve every layer type before compiled profile execution can - # skip its attention body. This tensor describes heads/device only. - self._reserve_profile_workspace( - torch.empty((0, self.padded_heads, _DSV4_HEAD_DIM), device=device) - ) - def _get_cache_page_view( self, cache: torch.Tensor, @@ -1312,81 +1281,6 @@ def _get_cache_page_view( self._b12x_cache_page_views[key] = view return view - def _reserve_profile_workspace(self, q: torch.Tensor) -> None: - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) - - module = _require_b12x_compressed_sparse_mla() - indexed_width = 0 - if self.compress_ratio == 4: - if self.topk_indices_buffer is not None: - indexed_width = int(self.topk_indices_buffer.shape[-1]) - elif self.indexer is not None: - indexed_width = int(self.indexer.topk_tokens) - elif self.compress_ratio > 1: - indexed_width = _c128a_topk_width( - self.max_model_len, - self.compress_ratio, - ) - - indexed_widths = ( - _c128a_profile_widths(indexed_width) - if self.compress_ratio > 4 - else (indexed_width,) - ) - swa_widths = { - int(self.window_size), - int(self.window_size) + self.max_image_tokens, - } - speculative_config = self.vllm_config.speculative_config - if speculative_config is not None and speculative_config.use_dspark(): - swa_widths.add( - get_dspark_swa_index_width( - int(self.window_size), - speculative_config.num_speculative_tokens or 0, - ) - ) - swa_width = max(swa_widths) - width = max(swa_width + indexed_width, 1) - rows = max(int(self.max_num_batched_tokens), 1) - decode_row_capacity = _get_dspark_decode_row_capacity(self.vllm_config) - max_chunks_per_row = module.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - # Split count is not monotonic in index width: a shorter prefix may - # use 12-token chunks while a longer prefix uses 64-token chunks. - # Cover every metadata width, not only the longest supported context. - max_q_chunks = max( - _max_q_chunks( - rows, - max(swa + indexed, 1), - module.split_chunks_for_contract, - decode_row_capacity, - ) - for swa in swa_widths - for indexed in indexed_widths - ) - # Allocation profiling needs geometry, not an executable declaration - # tied to live cache storage or an autotuning preparation session. - plan = plan_compressed_sparse_mla_scratch( - module.Caps( - device=q.device, - num_q_heads=int(q.shape[1]), - max_q_rows=rows, - max_width=width, - head_dim=_DSV4_HEAD_DIM, - v_head_dim=_DSV4_HEAD_DIM, - page_size=int(self.swa_cache_layer.block_size), - max_chunks_per_row=max_chunks_per_row, - max_q_chunks=max_q_chunks, - decode_row_capacity=decode_row_capacity, - ) - ) - current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) - def forward_mqa( self, q: torch.Tensor, @@ -1404,7 +1298,8 @@ def forward_mqa( attn_metadata = get_forward_context().attn_metadata if attn_metadata is None: output.zero_() - self._reserve_profile_workspace(q) + # State preparation reserves each declared operation's scratch + # before KV admission. This metadata-free pass skips attention. return assert isinstance(attn_metadata, dict) From c72d80c234fa0912bedd807132f62daa209a45dd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:20:04 +0000 Subject: [PATCH 009/365] test(ds4): retain query padding reservation with prepared scratch Verify that the B12X attention class inherits per-microbatch padded query reservation while attention scratch remains owned by prepared declarations. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 842bac6e0232a7796a01c9a3d0456e277a2ed9d9) --- tests/v1/worker/test_workspace.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index d31428e706e7..e5df32783a15 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -358,3 +358,31 @@ def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( torch.testing.assert_close(q, original_q, rtol=0, atol=0) assert torch.count_nonzero(output) == 0 assert requested == [] + + +def test_dsv4_profile_retains_padded_query_reservation() -> None: + from vllm.models.deepseek_v4.attention import DeepseekV4Attention + from vllm.models.deepseek_v4.nvidia.b12x import DeepseekV4B12xAttention + + assert ( + DeepseekV4B12xAttention.reserve_profile_scratch + is DeepseekV4Attention.reserve_profile_scratch + ) + layer = DeepseekV4B12xAttention.__new__(DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.kv_cache_torch_dtype = torch.uint8 + device = torch.device("cuda:0") + layer.q_norm = SimpleNamespace(weight=SimpleNamespace(device=device)) + layer._q_padded_scratch_num_ubatches = 2 + layer.max_num_batched_tokens = 4096 + layer.padded_heads = 32 + layer.head_dim = 512 + layer._q_padded_scratch_dtype = torch.bfloat16 + reserved = [] + layer._reserve_q_padded_scratch_buffer = lambda *args: reserved.append(args) + + layer.reserve_profile_scratch() + + assert reserved == [ + (4096, 32, 512, torch.bfloat16, device, ubatch) for ubatch in range(2) + ] From 9487ef5fe8aea53ad8a05a3ad20d1a06d98eca6b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 06:23:19 +0000 Subject: [PATCH 010/365] test(ds4): validate declared attention workspace envelopes Exercise compressed-prefix, speculative, and multimodal scratch declarations through the preparation coordinator. Verify that both microbatch workspaces cover serving shapes after locking, without invoking the removed legacy split-envelope allocator. Import SimpleNamespace in the standalone workspace suite. Production code and serving behavior are unchanged. Validation: standalone test failure reproducer and pre-commit hooks. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 4b7b86dd057e334bc5fb953eea732ee395bfdc15) --- .../attention/test_ds4_workspace_envelope.py | 149 +++++++++++------- 1 file changed, 94 insertions(+), 55 deletions(-) diff --git a/tests/v1/attention/test_ds4_workspace_envelope.py b/tests/v1/attention/test_ds4_workspace_envelope.py index 404915c1469f..71c8a568b776 100644 --- a/tests/v1/attention/test_ds4_workspace_envelope.py +++ b/tests/v1/attention/test_ds4_workspace_envelope.py @@ -1,76 +1,115 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from math import prod from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from vllm.models.deepseek_v4.nvidia import b12x as adapter +from vllm.utils.b12x import B12xWorkload +from vllm.v1.worker import workspace as workspace_module +from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator +from vllm.v1.worker.workspace import WorkspaceManager @pytest.mark.skipif(not torch.cuda.is_available(), reason="B12X device plan required") @pytest.mark.parametrize("max_length", [786_688, 1_048_576]) @pytest.mark.parametrize("drafts", [0, 5]) -def test_profile_covers_shorter_compressed_prefixes(monkeypatch, max_length, drafts): - """Shorter prefixes can select more split chunks than the longest prefix.""" - from b12x.attention import compressed_sparse_mla as mla - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_declared_workspace_covers_shorter_compressed_prefixes( + monkeypatch, max_length, drafts, image_tokens +): + """Reserve decode-prefix and capacity-prefill plans before workspace lock. + Native plan scratch requirements are evaluated without compiling attention + kernels. The job stub completes only the coordinator's control protocol; + declaration, plan selection, and workspace reservation are real. + """ + device = torch.device("cuda", torch.accelerator.current_device_index()) spec = SimpleNamespace(use_dspark=lambda: True, num_speculative_tokens=drafts) - layer = SimpleNamespace( - compress_ratio=128, + layer = adapter.DeepseekV4B12xAttention.__new__(adapter.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.prefix = "model.layers.0.self_attn" + layer.indexer = None + layer.compress_ratio = 128 + layer.max_model_len = max_length + layer.max_num_batched_tokens = 4096 + layer.max_image_tokens = image_tokens + layer.window_size = 128 + layer.padded_heads = 32 + layer._b12x_cache_page_views = {} + layer._b12x_mla_plans = {} + layer.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=256), + speculative_config=spec if drafts else None, + scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=4096), + ) + page_bytes = adapter._DSV4_CACHE_BYTES_PER_TOKEN + layer.swa_cache_layer = SimpleNamespace( + block_size=256, + kv_cache=torch.empty((2, 256 * page_bytes), dtype=torch.uint8, device=device), + ) + layer.kv_cache = torch.empty((2, 2 * page_bytes), dtype=torch.uint8, device=device) + workload = B12xWorkload( + stage="state", + token_counts=(1, 8, 48, 128, 256, 4096), + fixed_token_counts=(1, 8, 48, 128, 256), + output_dtype=torch.bfloat16, + max_tokens=4096, + max_seqs=8, max_model_len=max_length, - max_num_batched_tokens=4096, - max_image_tokens=0, - window_size=128, - vllm_config=SimpleNamespace( - speculative_config=spec if drafts else None, - scheduler_config=SimpleNamespace( - max_num_seqs=8, max_num_batched_tokens=4096 - ), - ), - swa_cache_layer=SimpleNamespace(block_size=256), + speculative_tokens=drafts, ) - reserved = [] - monkeypatch.setattr( - adapter, - "current_workspace_manager", - lambda: SimpleNamespace( - get_simultaneous=lambda *items: reserved.append( - sum(prod(shape) * dtype.itemsize for shape, dtype in items) - ) - ), + (unit,) = layer.get_b12x_preparation_units(layer, workload) + queries = [request.plan.query for request in unit.requests] + widths = {128, 256, 512, 1024, 2048, 4096} + widths.add(6272 if max_length == 786_688 else 8192) + assert {q.indexed_width for q in queries} == widths + assert {q.query_rows for q in queries if q.mode == "decode"} == set(range(1, 257)) + assert {q.query_rows for q in queries if q.mode == "extend"} == {4096} + + job = Mock() + job.advance.return_value = SimpleNamespace( + done=True, pending_compilation=False, ready_collectives=() ) - adapter.DeepseekV4B12xAttention._reserve_profile_workspace( - layer, torch.empty((0, 32, 512), device="cuda") + session = Mock() + session.begin.return_value = job + workspace = WorkspaceManager(device, num_ubatches=2) + coordinator = B12xPreparationCoordinator( + session, + [(unit.requests, False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=workspace, ) - for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): - for index_width in (128, 256, 512, 1024, 2048, 4096): - width = 128 + index_width - capacity = 8 * (1 + drafts) if drafts else None - plan = plan_compressed_sparse_mla_scratch( - mla.Caps( - device=torch.device("cuda"), - num_q_heads=32, - max_q_rows=rows, - max_width=width, - head_dim=512, - v_head_dim=512, - page_size=256, - decode_row_capacity=capacity, - max_chunks_per_row=mla.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=capacity, - ), + outcome = coordinator.advance() + assert outcome["done"] and outcome["error"] is None, outcome + workspace.lock() + pointers = [buffer.data_ptr() for buffer in workspace._current_workspaces] + sizes = [buffer.numel() for buffer in workspace._current_workspaces] + assert sizes[0] == sizes[1] > 0 + + for ubatch in range(2): + monkeypatch.setattr( + workspace_module, "dbo_current_ubatch_id", lambda ubatch=ubatch: ubatch + ) + for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): + mode = "decode" if rows <= 256 else "extend" + swa_widths = {q.swa_width for q in queries if q.mode == mode} + for swa_width in swa_widths: + for index_width in widths: + plan = layer._b12x_mla_plan( + mode, + rows, + torch.empty((rows, swa_width), device="meta"), + torch.empty((rows, index_width), device="meta"), + ) + workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in plan.scratch_specs()) + ) + assert ( + workspace._current_workspaces[ubatch].data_ptr() == pointers[ubatch] ) - ) - required = sum( - prod(shape) * dtype.itemsize - for shape, dtype in plan.shapes_and_dtypes() - ) - assert max(reserved) >= required, (rows, width, reserved, required) From ee49230d89ebf644f882d9d8f8d183e779c57256 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 09:53:36 +0000 Subject: [PATCH 011/365] fix(deepseek-v41): preserve tool namespaces and reminder messages Canonicalize explicit namespaces before Chat Completions validation can discard them. Keep definitions, named choices, assistant history and streamed responses on the same qualified function identity. Reject conflicting declarations without mutating input dictionaries, and accept the reference latest_reminder role in the V4.1 tokenizer. Retain ordinary tool names, reasoning budgets, image ordering and the existing parser. Add pinned reference prompt fixtures, API validation and streaming roundtrip coverage. Validation: 179 unit tests, pre-commit hooks, TP4 temperature-1 Chat Completions tool/history checks and image smoke checks pass. Reference contract: Qizhou Guo / DeepSeek, HF DeepSeek-V4.1-Flash dba1be0a40aa45a94ad051997016db3960a90277. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5b5dd1c055d4f9561d7bbfe221f004da73815448) --- docs/features/tool_calling.md | 42 ++ .../openai/test_tool_calls_serialization.py | 162 ++++++- tests/parser/engine/test_deepseek_v4.py | 93 ++++ .../fixtures/deepseek_v41/encoding.json | 410 ++++++++++++++++++ tests/tokenizers_/test_deepseek_v41.py | 66 +++ .../openai/chat_completion/protocol.py | 33 +- vllm/tokenizers/deepseek_v41.py | 9 +- vllm/tokenizers/deepseek_v41_encoding.py | 8 +- vllm/utils/tool_names.py | 67 +++ 9 files changed, 882 insertions(+), 8 deletions(-) create mode 100644 tests/tokenizers_/fixtures/deepseek_v41/encoding.json create mode 100644 vllm/utils/tool_names.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 9628bc431a1e..2e8178d392ad 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -338,6 +338,48 @@ Supported models: Flags: `--tool-call-parser deepseek_v31 --chat-template {see_above}` +### DeepSeek-V4.1 Models (`deepseek_v41`) + +Status: implemented. The `deepseek_v41` tokenizer and parser support spaced +DSML tool calls and namespace-qualified function names for +`deepseek-ai/DeepSeek-V4.1-Flash`. + +Use `--tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 +--reasoning-parser deepseek_v41 --enable-auto-tool-choice`. + +Chat Completions accepts a namespace on a function tool: + +```json +{ + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "lookup", + "description": "Look up an item.", + "parameters": {"type": "object", "properties": {"sku": {"type": "string"}}} + } +} +``` + +The namespace may also be a string, or appear inside `function`. Before +validation and rendering, it becomes `function.name="inventory::lookup"`. +The namespace description is prepended to the function description. +The same normalization applies to assistant tool-call history and named +`tool_choice` objects. A qualified name must agree with any explicit namespace; +conflicting declarations are rejected. Tools without namespaces are unchanged. + +Complete and streamed Chat Completions responses retain the qualified +`function.name`; they do not add a separate response `namespace` field. Dispatch +tools by that qualified name and preserve it when sending assistant history. +This extension is separate from the Responses API's `type="namespace"` schema. + +The tokenizer also accepts `{"role":"latest_reminder","content":"..."}` +between a system message and a user message, rendering the reference +`<｜latest_reminder｜>` marker for date, locale or similar context. This is a +DeepSeek-specific role, not a standard OpenAI message role. Other unknown roles +remain invalid. The encoding contract follows the +[DeepSeek V4.1 reference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/dba1be0a40aa45a94ad051997016db3960a90277/encoding/README.md). + ### OpenAI OSS Models (`openai`) Supported models: diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index cedc2575b80e..478379f5e13b 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -15,9 +15,16 @@ "ValueError: Unexpected tool call id ...". """ +import copy + import pytest +from pydantic import ValidationError -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.exceptions import VLLMValidationError def _make_tool_call(tc_id: str, name: str, args: str) -> dict: @@ -148,3 +155,156 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): # Verify after model_dump_json too _ = req.model_dump_json() assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +@pytest.mark.parametrize("iterator", [False, True]) +def test_namespaces_survive_definitions_choices_and_history( + location, namespace, iterator +): + def qualify(item, value): + target = item if location == "tool" else item["function"] + target["namespace"] = value + return item + + tools = [ + qualify({"type": "function", "function": {"name": "lookup"}}, namespace), + qualify({"type": "function", "function": {"name": "lookup"}}, "billing"), + ] + calls = [qualify(_make_tool_call("call_1", "lookup", "{}"), namespace)] + payload = { + "model": "DeepSeek-V4.1-Flash", + "tools": tools, + "tool_choice": qualify( + {"type": "function", "function": {"name": "lookup"}}, namespace + ), + "messages": [ + {"role": "user", "content": "Look up stock."}, + {"role": "assistant", "content": None, "tool_calls": calls}, + ], + } + original = copy.deepcopy(payload) + if iterator: + payload["messages"][1]["tool_calls"] = iter(calls) + request = ChatCompletionRequest.model_validate(payload) + assert [tool.function.name for tool in request.tools] == [ + "inventory::lookup", + "billing::lookup", + ] + assert request.tool_choice.function.name == "inventory::lookup" + assert request.messages[1]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + serialized = request.model_dump_json() + assert request.model_dump_json() == serialized + restored = ChatCompletionRequest.model_validate_json(serialized) + assert restored.model_dump_json() == serialized + if not iterator: + assert payload == original + assert calls == original["messages"][1]["tool_calls"] + + +def test_namespace_description_is_preserved_exactly_once(): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "inventory::lookup", + "description": "Find a SKU.", + "strict": True, + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == "Stock operations.\nFind a SKU." + assert parsed.function.strict is True + assert tool == original + assert ChatCompletionToolsParam.model_validate( + parsed.model_dump() + ).model_dump() == (parsed.model_dump()) + + +@pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) +@pytest.mark.parametrize( + "namespace, name", + [ + ("inventory", "billing::lookup"), + ("a::b", "lookup"), + ("", "lookup"), + ({}, "lookup"), + ({"name": 7}, "lookup"), + (False, "lookup"), + ("inventory", "inventory::x::lookup"), + ("inventory", "inventory::"), + ], +) +def test_invalid_namespaces_are_client_errors(namespace, name, tool_choice): + with pytest.raises(ValidationError, match="[Nn]amespace"): + ChatCompletionRequest.model_validate( + { + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "namespace": namespace, + "function": {"name": name}, + } + ], + "tool_choice": tool_choice, + } + ) + + +def test_conflicting_namespace_locations_are_rejected(): + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup", "namespace": "billing"}, + } + ) + + +def test_consistent_namespace_locations_do_not_duplicate_the_prefix(): + parsed = ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": {"name": "inventory"}, + "function": {"name": "inventory::lookup", "namespace": "inventory"}, + } + ) + assert parsed.function.name == "inventory::lookup" + + +def test_conflicting_history_namespace_is_not_silently_dropped(): + call = _make_tool_call("call_a", "inventory::lookup", "{}") + call["namespace"] = "billing" + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + _make_request([{"role": "assistant", "content": None, "tool_calls": [call]}]) + + +def test_bare_choice_cannot_select_a_namespaced_tool(): + with pytest.raises(VLLMValidationError, match="does not match"): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "lookup"}], + "tools": [ + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup"}, + } + ], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + } + ) + + +def test_plain_and_already_qualified_names_keep_their_spelling(): + for name in ("lookup", "inventory::lookup", "opaque::legacy::name"): + tool = {"type": "function", "function": {"name": name}} + assert ChatCompletionToolsParam.model_validate(tool).function.name == name diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 7cbf9791a7b0..9180a6f4b2d5 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -21,6 +21,7 @@ simulate_reasoning_streaming, simulate_tool_streaming, ) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.deepseek_v4 import ( DSML_INVOKE_END, @@ -36,10 +37,13 @@ _unwrap_wrapper_args, deepseek_v4_config, ) +from vllm.parser.deepseek_v41 import DeepSeekV41Parser +from vllm.parser.engine.adapters import make_adapters from vllm.parser.engine.registered_adapters import ( DeepSeekV4ParserReasoningAdapter, DeepSeekV4ParserToolAdapter, ) +from vllm.tokenizers.deepseek_v41_encoding import encode_messages _THINK_START_ID = 50 _THINK_END_ID = 51 @@ -48,6 +52,95 @@ _PARAM_CLOSE = "" +@pytest.mark.parametrize("chunk_size", [1, 2, 7, 17, 1000]) +def test_v41_namespaced_tools_keep_identity_in_full_stream_and_history(chunk_size): + names = ["inventory::lookup", "billing::lookup"] + request = ChatCompletionRequest( + messages=[{"role": "user", "content": "Look up A1."}], + tools=[ + { + "type": "function", + "namespace": namespace, + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": {"sku": {"type": "string"}}, + }, + }, + } + for namespace in ("inventory", "billing") + ], + ) + tokenizer = make_mock_tokenizer({"": 50, "": 51}) + _, adapter = make_adapters(DeepSeekV41Parser) + + def parser(): + return adapter( + tokenizer, request.tools, chat_template_kwargs={"thinking": False} + ) + + text = ( + "<｜DSML｜ calls>\n" + + "\n".join( + f'<｜DSML｜ invoke name="{name}">\n' + '<｜DSML｜ parameter name="sku" string="true">A1\n' + "" + for name in names + ) + + "\n" + ) + calls = parser().extract_tool_calls(text, request).tool_calls + assert [call.function.name for call in calls] == names + assert [json.loads(call.function.arguments) for call in calls] == [ + {"sku": "A1"} + ] * 2 + stream = parser() + collected: dict[int, dict[str, str]] = {} + + def collect(delta): + if delta is None: + return + for call in delta.tool_calls: + value = collected.setdefault(call.index, {"name": "", "arguments": ""}) + if call.function: + value["name"] += call.function.name or "" + value["arguments"] += call.function.arguments or "" + + previous = "" + for offset in range(0, len(text), chunk_size): + chunk = text[offset : offset + chunk_size] + collect( + stream.extract_tool_calls_streaming( + previous, + previous + chunk, + chunk, + (), + (), + (), + request, + ) + ) + previous += chunk + collect(stream.finish_streaming()) + assert [collected[i]["name"] for i in sorted(collected)] == names + assert [json.loads(collected[i]["arguments"]) for i in sorted(collected)] == ( + [{"sku": "A1"}] * 2 + ) + prompt = encode_messages( + [ + {"role": "user", "content": "Look up A1."}, + { + "role": "assistant", + "content": "", + "tool_calls": [call.model_dump(exclude_none=True) for call in calls], + }, + ], + thinking_mode="chat", + ) + assert all(f'name="{name}"' in prompt for name in names) + + def _param(name: str, is_str: str, value: str) -> str: return f"<{_PARAM_OPEN.format(name=name, is_str=is_str)}{value}{_PARAM_CLOSE}" diff --git a/tests/tokenizers_/fixtures/deepseek_v41/encoding.json b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json new file mode 100644 index 000000000000..532e842e96aa --- /dev/null +++ b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json @@ -0,0 +1,410 @@ +{ + "reference": "https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/raw/dba1be0a40aa45a94ad051997016db3960a90277/encoding/encoding.py", + "reference_author": "DeepSeek; namespace contract by Qizhou Guo", + "cases": [ + { + "name": "reminder_chat", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>You are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "reminder_thinking", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\nYou are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "mid_system_chat", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "mid_system_thinking", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "reasoning_history_chat", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "reasoning_history_thinking", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "interleaved_images_chat", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "interleaved_images_thinking", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "namespaced_tool_history_false", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "namespaced_tool_history_true", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>Consult both tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "action_task", + "messages": [ + { + "role": "user", + "content": "Find the documentation.", + "task": "action" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Find the documentation.<｜Assistant｜><｜action｜>" + }, + { + "name": "title_task", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi.", + "task": "title" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜title｜>" + } + ] +} diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index d0916fb6077b..5035a7e71431 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/tests/tokenizers_/test_deepseek_v41.py @@ -4,12 +4,14 @@ import copy import json from pathlib import Path +from typing import Any import pytest from tests.tokenizers_.test_deepseek_v4 import FakeHfTokenizer from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tokenizers.deepseek_v41 import get_deepseek_v41_tokenizer +from vllm.tokenizers.deepseek_v41_encoding import encode_messages FIXTURES = Path(__file__).parent / "fixtures" @@ -185,3 +187,67 @@ def image_part(color): def test_unsupported_media_is_a_request_error(): with pytest.raises(ValueError, match="text and image content only"): render([{"role": "user", "content": [{"type": "input_audio"}]}]) + + +_FIXTURES = json.loads( + (Path(__file__).parent / "fixtures/deepseek_v41/encoding.json").read_text() +) + + +@pytest.mark.parametrize("case", _FIXTURES["cases"], ids=lambda case: case["name"]) +def test_prompt_matches_reference_fixture(case): + original = copy.deepcopy(case["messages"]) + request = ChatCompletionRequest( + messages=case["messages"], + tools=case.get("tools"), + ) + prompt = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tools=[tool.model_dump(exclude_none=True) for tool in request.tools] + if request.tools + else None, + tokenize=False, + **case["kwargs"], + ) + assert prompt == case["expected"] + assert case["messages"] == original + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +def test_direct_encoder_and_api_preserve_the_same_tool_identity(location, namespace): + tool: dict[str, Any] = {"type": "function", "function": {"name": "lookup"}} + call: dict[str, Any] = { + "id": "call_a", + "type": "function", + "function": {"name": "lookup", "arguments": '{"sku":"A1"}'}, + } + for item in (tool, call): + (item if location == "tool" else item["function"])["namespace"] = namespace + messages = [ + {"role": "system", "content": "Use tools.", "tools": [tool]}, + {"role": "user", "content": "Find A1."}, + {"role": "assistant", "content": "", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call_a", "content": "In stock."}, + ] + original = copy.deepcopy(messages) + direct = encode_messages(messages, thinking_mode="chat") + request = ChatCompletionRequest(messages=messages) + rendered = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tokenize=False, + thinking=False, + ) + assert rendered == direct + assert '"name": "inventory::lookup"' in rendered + assert 'name="inventory::lookup"' in rendered + assert "In stock." in rendered + assert messages == original + + +def test_reminder_does_not_enable_arbitrary_message_roles(): + with pytest.raises(ValueError, match="Invalid role: SYSTEM"): + get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + [{"role": "SYSTEM", "content": "Hello"}], + tokenize=False, + ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index badd5a0344c6..f9d07659cdb8 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -51,6 +51,7 @@ ThinkingTokenBudget, ) from vllm.utils import random_uuid +from vllm.utils.tool_names import normalize_tool_namespace logger = init_logger(__name__) @@ -186,6 +187,11 @@ class ChatCompletionToolsParam(OpenAIBaseModel): function: FunctionDefinition defer_loading: bool | None = None + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + @model_validator(mode="after") def _propagate_defer_loading(self) -> "ChatCompletionToolsParam": if self.defer_loading is not None and self.function.defer_loading is None: @@ -209,6 +215,11 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): function: ChatCompletionNamedFunction type: Literal["function"] = "function" + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + class ChatCompletionRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation @@ -534,23 +545,31 @@ def check_cache_salt_support(cls, data: Any) -> Any: def _normalize_messages_before(cls, data: Any) -> Any: """Pre-process message dicts before Pydantic field validation. - Performs two normalizations in a single pass: + Normalizes tool identities and message fields in a single pass: - Converts tool_calls generators/iterators to lists so one-shot generators are not consumed during union type matching. - Renames the deprecated ``reasoning_content`` field to ``reasoning`` so downstream code only needs to check one field. + - Qualifies tool-call namespaces before union validation can drop them. """ if not isinstance(data, dict): return data messages = data.get("messages") if not isinstance(messages, list): return data + data = dict(data) + messages = [dict(msg) if isinstance(msg, dict) else msg for msg in messages] + data["messages"] = messages for msg in messages: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None and not isinstance(tool_calls, list): - msg["tool_calls"] = list(tool_calls) + if tool_calls is not None: + msg["tool_calls"] = [ + normalize_tool_namespace(call) for call in tool_calls + ] + if isinstance(msg.get("tools"), list): + msg["tools"] = [normalize_tool_namespace(tool) for tool in msg["tools"]] reasoning_content = msg.pop("reasoning_content", None) if reasoning_content is not None and msg.get("reasoning") is None: msg["reasoning"] = reasoning_content @@ -916,6 +935,14 @@ def check_tool_usage(cls, data): if not isinstance(data, dict): return data + # Named choices must be compared with qualified identities before + # Pydantic serializes function definitions and discards extra fields. + data = dict(data) + if isinstance(data.get("tools"), list): + data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + if isinstance(data.get("tool_choice"), dict): + data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) + # Reject empty tools array, matching OpenAI API behavior if data.get("tools") == []: raise VLLMValidationError( diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index 3405cbbbe7dc..68ba93e5ec8a 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -23,7 +23,14 @@ def _normalize_messages( result = [dict(message) for message in copy.deepcopy(messages)] for message in result: role = message.get("role") - if role not in ("system", "developer", "user", "assistant", "tool"): + if role not in ( + "system", + "developer", + "user", + "assistant", + "tool", + "latest_reminder", + ): raise ValueError(f"Invalid role: {role}") if role == "developer": message["role"] = "system" diff --git a/vllm/tokenizers/deepseek_v41_encoding.py b/vllm/tokenizers/deepseek_v41_encoding.py index 995115f1f7d0..89907d152bef 100644 --- a/vllm/tokenizers/deepseek_v41_encoding.py +++ b/vllm/tokenizers/deepseek_v41_encoding.py @@ -11,6 +11,8 @@ import json from typing import Any, Dict, List, Optional, Union +from vllm.utils.tool_names import normalize_tool_namespace + bos_token: str = "<｜begin▁of▁sentence｜>" eos_token: str = "<｜end▁of▁sentence｜>" thinking_start_token: str = "" @@ -55,8 +57,8 @@ def to_json(value: Any) -> str: def tools_from_openai_format(tools): - """Extract function definitions from OpenAI-format tool list.""" - return [tool["function"] for tool in tools] + """Extract function definitions with namespace-qualified identities.""" + return [normalize_tool_namespace(tool)["function"] for tool in tools] def tool_calls_from_openai_format(tool_calls): @@ -66,7 +68,7 @@ def tool_calls_from_openai_format(tool_calls): "name": tool_call["function"]["name"], "arguments": tool_call["function"]["arguments"], } - for tool_call in tool_calls + for tool_call in map(normalize_tool_namespace, tool_calls) ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py new file mode 100644 index 000000000000..02fcbc9e7c2e --- /dev/null +++ b/vllm/utils/tool_names.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Canonical function identities for namespace-bearing chat tools. + +The namespace spelling follows Qizhou Guo's DeepSeek V4.1 reference encoder: +https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/commit/dba1be0a40aa45a94ad051997016db3960a90277 +""" + +from typing import Any + + +def _namespace_parts(namespace: Any) -> tuple[str, str | None]: + description = None + if isinstance(namespace, dict): + description = namespace.get("description") + namespace = namespace.get("name") + if not isinstance(namespace, str) or not namespace or "::" in namespace: + raise ValueError("Tool namespace must be a nonempty string without '::'") + if description is not None and not isinstance(description, str): + raise ValueError("Tool namespace description must be a string") + return namespace, description + + +def normalize_tool_namespace(tool: Any) -> Any: + """Qualify explicit namespaces without mutating a definition or tool call. + + Accept a string or {name, description} namespace on the tool or function. + Consistent duplicate declarations are allowed; conflicting identities are + rejected. Remove namespace metadata after folding it into the function name + and description, so normalization is idempotent across API and tokenizer + boundaries. Inputs without namespace metadata retain their existing names. + """ + if not isinstance(tool, dict) or not isinstance(tool.get("function"), dict): + return tool + function = tool["function"] + declarations = [ + _namespace_parts(value) + for value in (tool.get("namespace"), function.get("namespace")) + if value is not None + ] + if not declarations: + return tool + namespace, description = declarations[0] + if any(name != namespace for name, _ in declarations[1:]): + raise ValueError("Conflicting tool namespaces on the tool and function") + + name = function.get("name") + if not isinstance(name, str) or not name: + raise ValueError("Namespaced function name must be a nonempty string") + if "::" in name: + prefix, name = name.split("::", 1) + if prefix != namespace: + raise ValueError("Conflicting tool namespaces in the qualified name") + if not name or "::" in name: + raise ValueError("Namespaced function name must be nonempty without '::'") + + normalized = dict(tool) + normalized.pop("namespace", None) + normalized["function"] = function = dict(function) + function.pop("namespace", None) + function["name"] = f"{namespace}::{name}" + if description: + detail = function.get("description") + if detail is not None and not isinstance(detail, str): + raise ValueError("Function description must be a string") + function["description"] = description + "\n" + (detail or "") + return normalized From c5596fa2023fdd42b8f05e9f2bb92b914d19ff63 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 10:55:43 +0000 Subject: [PATCH 012/365] fix(tools): preserve validation paths and namespace descriptions Normalize supported tool-call iterables without preempting scalar field validation. Retain the first non-null namespace description while rejecting conflicting names. Focused regressions reproduce both review findings; 50 tool-call serialization tests pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit fb50f44a5f025c9c2eed857fbac938e8513227a7) --- .../openai/test_tool_calls_serialization.py | 37 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 5 ++- vllm/utils/tool_names.py | 7 +++- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 478379f5e13b..156d6bb140c4 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -226,6 +226,43 @@ def test_namespace_description_is_preserved_exactly_once(): ).model_dump() == (parsed.model_dump()) +@pytest.mark.parametrize("tool_calls", [1, False, 1.5]) +def test_invalid_tool_calls_shape_reports_the_field(tool_calls): + with pytest.raises(ValidationError) as error: + _make_request( + [{"role": "assistant", "content": None, "tool_calls": tool_calls}] + ) + assert any("tool_calls" in item["loc"] for item in error.value.errors()) + + +@pytest.mark.parametrize( + "outer, inner, expected", + [ + (None, "Stock operations.", "Stock operations.\nFind a SKU."), + ("Outer description.", "Inner description.", "Outer description.\nFind a SKU."), + ("", "Inner description.", "Find a SKU."), + ], +) +def test_consistent_namespaces_preserve_first_non_null_description( + outer, inner, expected +): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": outer}, + "function": { + "name": "lookup", + "namespace": {"name": "inventory", "description": inner}, + "description": "Find a SKU.", + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == expected + assert tool == original + assert ChatCompletionToolsParam.model_validate(parsed.model_dump()) == parsed + + @pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) @pytest.mark.parametrize( "namespace, name", diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index f9d07659cdb8..e5913efef6c6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -4,6 +4,7 @@ # Adapted from # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py import time +from collections.abc import Iterable from typing import Annotated, Any, ClassVar, Literal from openai.types.chat.chat_completion_audio import ( @@ -564,7 +565,9 @@ def _normalize_messages_before(cls, data: Any) -> Any: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None: + if isinstance(tool_calls, Iterable) and not isinstance( + tool_calls, (str, bytes, dict) + ): msg["tool_calls"] = [ normalize_tool_namespace(call) for call in tool_calls ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py index 02fcbc9e7c2e..1cf3f368d464 100644 --- a/vllm/utils/tool_names.py +++ b/vllm/utils/tool_names.py @@ -41,8 +41,11 @@ def normalize_tool_namespace(tool: Any) -> Any: if not declarations: return tool namespace, description = declarations[0] - if any(name != namespace for name, _ in declarations[1:]): - raise ValueError("Conflicting tool namespaces on the tool and function") + for declared_name, declared_description in declarations[1:]: + if declared_name != namespace: + raise ValueError("Conflicting tool namespaces on the tool and function") + if description is None: + description = declared_description name = function.get("name") if not isinstance(name, str) or not name: From 811dbf09e8b2f7ed3adbd07ff8ec247d8aff31c6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 11:13:52 +0000 Subject: [PATCH 013/365] perf(comm): index declared PCIe collective plans by their full call contract Bound lookup storage to declared shapes and avoid scans and exception construction when unsupported prefill shapes fall through to another transport. Preserve normalization weight identity, epsilon, strides and first-declaration precedence. Inspired by fernandaspets/vllm_sm120#2. Native suite: 36 passed including TP2 fused collective replay, one four-GPU case skipped. Serving throughput is not yet qualified. Co-authored-by: neolithic5452 Co-authored-by: Neo Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5e693878601f1145e41582999439080687cc0bf1) --- benchmarks/benchmark_b12x_plan_lookup.py | 81 ++++++++++++++++ .../distributed/test_b12x_pcie_all_reduce.py | 86 +++++++++++++++++ .../b12x_pcie_all_reduce.py | 93 ++++++++++++------- 3 files changed, 229 insertions(+), 31 deletions(-) create mode 100644 benchmarks/benchmark_b12x_plan_lookup.py diff --git a/benchmarks/benchmark_b12x_plan_lookup.py b/benchmarks/benchmark_b12x_plan_lookup.py new file mode 100644 index 000000000000..fc5c8fe0dcba --- /dev/null +++ b/benchmarks/benchmark_b12x_plan_lookup.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure CPU-only lookup cost for declared B12X PCIe collective shapes. + +No GPU operation is executed. Use --module-file to compare pinned source files +in the same Python/runtime environment; these timings are not serving throughput. +""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +import sys +import timeit +from pathlib import Path + +import torch + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--module-file", type=Path, required=True) + parser.add_argument("--declarations", type=int, nargs="+", default=[64, 512, 4096]) + parser.add_argument("--iterations", type=int, default=1000) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + if min(*args.declarations, args.iterations, args.repeats) < 1: + parser.error("counts must be positive") + spec = importlib.util.spec_from_file_location("pcie_lookup_probe", args.module_file) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + report = { + "source_sha256": hashlib.sha256(args.module_file.read_bytes()).hexdigest(), + "iterations": args.iterations, + "repeats": args.repeats, + "scope": "CPU metadata lookup; no collective or model execution", + "results": [], + } + for count in args.declarations: + communicator = object.__new__(module.B12xPcieAllReduce) + invocations = [ + module.B12xPcieInvocation( + name=str(rows), + operation="all_reduce", + shape=(rows, 16), + dtype=torch.bfloat16, + ) + for rows in range(1, count + 1) + ] + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + if hasattr(communicator, "_index_declared_plans"): + communicator._index_declared_plans() + for case, rows, expected in (("hit", count, True), ("miss", count + 1, False)): + source = torch.empty((rows, 16), dtype=torch.bfloat16) + + def lookup(communicator=communicator, source=source): + return communicator._has_plan_for(source) + + assert lookup() is expected + for _ in range(100): + lookup() + samples = timeit.repeat(lookup, number=args.iterations, repeat=args.repeats) + report["results"].append( + { + "declarations": count, + "case": case, + "microseconds_per_call": [ + s * 1e6 / args.iterations for s in samples + ], + "median_us": statistics.median(samples) * 1e6 / args.iterations, + } + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_b12x_pcie_all_reduce.py b/tests/distributed/test_b12x_pcie_all_reduce.py index 4a4b6a21eeda..017883b4a3e5 100644 --- a/tests/distributed/test_b12x_pcie_all_reduce.py +++ b/tests/distributed/test_b12x_pcie_all_reduce.py @@ -58,9 +58,11 @@ def _make_communicator( communicator.twoshot_max_bytes = 0 plan = object() communicator._plans = {"prepared": plan} + communicator._plan_index = {} communicator._routes = {} communicator._invocations = {} communicator._plan_for = MagicMock(return_value=plan) + communicator._lookup_plan = MagicMock(return_value=plan) return communicator, runtime @@ -323,6 +325,7 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: communicator, _ = _make_communicator() communicator._plans = {} communicator._invocations = {} + communicator._lookup_plan.return_value = None with pytest.raises( PreparationResourceUnavailableError, @@ -331,6 +334,89 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: B12xPcieAllReduce._plan_for(communicator, torch.randn(2, 4)) +def _indexed_communicator(invocations): + communicator = object.__new__(B12xPcieAllReduce) + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + communicator._index_declared_plans() + return communicator + + +def test_declared_plan_index_preserves_complete_fused_identity(): + operation = "all_reduce_fused_add_rms_norm" + weight_a = torch.ones(4) + weight_b = weight_a.clone() + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocations = [ + b12x_pcie_all_reduce.B12xPcieInvocation( + name=name, + operation=operation, + shape=(2, 4), + dtype=source.dtype, + norm_weight=weight, + epsilon=epsilon, + ) + for name, weight, epsilon in ( + ("a", weight_a, 1e-6), + ("b", weight_b, 1e-6), + ("epsilon", weight_a, 1e-5), + ) + ] + communicator = _indexed_communicator(invocations) + for invocation in invocations: + assert ( + communicator._plan_for( + source, + operation=operation, + weight=invocation.norm_weight, + epsilon=invocation.epsilon, + ) + is communicator._plans[invocation.name] + ) + assert not communicator._has_plan_for(source) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a.clone(), epsilon=1e-6 + ) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a, epsilon=1e-4 + ) + + +def test_declared_plan_index_matches_shape_dtype_stride_and_first_declaration(): + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocation = b12x_pcie_all_reduce.B12xPcieInvocation + communicator = _indexed_communicator( + [ + invocation("first", "all_reduce", (2, 4), source.dtype), + invocation( + "equivalent", "all_reduce", (2, 4), source.dtype, strides=(4, 1) + ), + invocation("strided", "all_reduce", (2, 4), source.dtype, strides=(8, 1)), + ] + ) + assert communicator._plan_for(source) is communicator._plans["first"] + strided = torch.empty((2, 8), dtype=source.dtype)[:, :4] + assert communicator._plan_for(strided) is communicator._plans["strided"] + assert not communicator._has_plan_for(source.float()) + assert not communicator._has_plan_for(source.flatten()) + communicator._invocations = {} + assert communicator._plan_for(source) is communicator._plans["first"] + communicator._index_declared_plans() + assert not communicator._has_plan_for(source) + + +def test_undeclared_plan_probe_neither_raises_nor_grows_the_index(): + communicator = _indexed_communicator([]) + for rows in range(1, 65): + source = torch.empty((rows, 4), dtype=torch.bfloat16) + assert not communicator._has_plan_for(source) + with pytest.raises( + PreparationResourceUnavailableError, match="no declared plan" + ): + communicator._plan_for(source) + assert communicator._plan_index == {} + + def test_fused_allreduce_has_an_independent_cutoff() -> None: communicator, runtime = _make_communicator( allreduce_max_bytes=16, fused_max_bytes=64 diff --git a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py index 16cf5b31fadf..e82d1510731f 100644 --- a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py @@ -248,6 +248,7 @@ def __init__( ] ] = [] self._plans: dict[str, object] = {} + self._plan_index: dict[tuple, object] = {} self._invocations: dict[str, B12xPcieInvocation] = {} self._routes: dict[str, str] = {} if len(self.global_ranks) != self.world_size: @@ -596,12 +597,14 @@ def get_b12x_preparation_units( plans: dict[str, object] = {} routes: dict[str, str] = {} collective_ranks = tuple(sorted(self.global_ranks)) - declarations = {} + declarations: dict[tuple, Any] = {} + prepare: Callable[..., Any] for invocation in invocations: route = self._route_invocation(invocation) if route is None: continue routes[invocation.name] = route + assert self._runtime is not None target = self._runtime._prepared_channel_for_stream( None, invocation.channel_id ) @@ -630,7 +633,7 @@ def get_b12x_preparation_units( plan = _oneshot_preparation.plan(query, runtime=target) declarations[key] = plan - def prepare(state, invocation=invocation, query=query): + def prepare_oneshot(state, invocation=invocation, query=query): registered = bool(query.setup["registered"]) inp = self._request_input(invocation, registered=registered) reset, produce, restore, owners = self._oneshot_prime_input( @@ -665,6 +668,8 @@ def produce_fused() -> None: restore=restore, owners=owners, ) + + prepare = prepare_oneshot elif route == "twoshot": assert self._twoshot is not None query = _twoshot_preparation.query_from_metadata( @@ -681,13 +686,15 @@ def produce_fused() -> None: plan = _twoshot_preparation.plan(query, runtime=self._twoshot) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_twoshot(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) return _twoshot_preparation.prepared_call( state, payload=inp, out=torch.empty_like(inp) ) + + prepare = prepare_twoshot else: assert self._dma is not None query = _dma_preparation.query_from_metadata( @@ -703,7 +710,7 @@ def prepare(state, invocation=invocation): plan = _dma_preparation.plan(query, runtime=self._dma) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_dma(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) @@ -711,6 +718,7 @@ def prepare(state, invocation=invocation): state, inp=inp, out=torch.empty_like(inp) ) + prepare = prepare_dma plans[invocation.name] = plan requests.append( plan.request( @@ -732,6 +740,7 @@ def prepare(state, invocation=invocation): } self._plans = plans self._routes = routes + self._index_declared_plans() if not requests: return () return ( @@ -743,6 +752,47 @@ def prepare(state, invocation=invocation): ), ) + @staticmethod + def _plan_key(operation, shape, dtype, strides, weight=None, epsilon=None): + norm = None if operation == "all_reduce" else (id(weight), epsilon) + return operation, tuple(shape), dtype, tuple(strides), norm + + def _index_declared_plans(self) -> None: + index: dict[tuple, object] = {} + for name, invocation in self._invocations.items(): + key = self._plan_key( + invocation.operation, + invocation.shape, + invocation.dtype, + invocation.strides + if invocation.strides is not None + else self._contiguous_strides(invocation.shape), + invocation.norm_weight, + invocation.epsilon, + ) + # Equivalent declarations keep the first plan, matching model order. + index.setdefault(key, self._plans[name]) + self._plan_index = index + + def _lookup_plan( + self, + inp: torch.Tensor, + *, + operation: str = "all_reduce", + weight: torch.Tensor | None = None, + epsilon: float | None = None, + ): + return self._plan_index.get( + self._plan_key( + operation, + inp.shape, + inp.dtype, + inp.stride(), + weight, + epsilon, + ) + ) + def _plan_for( self, inp: torch.Tensor, @@ -751,32 +801,16 @@ def _plan_for( weight: torch.Tensor | None = None, epsilon: float | None = None, ): - strides = tuple(inp.stride()) - for name, invocation in self._invocations.items(): - if ( - invocation.operation == operation - and tuple(inp.shape) == invocation.shape - and inp.dtype == invocation.dtype - and strides - == ( - self._contiguous_strides(invocation.shape) - if invocation.strides is None - else invocation.strides - ) - and ( - operation == "all_reduce" - or ( - invocation.norm_weight is weight - and invocation.epsilon == epsilon - ) - ) - ): - return self._plans[name] + plan = self._lookup_plan( + inp, operation=operation, weight=weight, epsilon=epsilon + ) + if plan is not None: + return plan from vllm.utils.b12x import PreparationResourceUnavailableError raise PreparationResourceUnavailableError( f"PCIe collective has no declared plan for {operation} on shape " - f"{tuple(inp.shape)} {inp.dtype} strides {strides}; declared: " + f"{tuple(inp.shape)} {inp.dtype} strides {tuple(inp.stride())}; declared: " + ", ".join( f"{invocation.operation}{invocation.shape}" for invocation in self._invocations.values() @@ -791,11 +825,7 @@ def _has_plan_for(self, inp: torch.Tensor, **kwargs) -> bool: the captured and planned serving shapes, so an undeclared shape is a prefill chunk or warm-up batch outside the graph-replayed set. """ - from vllm.utils.b12x import PreparationResourceUnavailableError - - try: - self._plan_for(inp, **kwargs) - except PreparationResourceUnavailableError: + if self._lookup_plan(inp, **kwargs) is None: logger.debug( "b12x PCIe all-reduce declines undeclared shape %s %s", tuple(inp.shape), @@ -953,6 +983,7 @@ def close(self) -> None: self._runtime.close() self._runtime = None self._plans.clear() + self._plan_index.clear() self._invocations.clear() self._routes.clear() self.disabled = True From b1247feb061afc6a055400994035f811fab1ae5e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:36:27 +0000 Subject: [PATCH 014/365] fix(glm): invalidate physical-selection plans with cache bindings Preserve the KK prepared-state call interface and invalidate geometry-dependent declarations on cache bind/unbind. Extend the existing construction test to cover declaration lifetime. Reconciles the remaining ownership part of vLLM #789; native qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 8 ++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 5 +++++ 2 files changed, 13 insertions(+) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 843fb7e1356d..ffb61df5ec9e 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -751,9 +751,17 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( _, main = _packed_main_cache( device=torch.device("cpu"), blocks=2, layers=1, block_size=256, layer=0 ) + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() indexer.bind_main_kv_cache(main) assert indexer.indexer_op._num_q_heads == indexer._q_fp8.shape[1] == 32 + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() + indexer.unbind_main_kv_cache() + assert indexer._physical_selection_plan is None + assert indexer._index_cache is None + assert indexer._main_cache_num_blocks == 0 def test_glm53_packed_tail_accepts_nvfp4_main_record() -> None: diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index cc0311138364..3f15e9b2ce77 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -143,6 +143,7 @@ def __init__( self._expand_pooled_topk_to_physical_slots = ( b12x_sparse_mla.expand_pooled_topk_to_physical_slots ) + self._physical_selection_plan = None self.index_kpool_compress_ape = nn.Parameter( allocate_weights( torch.empty, @@ -360,6 +361,7 @@ def _index_cache_view( return index_cache, subpages_per_parent, parent_stride_pages def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: + self._physical_selection_plan = None index_cache, subpages, parent_stride_pages = self._index_cache_view(main_cache) block_size = int(main_cache.shape[1]) parent_table_width = self._max_parent_table_width( @@ -393,6 +395,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: def unbind_main_kv_cache(self) -> None: self._index_cache = None self._main_cache_num_blocks = 0 + self._physical_selection_plan = None def make_b12x_physical_selection_prepare_call( self, output: torch.Tensor, active_counts: torch.Tensor, *, state=None @@ -707,6 +710,8 @@ def forward( if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: if self._main_cache_num_blocks < 1: raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From 04bea2ea87b4ad493ed7e2441b467fdcfd366aab Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:49:49 +0000 Subject: [PATCH 015/365] test(b12x): prepare native cache writers in graph-lifetime coverage Keep Karmic cache preparation contracts while checking that graph resources release caller activations. Preserve parser stream chunks as explicit concatenated values. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/parser/engine/test_deepseek_v4.py | 12 ++++++--- tests/v1/attention/test_b12x_v41_workspace.py | 25 ++++++++++--------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 9180a6f4b2d5..5eb369cee969 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -305,8 +305,10 @@ def test_streaming_split_next_parameter_tag_is_buffered( chunks = [ DSML_TOOL_START, f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n", - f"<{_PARAM_OPEN.format(name='location', is_str='true')}" - "Paris a<｜DSML｜parameter", + ( + f"<{_PARAM_OPEN.format(name='location', is_str='true')}" + "Paris a<｜DSML｜parameter" + ), ' name="date" string="true">tomorrow', _PARAM_CLOSE, DSML_INVOKE_END, @@ -376,8 +378,10 @@ def test_streaming_with_trailing_content(self, mock_tokenizer, mock_request): parser = DeepSeekV4Parser(mock_tokenizer) chunks = [ DSML_TOOL_START, - f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" - f"{_param('location', 'true', 'NYC')}\n", + ( + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" + f"{_param('location', 'true', 'NYC')}\n" + ), DSML_TOOL_END, "Done.", ] diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 260f386aa122..847536b72ae4 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -787,12 +787,13 @@ def metadata(page): device=device, dtype=torch.uint8, ) - from b12x.attention._shared.mla.compressed_reference import ( - pack_deepseek_v41_cache_reference, - ) - - cache[1:].copy_( - pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) + _write_cache( + kv, + cache, + torch.arange(length, device=device) + page, + page_size=page, + cache_kind=kind, + cache_format="deepseek_v41", ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -813,6 +814,12 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 + _write_index_keys( + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + page_size=main_page, + ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -836,12 +843,6 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) - attention.dsa_indexer.quantize_write_index_k_mxfp4( - layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - ) activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): From 8c5ec192580f8944846e89167ee9547677e3e513 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:23:02 +0000 Subject: [PATCH 016/365] test(ds41): declare attention bounds in WO capture fixtures Prepare the output-projection fixture with the same declared attention envelope as the model. Both graph-lifetime cases pass on SM120 with the KK native wheel. Co-authored-by: OpenAI Codex --- tests/v1/attention/test_b12x_v41_workspace.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 847536b72ae4..c0c33fa5a341 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -1047,6 +1047,7 @@ def test_output_projection_capture_releases_caller_activations( ), ) layer.setup_wo_projection() + layer._declare_attention(device) workload = B12xWorkload( stage="weights", token_counts=(8, 16), From 3eefaaf43bc896098b7745f9b22d3aa5f3ab7ed6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:37:43 +0000 Subject: [PATCH 017/365] test(b12x): preserve Karmic preparation and precision contracts Observe caller-owned MHC outputs while requiring preparation not to retain them; retain allocation, pointer and graph replay assertions. Exercise the generic ModelOpt W4A16 method and padded native scales. Initialize the complete NVFP4 fixture. Use an IEEE FP32 reference instead of TF32 for unquantized projections; existing numerical tolerances are unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../kernels/test_b12x_linear.py | 52 ++++++++++++++----- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 6a1a540400e3..5a2e6d0a7db7 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -493,7 +493,8 @@ def pack(weight: torch.Tensor, weight_scale: torch.Tensor): assert len(calls) == 1 weight, weight_scale = calls[0] assert weight.shape == (48, 128) - assert weight_scale.shape == (48, 4) + # Preserve the checkpoint's padded row layout for native packed scales. + assert weight_scale.shape == (64, 4) assert weight.dtype == torch.float8_e4m3fn assert weight_scale.dtype == torch.uint8 assert layer.weight.numel() == 0 @@ -933,7 +934,14 @@ def test_b12x_fp4_processes_scale_and_preserves_loader( if kernel_cls is B12xNvFp4LinearKernel: layer.weight = torch.empty((48, 64), dtype=torch.uint8) layer.weight_global_scale = torch.tensor(0.5) - packed = object() + layer.input_global_scale_inv = torch.tensor(1.0) + + @dataclass(frozen=True) + class PackedWeight: + in_features: int + values: torch.Tensor + + packed = PackedWeight(0, layer.weight) def pack_weight(weight, scale, *, recipe, global_scale): assert weight.data_ptr() == layer.weight.data_ptr() @@ -948,13 +956,16 @@ def pack_weight(weight, scale, *, recipe, global_scale): lambda: types.SimpleNamespace(pack_weight=pack_weight), ) kernel = object.__new__(kernel_cls) + if kernel_cls is B12xNvFp4LinearKernel: + kernel.config = NvFp4LinearLayerConfig(use_a16=False) kernel.process_weights_after_loading(layer) assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() assert layer.weight_scale.weight_loader is weight_loader if kernel_cls is B12xNvFp4LinearKernel: - assert layer.b12x_nvfp4_packed_weight is packed + assert layer.b12x_nvfp4_packed_weight.in_features == 128 + assert layer.b12x_nvfp4_packed_weight.values is packed.values def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: @@ -1055,7 +1066,7 @@ def test_b12x_w4a16_modelopt_vision_width_preserves_bf16_activations(monkeypatch from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, - ModelOptNvFp4W4A16LinearMethod, + build_linear_method, ) from vllm.v1.worker.workspace import ( current_workspace_manager, @@ -1077,11 +1088,13 @@ def reject_activation_quantization(*args, **kwargs): n, k = 1152, 4304 config = VllmConfig(kernel_config=KernelConfig(linear_backend="b12x")) with set_current_vllm_config(config), torch.no_grad(): - method = ModelOptNvFp4W4A16LinearMethod( + method = build_linear_method( ModelOptNvFp4Config( quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True, - ) + ), + "W4A16_NVFP4", + "vision_tower.mlp", ) layer = torch.nn.Module() with torch.device(device): @@ -1932,10 +1945,13 @@ def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( name = f"_b12x_mhc_{native_operation}_impl" native = getattr(_impl, name) captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + captured_outputs = [] def observe(*args, **kwargs): captured_call[:] = [(args, kwargs)] - return native(*args, **kwargs) + result = native(*args, **kwargs) + captured_outputs[:] = outputs(result) + return result def outputs(value): return value if isinstance(value, tuple) else (value,) @@ -1948,7 +1964,10 @@ def outputs(value): with monkeypatch.context() as hooks: hooks.setattr(_impl, name, observe) call.produce() - first = outputs(call.invoke()) + assert call.invoke() is None + first = tuple(captured_outputs) + assert call.invoke() is None + second = tuple(captured_outputs) args, kwargs = captured_call.pop() reference_kwargs = { key: value @@ -1956,7 +1975,6 @@ def outputs(value): if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} } pointers = [tensor.data_ptr() for tensor in first] - second = outputs(call.invoke()) assert [tensor.data_ptr() for tensor in second] == pointers for actual, expected in zip( second, outputs(native(*args, **reference_kwargs)), strict=True @@ -1975,9 +1993,9 @@ def outputs(value): graph.replay() torch.accelerator.synchronize() assert torch.accelerator.memory_allocated(device) == allocated - assert [ - tensor.data_ptr() for tensor in outputs(call.output) - ] == pointers + # Preparation must not pin full-prefill results in the session. + assert call.output is None + assert [tensor.data_ptr() for tensor in first] == pointers expected = outputs(native(*args, **reference_kwargs)) for actual, wanted in zip(first, expected, strict=True): assert torch.isfinite(actual).all() @@ -2722,7 +2740,7 @@ def expected(rows): assert set(plans) == {4, capacity} actual = run(11) assert plans[11].prepared is not None - assert plans[11].selection.source == "fixed" + assert plans[11].selection.source == "default" torch.testing.assert_close(actual, expected(11), rtol=0.02, atol=0.125) session.freeze() with kernel_resolution_guard("FP8 prepared exact-M execution"): @@ -2746,11 +2764,17 @@ def expected(rows): @pytest.mark.parametrize("output_dtype", [torch.bfloat16, torch.float32]) @torch.no_grad() -def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay(output_dtype): +def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay( + output_dtype, request +): if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: pytest.skip("native b12x projection requires SM12x") from vllm.models.deepseek_v4_1.b12x_layers import B12xLinearMethod + # A TF32 reference rounds FP32 inputs, unlike the native SIMT projection. + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = torch.nn.Module() layer.weight = torch.nn.Parameter( From ee97b07da6a0d8878bc582365ff08e5cc6733827 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 22:04:32 +0000 Subject: [PATCH 018/365] test(ds41): model native FP8 prefill arithmetic in CED oracle Compare compact replay with an independent BF16 QK / FP8 PV reference and assert the complete causal candidate set. Retain numerical tolerances, poisoned discarded rows, padded outputs, caller scratch and CUDA graph checks. Production kernels and precision are unchanged. Validation: 333 native tests passed, three multi-GPU tests separately passed; all three CED geometries pass at unchanged tolerances. Reference adapted from B12X commit 4af78b86771d8491ec7b9d27168fdfbc1d67a773. Existing mypy and CUDA-API hook findings are unchanged. Co-authored-by: Luke Alonso Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_v41_workspace.py | 129 ++++++++++++++---- 1 file changed, 106 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index c0c33fa5a341..66ea9c7bd3e6 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -701,6 +701,73 @@ def test_mhc_fixed_capacity_buckets_preserve_decode_policy( assert bool(torch.isfinite(y).all()) +def _v41_fp8_operands(cache, kind): + """Decode native cache rows into E4M3 operands and per-64 scales. + + This independent operand reference follows B12X's + tests/_reference/v41_fp8.py. SWA pairs its E8M0 scales; indexed NVFP4 + uses a power-of-two scale that bounds four adjacent scale groups. + """ + if kind == "swa": + rows = cache.reshape(-1, 528) + codes = rows[:, :512].contiguous().view(torch.float8_e4m3fn).double() + original = rows[:, 512:528].double() + exponent = original.reshape(-1, 8, 2).amax(-1).clamp_min(1) - 127 + values = codes * torch.exp2(original - 127).repeat_interleave(32, dim=1) + else: + rows = cache.reshape(-1, 288) + packed = rows[:, :256].long() + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(1) + lut = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], + dtype=torch.float64, + device=cache.device, + ) + original = rows[:, 256:288].contiguous().view(torch.float8_e4m3fn).double() + values = lut[codes] * original.repeat_interleave(16, dim=1) + bound = original.reshape(-1, 8, 4).amax(-1) * (6 / 448) + exponent = torch.ceil(torch.log2(bound.clamp_min(2.0**-126))) + scales = torch.exp2(exponent) + codes = (values / scales.repeat_interleave(64, dim=1)).float() + return codes.to(torch.float8_e4m3fn).float(), scales.float() + + +def _v41_fp8_prefill_reference(q, values, scales, valid, sink): + """Model BF16 Q/K and per-64-key FP8 probability/value products. + + Prefill computes Q/K with BF16 operands. Each output group independently + quantizes scaled probabilities to E4M3 before its product with FP8 V; + partial numerators remain FP32 until the final LSE merge. + """ + keys = (values * scales.repeat_interleave(64, dim=-1)).bfloat16().float() + logits = torch.einsum("rhd,rkd->rhk", q.float(), keys) * 512**-0.5 + logits.masked_fill_(~valid[:, None], -torch.inf) + partials, lses = [], [] + for first in range(0, values.shape[1], 64): + local = logits[:, :, first : first + 64] + maximum = local.amax(-1) + p = torch.exp(local - maximum[:, :, None]) + p = torch.where(valid[:, None, first : first + 64], p, 0) + denominator = p.sum(-1) + groups = [] + for group in range(8): + weighted = p * scales[:, None, first : first + 64, group] + pscale = weighted.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 + pq = (weighted / pscale).to(torch.float8_e4m3fn).float() + vq = values[:, first : first + 64, group * 64 : (group + 1) * 64] + groups.append(torch.einsum("rhk,rkd->rhd", pq, vq) * pscale) + partials.append(torch.cat(groups, -1) / denominator.clamp_min(1e-30)[..., None]) + lses.append( + torch.where(denominator > 0, maximum + denominator.log(), -torch.inf) + ) + lse = torch.stack(lses, -1) + maximum = torch.maximum(lse.amax(-1), sink.float()[None]) + weights = torch.where(torch.isfinite(lse), torch.exp(lse - maximum[..., None]), 0) + denominator = weights.sum(-1) + torch.exp(sink.float()[None] - maximum) + output = (torch.stack(partials, -2) * weights[..., None]).sum(-2) + return (output / denominator.clamp_min(1e-30)[..., None]).bfloat16().float() + + @pytest.mark.parametrize("main_page,swa_page", [(64, 32), (128, 64), (256, 128)]) @pytest.mark.parametrize( "is_decode,rows,live_rows", @@ -1561,14 +1628,14 @@ def test_ced_compact_attention_bounded_oracle_and_frozen_replay( ): from contextlib import ExitStack - from b12x.attention._shared.mla.compressed_reference import ( - unpack_deepseek_v41_cache_reference, - ) from b12x.preparation import PreparationSession from vllm.utils.b12x import B12xWorkload attention, manager, workspace = native_workspace + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) torch.manual_seed(713) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = _layer(attention, 20, swa_page=swa_page) @@ -1633,17 +1700,19 @@ def metadata(page): ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache - swa_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + swa_values, swa_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) # NaN FP8 payloads in every old SWA row: reading below the # replay boundary cannot accidentally look like a valid zero page. cache[1 : 1 + boundary // page].fill_(0x7F) else: layer.kv_cache = cache - main_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + main_values, main_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) layer.indexer.k_cache.kv_cache = torch.zeros( ( length // main_page + 1, @@ -1683,21 +1752,35 @@ def run(): ) def oracle(live): - values = torch.cat((swa_values[boundary:], main_values)) - logical = torch.cat( + selected = layer.topk_indices_buffer[:live].long() + # Candidate order affects per-tile FP8 rounding. Verify its complete + # causal set independently before using that order in the numeric oracle. + columns = torch.arange(selected.shape[1], device=device)[None] + expected_ids = columns.expand_as(selected).clone() + expected_ids.masked_fill_(columns > positions[:live, None], length) + actual_ids = selected.masked_fill(selected < 0, length).sort(dim=1).values + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + swa_logical = torch.arange(boundary, length, device=device)[None] + swa_valid = swa_logical <= positions[:live, None] + main_valid = (selected >= 0) & (selected <= positions[:live, None]) + valid = torch.cat((swa_valid, main_valid), dim=1) + values = torch.cat( ( - torch.arange(boundary, length, device=device), - torch.arange(length, device=device), - ) - ) - logits = torch.einsum("rhd,kd->rhk", q[:live].float(), values) * 512**-0.5 - logits.masked_fill_( - logical[None, None] > positions[:live, None, None], -torch.inf - ) - logits = torch.cat( - (logits, layer.attn_sink[None, :, None].expand(live, -1, -1)), -1 + swa_values[None, boundary:].expand(live, -1, -1), + main_values[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + scales = torch.cat( + ( + swa_scales[None, boundary:].expand(live, -1, -1), + main_scales[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + return _v41_fp8_prefill_reference( + q[:live], values, scales, valid, layer.attn_sink ) - return torch.einsum("rhk,kd->rhd", logits.softmax(-1)[..., :-1], values) workload = B12xWorkload( stage="state", @@ -1744,9 +1827,9 @@ def oracle(live): visible.copy_((positions + 1).clamp_min(0).int()) starts[1] = live q.normal_() - expected = oracle(live) # Eager execution also catches live-count specialization leaks. run() + expected = oracle(live) torch.testing.assert_close( out[:live].float(), expected, rtol=0.04, atol=0.025 ) From fb2de2bc7de3ae6556e7acbe756c8d409f02a448 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:11:18 +0000 Subject: [PATCH 019/365] fix(dflash): initialize fused MXFP8 projection with generic ModelOpt Recognize MXFP8 through its QuantSpec weight key and give the fused context K/V projection an independent kernel lifecycle. Preserve input/output dtypes and bias metadata without changing checkpoint precision or draft sampling. Keep BF16 loading unchanged and adapt the draft-loader fixture to the shared pipeline-safe loader helper. Validation: 34 DFlash tests pass. Native B12X fused projection passes exact representable-input checks at 1/8/32 rows for quantized and A16 activations, including changed-input CUDA graph replay with zero additional allocated bytes. Applicable pre-commit hooks pass. Whole-model MXFP8 DFlash qualification remains pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/spec_decode/test_dflash2.py | 100 +++++++++++++++++++-- vllm/model_executor/models/qwen3_dflash.py | 45 ++++++---- 2 files changed, 120 insertions(+), 25 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 315f28235456..09dccc5054e9 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -14,6 +14,8 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): + from vllm.config import LoadConfig + draft_load_config = object() draft_model_config = SimpleNamespace(hf_config=SimpleNamespace()) speculative_config = SimpleNamespace( @@ -25,6 +27,7 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): vllm_config = SimpleNamespace( attention_config=SimpleNamespace(), cache_config=SimpleNamespace(), + load_config=LoadConfig(load_format="fastsafetensors"), speculative_config=speculative_config, ) loaded = SimpleNamespace(model=SimpleNamespace()) @@ -42,17 +45,16 @@ def fake_get_model(**kwargs): monkeypatch.setattr(dflash_utils, "replace", fake_replace) monkeypatch.setattr(dflash_utils, "get_model", fake_get_model) monkeypatch.setattr( - dflash_utils, - "get_pp_group", + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.compilation.backends.set_model_tag", - lambda _tag: nullcontext(), + "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group", + lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.model_executor.models.qwen3_dflash.dflash_target_rope_is_neox_style", - lambda _model: None, + "vllm.compilation.backends.set_model_tag", + lambda _tag: nullcontext(), ) monkeypatch.setattr( "vllm.model_executor.models.qwen3_dflash.dflash_has_any_non_causal", @@ -61,6 +63,8 @@ def fake_get_model(**kwargs): assert dflash_utils.load_dflash_model(SimpleNamespace(), vllm_config) is loaded assert captured["load_config"] is draft_load_config + assert captured["vllm_config"].load_config.load_format == "auto" + assert vllm_config.load_config.load_format == "fastsafetensors" def test_dflash_reset_attn_releases_cache_layout_state(): @@ -158,7 +162,7 @@ def test_dflash_context_projection_rejects_mixed_quantization(mxfp8_layer: int): from torch import nn from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, + build_linear_method, ) from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model @@ -166,8 +170,8 @@ class UnreadableWeight: def __getitem__(self, _key): raise AssertionError("weights must not be read before validation") - mxfp8_method = object.__new__(ModelOptMxFp8LinearMethod) - methods = [None, None] + mxfp8_method = build_linear_method(None, "MXFP8", "") + methods = [build_linear_method(None, "FP8", ""), None] methods[mxfp8_layer] = mxfp8_method layers_attn = [ SimpleNamespace( @@ -186,6 +190,84 @@ def __getitem__(self, _key): model._build_context_kv_buffers(layers_attn, has_bias=False) +@pytest.mark.parametrize("mxfp8", [False, True]) +@pytest.mark.parametrize("has_bias", [False, True]) +def test_fused_context_projection_owns_its_linear_method(monkeypatch, mxfp8, has_bias): + """Fused K/V packing must not reconfigure the query projection's kernel.""" + from torch import nn + + import vllm.model_executor.layers.quantization.modelopt as modelopt + import vllm.model_executor.parameter as parameter + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + kernels = [] + + def select_kernel(spec, layer, runtime_dtypes, **kwargs): + kernel = SimpleNamespace( + process_weights_after_loading=lambda layer: None, + input_quant_key=lambda: None, + ) + kernels.append((kernel, layer.output_size_per_partition)) + assert runtime_dtypes.input_dtype == torch.bfloat16 + assert runtime_dtypes.out_dtype == torch.bfloat16 + assert runtime_dtypes.marlin_input_dtype == torch.float16 + assert layer.has_bias is has_bias + return kernel + + monkeypatch.setattr(modelopt, "select_linear_kernel", select_kernel) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + method = modelopt.build_linear_method(None, "MXFP8", "") if mxfp8 else None + original_kernel = object() + if method is not None: + method.kernel = original_kernel + method.input_dtype = method.out_dtype = torch.bfloat16 + method.marlin_input_dtype = torch.float16 + layers = [] + for index in range(2): + projection = nn.Module() + projection.quant_method = method + weight = torch.full((192, 64), index + 1, dtype=torch.bfloat16) + if mxfp8: + weight = weight.to(torch.float8_e4m3fn) + projection.weight_scale = nn.Parameter( + torch.full((192, 2), 127, dtype=torch.uint8), requires_grad=False + ) + projection.weight = nn.Parameter(weight, requires_grad=False) + if has_bias: + projection.bias = nn.Parameter(torch.full((192,), float(index + 1))) + attention = SimpleNamespace( + qkv_proj=projection, + q_size=128, + k_norm=SimpleNamespace(weight=torch.ones(32)), + ) + layers.append(SimpleNamespace(self_attn=attention)) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.layers = layers + model.hidden_norm = SimpleNamespace(weight=torch.ones(64, dtype=torch.bfloat16)) + model._fused_kv_linear = nn.Module() + model._fused_kv_quant_method = None + model._build_context_kv_buffers([layer.self_attn for layer in layers], has_bias) + expected = model._fused_kv_weight.clone() + model.process_weights_after_loading() + if mxfp8: + assert method is not None + assert model._fused_kv_quant_method is not method + assert method.kernel is original_kernel + assert kernels == [(model._fused_kv_quant_method.kernel, 128)] + assert model._fused_kv_linear.weight_block_size == [1, 32] + torch.testing.assert_close( + model._fused_kv_linear.weight.float(), expected.float() + ) + assert model._fused_kv_weight is None + assert model._fused_kv_weight_scale is None + else: + assert not kernels + assert model._fused_kv_quant_method is None + torch.testing.assert_close(model._fused_kv_weight, expected) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3c7c301814ff..0ea09746a4bb 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -27,6 +27,8 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -511,13 +513,11 @@ def _build_context_kv_buffers( layers_attn: list[nn.Module], has_bias: bool, ) -> None: - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_methods = [a.qkv_proj.quant_method for a in layers_attn] uses_mxfp8 = [ - isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods + isinstance(method, ModelOptLinearMethod) + and method.spec.weight == kMxfp8Static + for method in quant_methods ] if any(uses_mxfp8) and not all(uses_mxfp8): raise ValueError( @@ -593,12 +593,11 @@ def _build_fused_kv_buffers(self) -> None: def process_weights_after_loading(self) -> None: """Pack the serialized MXFP8 context projection for its GEMM backend.""" - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_method = self.layers[0].self_attn.qkv_proj.quant_method - if not isinstance(quant_method, ModelOptMxFp8LinearMethod): + if not ( + isinstance(quant_method, ModelOptLinearMethod) + and quant_method.spec.weight == kMxfp8Static + ): return if self._fused_kv_weight is None or self._fused_kv_weight_scale is None: raise RuntimeError( @@ -607,9 +606,23 @@ def process_weights_after_loading(self) -> None: ) output_size, input_size = self._fused_kv_weight.shape - self._fused_kv_linear.input_size_per_partition = input_size - self._fused_kv_linear.output_size_per_partition = output_size - self._fused_kv_linear.logical_widths = [output_size] + # Fused context K/V has an independent weight shape and kernel + # lifecycle; do not repack through the query projection's method. + fused_method = ModelOptLinearMethod( + quant_method.spec, quant_method.ctx, quant_method.fmt + ) + fused_method.input_dtype = quant_method.input_dtype + fused_method.out_dtype = quant_method.out_dtype + fused_method.marlin_input_dtype = quant_method.marlin_input_dtype + self._fused_kv_linear.has_bias = self._fused_kv_bias is not None + fused_method.create_weights( + self._fused_kv_linear, + input_size_per_partition=input_size, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=self.hidden_norm.weight.dtype, + ) self._fused_kv_linear.register_parameter( "weight", nn.Parameter(self._fused_kv_weight, requires_grad=False) ) @@ -617,13 +630,13 @@ def process_weights_after_loading(self) -> None: "weight_scale", nn.Parameter(self._fused_kv_weight_scale, requires_grad=False), ) - quant_method.process_weights_after_loading(self._fused_kv_linear) - self._fused_kv_quant_method = quant_method + fused_method.process_weights_after_loading(self._fused_kv_linear) + self._fused_kv_quant_method = fused_method self._fused_kv_weight = None self._fused_kv_weight_scale = None logger.info_once( "Using %s for the fused DFlash context K/V projection.", - type(quant_method.kernel).__name__, + type(fused_method.kernel).__name__, ) def _project_context_kv( From 08aee0efe31fc64eb6068f5b6d44a9e5579913b2 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:21:42 +0000 Subject: [PATCH 020/365] fix(v2): register native sampling warmup contracts Register the existing top-k/top-p compile-key declarations on sampling ranks. Seeded and processed-logprob requests can select native filtering even when unseeded warmup uses FlashInfer. This covers singleton split-row kernels omitted by the bounded 8/16-row warmup; sampling arithmetic and the production JIT monitor default are unchanged. Validation: six GPU warmup cases pass for vocabularies 257/129280 and row limits 1/7/128. Every reachable row count and top-k/top-p subset runs with a hook that rejects post-warmup compilation. Applicable pre-commit hooks pass. The DS4.1 strict-monitor reproducer requires a rebuilt-image rerun. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/sample/test_topk_topp_sampler.py | 24 +++++++++++++++++++---- vllm/v1/worker/gpu/model_runner.py | 5 +++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 4d4740524de0..cbf1d7655c9e 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from torch import Generator @@ -13,6 +15,7 @@ apply_top_k_top_p_probs, apply_top_k_top_p_pytorch, random_sample, + register_top_k_top_p_warmups, warmup_top_k_top_p, ) from vllm.v1.sample.sampler import Sampler @@ -74,21 +77,34 @@ def test_apply_top_k_top_p_probs_matches_processed_logits( reason="requires GPU Triton compilation", ) @pytest.mark.parametrize("vocab_size", [257, 129280]) +@pytest.mark.parametrize("max_num_rows", [1, 7, 128]) def test_filter_warmup_covers_speculative_rows_and_constraint_subsets( - monkeypatch: pytest.MonkeyPatch, vocab_size: int + monkeypatch: pytest.MonkeyPatch, vocab_size: int, max_num_rows: int ): - """Four K5 requests must not compile top-p-only filtering after warmup.""" + """Seeded filtering must cover singleton, split-row and large-batch keys.""" from triton import knobs + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry + device = torch.device(DEVICE_TYPE) - warmup_top_k_top_p(vocab_size, 24, device) + registry = JitWarmupRegistry( + SimpleNamespace( + model_config=SimpleNamespace(get_vocab_size=lambda: vocab_size), + scheduler_config=SimpleNamespace(max_num_seqs=max_num_rows), + num_speculative_tokens=0, + ) + ) + with registry.activate(): + register_top_k_top_p_warmups() + registry.warmup() + warmup_top_k_top_p(vocab_size, max_num_rows, device) torch.accelerator.synchronize() def unexpected_compile(**kwargs): pytest.fail("Filtering compiled a kernel after bounded sampler warmup") monkeypatch.setattr(knobs.runtime, "jit_post_compile_hook", unexpected_compile) - for num_rows in (8, 9, 15, 16, 17, 24): + for num_rows in range(1, max_num_rows + 1): logits = torch.zeros(num_rows, vocab_size, device=device) top_k = torch.full((num_rows,), 20, dtype=torch.int32, device=device) top_p = torch.full((num_rows,), 0.95, device=device) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index aba519befaff..5b57f9e55146 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,6 +78,7 @@ ModelRunnerOutput, RoutedExpertsTensors, ) +from vllm.v1.sample.ops.topk_topp_sampler import register_top_k_top_p_warmups from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.watermarking.spec_decode import ( @@ -472,6 +473,10 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: + # Seeded and processed-logprob requests use native filtering even + # when unseeded warmup requests select FlashInfer sampling. + with self.jit_warmup_registry.activate(): + register_top_k_top_p_warmups() sampler_kwargs: dict[str, Any] = { "max_num_reqs": self.max_num_reqs, "vocab_size": self.vocab_size, From 0d072a46630449361f338203bdb2693106d7ee67 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:41:37 +0000 Subject: [PATCH 021/365] fix(chat): normalize namespaced tool iterables before matching Materialize Python tool collections before qualifying namespace metadata and validating named choices. Lists, tuples and generators retain the same serialized function identities without mutating definitions. JSON request behavior is unchanged. Validation: tuple inputs reproduce a false tool-choice mismatch before the fix. The namespace/history suite passes for lists, tuples and generators; combined CPU suites pass 90 tests with 43 GPU-only skips. Applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/entrypoints/openai/test_tool_calls_serialization.py | 7 +++++-- vllm/entrypoints/openai/chat_completion/protocol.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 156d6bb140c4..c62bff2d1ab6 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -160,8 +160,9 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): @pytest.mark.parametrize("location", ["tool", "function"]) @pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) @pytest.mark.parametrize("iterator", [False, True]) +@pytest.mark.parametrize("tools_collection", [list, tuple, iter]) def test_namespaces_survive_definitions_choices_and_history( - location, namespace, iterator + location, namespace, iterator, tools_collection ): def qualify(item, value): target = item if location == "tool" else item["function"] @@ -185,6 +186,7 @@ def qualify(item, value): ], } original = copy.deepcopy(payload) + payload["tools"] = tools_collection(tools) if iterator: payload["messages"][1]["tool_calls"] = iter(calls) request = ChatCompletionRequest.model_validate(payload) @@ -200,8 +202,9 @@ def qualify(item, value): assert request.model_dump_json() == serialized restored = ChatCompletionRequest.model_validate_json(serialized) assert restored.model_dump_json() == serialized - if not iterator: + if not iterator and tools_collection is list: assert payload == original + assert tools == original["tools"] assert calls == original["messages"][1]["tool_calls"] diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index e5913efef6c6..3607ec221b6a 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -941,8 +941,11 @@ def check_tool_usage(cls, data): # Named choices must be compared with qualified identities before # Pydantic serializes function definitions and discards extra fields. data = dict(data) - if isinstance(data.get("tools"), list): - data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + tools = data.get("tools") + if isinstance(tools, Iterable) and not isinstance( + tools, (str, bytes, bytearray, dict) + ): + data["tools"] = [normalize_tool_namespace(tool) for tool in tools] if isinstance(data.get("tool_choice"), dict): data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) From 058802e06ba9d41e5dfa0421882f8c10e5308aac Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:42:16 +0000 Subject: [PATCH 022/365] fix(glm): reject undeclared selection before recurrent cache updates Check physical-selection resources before dispatching pooled-cache writes on pure DCP1 decode. An unbound main cache or missing selection plan now fails without advancing request state. Supported selection arithmetic and the hot-path checks are unchanged. Validation: both missing-resource cases reproduce an early cache-write attempt before the fix and pass afterwards. All 67 pooled-indexer tests pass on stock SM120, including changed-input replay and prepared selection; applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 60 +++++++++++++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 17 ++++-- 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index ffb61df5ec9e..5b6d3ca555bc 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -689,6 +689,66 @@ def test_glm53_physical_selection_provider_is_explicit() -> None: ) +@pytest.mark.parametrize( + "main_blocks, error", + [(0, "main cache is not bound"), (1, "physical selection has not been declared")], +) +def test_glm53_missing_physical_plan_does_not_mutate_recurrent_cache( + monkeypatch, main_blocks, error +) -> None: + from vllm.models.glm5next.nvidia import pooled_indexer as module + + hidden = torch.zeros((1, 128)) + metadata = SimpleNamespace( + num_reqs=1, + num_actual_tokens=1, + num_decode_tokens=1, + num_decodes=1, + max_query_len=1, + block_table=torch.zeros((1, 1), dtype=torch.int32), + selector_state_slot_ids=torch.zeros(1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + indexer = SimpleNamespace( + max_tokens=1, + wq_b=lambda value: (torch.zeros((1, 32 * 128)), None), + wk=lambda value: (hidden, None), + k_norm=lambda value: value, + index_kpool_compress_gate=torch.zeros((128, 128)), + index_kpool_compress_ape=torch.zeros((4, 128)), + _project_head_weights=lambda value: torch.zeros((1, 32)), + _q_fp8=torch.empty((1, 32, 128), dtype=torch.float8_e4m3fn), + _q_scale=torch.empty((1, 32)), + main_layer_name="attention", + _index_cache=torch.zeros((1, 64, 132), dtype=torch.uint8), + _tail=torch.zeros((1, 2, 4, 128)), + _state_slots=Glm5NextPooledIndexer._state_slots, + _parent_table_width=1, + _parent_stride_pages=1, + block_size=256, + dcp_world_size=1, + _emit_physical_selection=True, + _main_cache_num_blocks=main_blocks, + _physical_selection_plan=None, + ) + monkeypatch.setattr(module, "fwht128_quant_fp8", lambda *args, **kwargs: None) + monkeypatch.setattr( + module, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={"attention": metadata}), + ) + + def reject_cache_update(*args, **kwargs): + pytest.fail("Missing physical-selection resources must not update the cache") + + monkeypatch.setattr(module, "update_decode_pools", reject_cache_update) + with pytest.raises(RuntimeError, match=error): + Glm5NextPooledIndexer.forward( + indexer, hidden, hidden, torch.zeros(1, dtype=torch.int64), None + ) + + def _packed_main_cache( *, device: torch.device, diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 3f15e9b2ce77..63ee0d63546f 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -568,6 +568,16 @@ def forward( f"actual={actual_table_width}, required={self._parent_table_width}" ) + decode_only = decode_rows == live_rows + emit_physical_selection = ( + decode_only and self.dcp_world_size == 1 and self._emit_physical_selection + ) + if emit_physical_selection: + if self._main_cache_num_blocks < 1: + raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") + update_decode_pools( index_cache, self._tail, @@ -586,7 +596,6 @@ def forward( ) parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] seq_lens = self._pool_seq_lens[:live_rows] - decode_only = decode_rows == live_rows decode_table = self._decode_block_table[:decode_rows] if decode_only: prepare_c4_decode_metadata( @@ -707,11 +716,7 @@ def forward( output = self.topk_indices_buffer[:rows] if live_rows < rows: output[live_rows:].fill_(-1) - if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: - if self._main_cache_num_blocks < 1: - raise RuntimeError("GLM selector main cache is not bound") - if self._physical_selection_plan is None: - raise RuntimeError("GLM physical selection has not been declared") + if emit_physical_selection: self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From a038b41614e5e0d37d2619a9a9408fd52c04ac31 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 20:40:28 +0000 Subject: [PATCH 023/365] fix(ds4.1): release transient MHC and WO graph activations Retain caller-owned scratch without extending activation lifetimes through B12X bindings. Keep prepared plans, fresh bindings, precision and CUDA graph coverage unchanged. Native GPU tests reproduce retained operands in both MHC modes and both WO modes. All four corrected cases and the block-linear regression pass exact changed-input graph replay. TP4 DSpark K7 serving restores logical KV capacity from 1,920,034 to 4,466,215 tokens and passes bounded text, prefix and vision checks. C1 is -0.65%, C8 -2.32% and 32k prefill +0.44% in sequential samples; regression-free throughput is not claimed. Ruff findings and eleven source mypy findings reproduce on the untouched base. Tests pass mypy. Skip those pre-existing hook failures and whole-file formatting for this isolated local commit; other commit hooks remain enabled. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit a28be967e5779629dbc212a312a3ad8924c0623f) --- tests/v1/attention/test_b12x_v41_workspace.py | 241 ++++++++++++++++++ vllm/models/deepseek_v4_1/attention.py | 7 +- vllm/models/deepseek_v4_1/b12x_layers.py | 4 +- 3 files changed, 249 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index e6bd78e6719d..1a87261c3e7a 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -422,6 +422,144 @@ def run(rows): session.close() +@pytest.mark.parametrize("operation", ["pre", "post_pre"]) +def test_mhc_capture_retains_scratch_not_caller_activations( + native_workspace, monkeypatch, operation +): + """Intermediate MHC outputs must be reusable across shared-pool graphs.""" + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4_1 import b12x_layers + from vllm.utils.b12x import B12xWorkload, register_b12x_layer + + _, manager, workspace = native_workspace + monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (8, 16)) + device = torch.device("cuda") + hidden = 5120 + torch.manual_seed(146) + layer = torch.nn.Module() + for kind in ("attn", "ffn"): + setattr(layer, f"hc_{kind}_fn", torch.randn(24, hidden * 4, device=device) / 64) + setattr(layer, f"hc_{kind}_scale", torch.ones(3, device=device)) + setattr(layer, f"hc_{kind}_base", torch.zeros(24, device=device)) + setattr( + layer, + f"{kind}_norm", + SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device) + ), + ) + layer.hc_attn_fn_broadcast = None + module = layer._b12x_mhc = b12x_layers.B12xMHC( + SimpleNamespace( + hidden_size=hidden, + rms_norm_eps=1e-20, + hc_eps=1e-6, + hc_sinkhorn_iters=20, + hc_mult=4, + ) + ) + name = f"test.mhc.capture.{operation}" + module.bind_layer_name(name) + register_b12x_layer(name, layer) + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=16, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare( + tuple( + request + for unit in module.get_b12x_preparation_units(layer, workload) + for request in unit.requests + ) + ) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in module._plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + residual = torch.randn(16, 4, hidden, dtype=torch.bfloat16, device=device) + previous = torch.randn(16, hidden, dtype=torch.bfloat16, device=device) + pre = torch.full((16, 4), 0.25, device=device) + previous_post = pre.clone() + comb = torch.eye(4, device=device).expand(16, -1, -1).contiguous() + destinations = [ + torch.empty_like(residual), + torch.empty_like(pre), + torch.empty_like(comb), + torch.empty_like(previous), + torch.empty_like(pre), + ] + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + kwargs = ( + {} + if operation == "pre" + else dict( + previous_output=previous[:rows], + previous_post=previous_post[:rows], + previous_comb=comb[:rows], + ) + ) + values = module.pre( + residual[:rows], + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + layer.attn_norm.weight, + pre[:rows], + **kwargs, + ) + for destination, value in zip(destinations, values, strict=True): + references.append(weakref.ref(value)) + destination[:rows].copy_(value) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + residual.normal_() + previous.normal_() + run(rows) + expected = [output[:rows].clone() for output in destinations] + for output in destinations: + output.fill_(float("nan")) + graphs[rows].replay() + for output, reference in zip(destinations, expected, strict=True): + torch.testing.assert_close(output[:rows], reference, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + @pytest.mark.parametrize("draft_tokens", [5, 7]) def test_parallel_draft_reservation_keeps_decode_split_parallelism( native_workspace, draft_tokens @@ -824,6 +962,109 @@ def projection(n, k, value): del resources +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_output_projection_capture_releases_caller_activations( + native_workspace, is_prefill +): + """WO bindings may borrow inputs and results without pinning them per graph.""" + from b12x.preparation import PreparationSession + + from vllm.utils.b12x import B12xWorkload + + attention, manager, workspace = native_workspace + device = torch.device("cuda") + torch.manual_seed(147) + layer = _layer(attention) + layer.capacity = 16 + layer.n_local_groups = 2 + layer.n_local_heads = 16 + layer.head_dim = 512 + layer.rope_head_dim = 64 + layer.o_lora_rank = 1024 + layer.hidden_size = 5120 + layer._wo_plans = {} + angles = torch.randn(32, 32, device=device) + layer.rotary_emb = SimpleNamespace( + cos_sin_cache=torch.cat((angles.cos(), angles.sin()), dim=-1) + ) + for name, shape in (("wo_a", (2048, 4096)), ("wo_b", (5120, 2048))): + setattr( + layer, + name, + SimpleNamespace( + weight=(torch.randn(shape, device=device) / 32).to(torch.float8_e4m3fn), + weight_scale_inv=torch.ones( + shape[0] // 32, shape[1] // 32, device=device + ).to(torch.float8_e8m0fnu), + ), + ) + layer.setup_wo_projection() + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=32, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare(layer._wo_preparation_unit(workload).requests) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in layer._wo_plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + inputs = torch.randn(16, 16, 512, dtype=torch.bfloat16, device=device) + positions = torch.arange(16, device=device) + output = torch.empty(16, 5120, dtype=torch.bfloat16, device=device) + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + source = inputs[:rows] + 1 + result = layer._o_proj(source, positions[:rows], is_prefill=is_prefill) + references.extend((weakref.ref(source), weakref.ref(result))) + if result._base is not None: + references.append(weakref.ref(result._base)) + output[:rows].copy_(result) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + inputs.normal_() + positions.copy_(torch.randperm(16, device=device)) + run(rows) + expected = output[:rows].clone() + output.fill_(float("nan")) + graphs[rows].replay() + torch.testing.assert_close(output[:rows], expected, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + def test_output_projection_uses_fused_block32_path(native_workspace, monkeypatch): from dataclasses import dataclass diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 661e1f862bee..06c590ce5bc8 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1706,9 +1706,10 @@ def _o_proj(self, o, positions, *, is_prefill=False): ) require_prepared(plan, "gemm.wo_projection", o.device) weights = self._wo_projection_weights + scratch = _scratch(plan) binding = wo_projection.bind_inv_rope( plan, - scratch=_scratch(plan), + scratch=scratch, o=o, positions=positions, cos_sin_cache=self.rotary_emb.cos_sin_cache, @@ -1725,7 +1726,9 @@ def _o_proj(self, o, positions, *, is_prefill=False): device=o.device, ) binding = replace(binding, output=output) - retain_cuda_graph_capture_resource(binding) + # The model owns the plan and weights; graph pools manage activation + # lifetimes. Keep only the borrowed workspace, not layer activations. + retain_cuda_graph_capture_resource(scratch) local = wo_projection.run_inv_rope( binding=binding, plan=plan, diff --git a/vllm/models/deepseek_v4_1/b12x_layers.py b/vllm/models/deepseek_v4_1/b12x_layers.py index de605dfd11c0..57292c9db2e1 100644 --- a/vllm/models/deepseek_v4_1/b12x_layers.py +++ b/vllm/models/deepseek_v4_1/b12x_layers.py @@ -586,7 +586,9 @@ def _mhc_pre( comb=comb, pre_out=pre_out, ) - retain_cuda_graph_capture_resource(binding) + # Caller-owned outputs follow PyTorch's graph-pool lifetimes. Retaining + # them through the binding prevents reuse across layers and graph shapes. + retain_cuda_graph_capture_resource(scratch) if previous_output is None: mhc.run_pre( residual, From 51b4a522be82835a2a341ab348b24f811e62c90d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 21:16:53 +0000 Subject: [PATCH 024/365] fix(ds4.1): release captured query and context activations Retain caller-owned scratch and attention metadata without retaining transient query/index operands or DSpark context-projection results. Prepared plans, per-call bindings, numerical kernels, cache geometry and graph coverage remain unchanged. Real SM120 tests reproduce the immutable-image ownership defects, including detached query storage. Twenty-four native lifetime and exact changed-input replay cases pass across decode/prefill and three page geometries. TP4 DSpark K7 memory snapshots remove 858657792 additional live graph bytes per rank and pass text, prefix and vision checks. Independent uninstrumented throughput qualification is recorded separately. The test file passes mypy; existing source typing findings remain. Ruff diagnostics match the parent revision. Skip the reproducing Ruff/mypy findings and unrelated whole-file formatting while retaining the remaining commit hooks. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 42627f07f551a5a38575eb1baca9693d7856c35e) --- tests/v1/attention/test_b12x_v41_workspace.py | 86 +++++++++++++------ vllm/models/deepseek_v4_1/attention.py | 22 ++--- vllm/models/deepseek_v4_1/nvidia/dspark.py | 2 +- 3 files changed, 72 insertions(+), 38 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 1a87261c3e7a..549ae1e88fec 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -4,10 +4,12 @@ import gc import weakref +from functools import partial from types import SimpleNamespace import pytest import torch +from torch.multiprocessing.reductions import StorageWeakRef from vllm.config import CUDAGraphMode @@ -336,9 +338,11 @@ def test_index_preparation_reuses_reserved_workspace(native_workspace): torch.testing.assert_close(cache, torch.zeros_like(cache), rtol=0, atol=0) +@pytest.mark.parametrize("projection_kind", ["linear", "dspark_context"]) def test_block_linear_capture_retains_scratch_not_caller_activations( native_workspace, monkeypatch, + projection_kind, ): from vllm.models.deepseek_v4_1 import b12x_layers @@ -353,10 +357,22 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( layer.weight = torch.randn(256, 256, device=device).to(torch.float8_e4m3fn) layer.weight_scale_inv = torch.ones(8, 8, device=device).to(torch.float8_e8m0fnu) monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (16,)) - method = b12x_layers.B12xFP8LinearMethod( - SimpleNamespace(weight_block_size=[32, 32]) - ) - method.process_weights_after_loading(layer) + if projection_kind == "linear": + method = b12x_layers.B12xFP8LinearMethod( + SimpleNamespace(weight_block_size=[32, 32]) + ) + method.process_weights_after_loading(layer) + project = partial(method.apply, layer) + plans = layer.b12x_plans + else: + from vllm.models.deepseek_v4_1.nvidia import dspark + + monkeypatch.setattr(dspark, "_execution_capacities", lambda: (16,)) + method = dspark._ContextKVProjection( + SimpleNamespace(fused_wqa_wkv=layer, q_lora_rank=0), 16 + ) + project = method + plans = method.plans workload = B12xWorkload( stage="weights", token_counts=(8, 16), @@ -376,7 +392,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( ) session.freeze() manager.reserve_all( - *((spec.shape, spec.dtype) for spec in layer.b12x_plans[0].scratch_specs()) + *((spec.shape, spec.dtype) for spec in plans[0].scratch_specs()) ) manager.lock() inputs = torch.randn(16, 256, dtype=torch.bfloat16, device=device) @@ -385,7 +401,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( def run(rows): source = inputs[:rows] + 1 - output = method.apply(layer, source) + output = project(source) references.extend((weakref.ref(source), weakref.ref(output))) outputs[:rows].copy_(output) @@ -695,6 +711,7 @@ def test_attention_shared_scratch_graph_replay( device = torch.device("cuda", torch.accelerator.current_device_index()) torch.manual_seed(142) layer = _layer(attention, swa_page=swa_page) + layer.rotary_emb = SimpleNamespace(cos_sin_cache=torch.empty(0, device=device)) layer.config.cache_config.block_size = main_page layer.max_model_len = 4096 if rows == 36: @@ -755,13 +772,12 @@ def metadata(page): device=device, dtype=torch.uint8, ) - _write_cache( - kv, - cache, - torch.arange(length, device=device) + page, - page_size=page, - cache_kind=kind, - cache_format="deepseek_v41", + from b12x.attention._shared.mla.compressed_reference import ( + pack_deepseek_v41_cache_reference, + ) + + cache[1:].copy_( + pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -782,12 +798,6 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 - _write_index_keys( - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - page_size=main_page, - ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -811,17 +821,35 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) + attention.dsa_indexer.quantize_write_index_k_mxfp4( + layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + ) + activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): - attention.dsa_indexer.quantize_q_mxfp4( - layer._index_plan("decode", min(rows, layer.DECODE_CHUNK)), - iq, - q_mxfp4=packed, - q_scales=scales, - ) - layer.forward_mqa( - q, None, positions, out, index_query=(packed, scales, weights) + query = q.clone() + index_query = packed.clone(), scales.clone(), weights.clone() + activation_refs.extend( + (name, StorageWeakRef(tensor.untyped_storage())) + for name, tensor in zip( + ("query", "index_query", "index_scales", "index_weights"), + (query, *index_query), + ) ) + mode = "decode" if is_decode else "prefill" + chunk = layer.DECODE_CHUNK if is_decode else layer.INDEX_CHUNK + for offset in range(0, rows, chunk): + end = min(offset + chunk, rows) + attention.dsa_indexer.quantize_q_mxfp4( + layer._index_plan(mode, end - offset), + iq[offset:end], + q_mxfp4=index_query[0][offset:end], + q_scales=index_query[1][offset:end], + ) + layer.forward_mqa(query, None, positions, out, index_query=index_query) run() batched_output = out.clone() @@ -845,6 +873,10 @@ def run(): torch.cuda.graph(graph, stream=stream), ): run() + # Graph resources own scratch, not the per-layer query activations. + gc.collect() + retained = [name for name, reference in activation_refs if not reference.expired()] + assert not retained, f"Capture resources retained caller activations: {retained}" # Changed queries exercise both selection and attention on replay, after # all plan storage has been reused by another operation. for step, seed in enumerate((143, 144)): diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 06c590ce5bc8..4dd404ee4853 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1454,16 +1454,17 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): candidate_lengths=source._candidate_lens[offset:end], ) index_plan = self._index_plan(index_mode, count) + scratch = current_workspace_manager().get_simultaneous( + *( + (spec.shape, spec.dtype) + for spec in dsa_indexer.scratch_specs( + index_plan, device=iq_data.device + ) + ) + ) binding = dsa_indexer.bind( index_plan, - scratch=current_workspace_manager().get_simultaneous( - *( - (spec.shape, spec.dtype) - for spec in dsa_indexer.scratch_specs( - index_plan, device=iq_data.device - ) - ) - ), + scratch=scratch, q_mxfp4=iq_data[offset:end], q_scales=iq_scale[offset:end], query_weights=iw[offset:end], @@ -1475,7 +1476,7 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): output_indices=self.topk_indices_buffer[offset:end], **candidate_args, ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) dsa_indexer.score(binding) dsa_indexer.select(binding) @@ -1534,7 +1535,8 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): swa_lengths=swa_lengths[:rows], **kwargs, ) - retain_cuda_graph_capture_resource(binding) + # Keep scratch and metadata, not aliases of the transient query storage. + retain_cuda_graph_capture_resource(buffers) mla.run( binding=binding, swa_k_cache=self.swa_cache_layer.kv_cache, diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v4_1/nvidia/dspark.py index 5251cc82d9e4..61ba555c05d3 100644 --- a/vllm/models/deepseek_v4_1/nvidia/dspark.py +++ b/vllm/models/deepseek_v4_1/nvidia/dspark.py @@ -186,7 +186,7 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor: packed_weight=self.weight, output=out.view(rows, self.weight.out_features, 1), ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) block_fp8_linear.run(binding=binding) return out From 0a28123d271a9fc9ae35128b5d97878f2fae0f07 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 12:25:32 +0000 Subject: [PATCH 025/365] test(ds41): verify MHC ownership across breakable graph segments Exercise pre and lagged post-pre outputs across an eager segment boundary and changed-input replay on a shared graph pool. Require transient Python owners to release while caller scratch remains retained. Native graph semantics and model code are unchanged. Applicable hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 3ead61584ac374ca5825ffb8aa21ad8a1a05beca) --- tests/v1/attention/test_b12x_v41_workspace.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 549ae1e88fec..260f386aa122 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -439,12 +439,14 @@ def run(rows): @pytest.mark.parametrize("operation", ["pre", "post_pre"]) +@pytest.mark.parametrize("capture_kind", ["full", "breakable"]) def test_mhc_capture_retains_scratch_not_caller_activations( - native_workspace, monkeypatch, operation + native_workspace, monkeypatch, operation, capture_kind ): """Intermediate MHC outputs must be reusable across shared-pool graphs.""" from b12x.preparation import PreparationSession + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture from vllm.models.deepseek_v4_1 import b12x_layers from vllm.utils.b12x import B12xWorkload, register_b12x_layer @@ -537,6 +539,10 @@ def run(rows): pre[:rows], **kwargs, ) + capture = BreakableCUDAGraphCapture.current() + if capture is not None: + # The lagged outputs cross a segment boundary before consumption. + capture.add_eager(lambda: None) for destination, value in zip(destinations, values, strict=True): references.append(weakref.ref(value)) destination[:rows].copy_(value) @@ -547,14 +553,23 @@ def run(rows): for rows in (16, 8): run(rows) torch.accelerator.synchronize() - graph = torch.cuda.CUDAGraph() + graph = ( + torch.cuda.CUDAGraph() + if capture_kind == "full" + else BreakableCUDAGraphCapture(pool=pool) + ) graphs[rows] = graph + context = ( + torch.cuda.graph(graph, pool=pool) if capture_kind == "full" else graph + ) with ( session.capture(), workspace.collect_cuda_graph_capture_resources() as resources, - torch.cuda.graph(graph, pool=pool), + torch.cuda.stream(torch.cuda.Stream()), + context, ): run(rows) + torch.accelerator.synchronize() owners.append(resources) gc.collect() assert all(reference() is None for reference in references) From 72085567f1cfc1d0b9e914d7f11d09c51e385c87 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:17:17 +0000 Subject: [PATCH 026/365] fix(ds4): reuse MHC autotuning output buffers Allocate one output set per preparation candidate so repeated CUDA graph timing samples do not accumulate functional-output copies beyond the primed memory budget. Serving kernels, operation queries, candidate search and caller-owned scratch are unchanged. Validation: 13 native SM120 tests pass, including exact functional parity and changed-input CUDA graph replay at 3 and 4096 tokens. The baseline fails output-pointer reuse. DS4 TP2 completes weights preparation; end-to-end serving remains gated by KV admission capacity. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 2b2fa507b89d9652c5d87db61c44ee89b5fe864f) --- .../kernels/test_b12x_linear.py | 112 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 22 +++- 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 72c4f6d130ff..6a1a540400e3 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -7,6 +7,7 @@ import sys import types from dataclasses import dataclass +from typing import Any import pytest import torch @@ -1874,6 +1875,117 @@ def test_b12x_mhc_declares_the_operations_each_layer_runs( mhc._plan_for("pre", 8) +@torch.inference_mode() +@pytest.mark.parametrize("operation", ["pre", "post_pre", "post_pre_bf16", "post"]) +@pytest.mark.parametrize("tokens", [3, 4096]) +def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( + monkeypatch, operation, tokens +): + """Tuning fixtures retain one output set per candidate, not per sample.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("native b12x mHC requires SM12x") + from b12x.norm.mhc import _impl + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x + + device = torch.device("cuda", torch.accelerator.current_device_index()) + hidden, mult = 4096, 4 + fn = torch.randn(24, mult * hidden, device=device) * 0.001 + norm = types.SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device), + variance_epsilon=1e-6, + ) + layer = types.SimpleNamespace( + hc_attn_fn=fn, + hc_ffn_fn=fn, + hc_attn_scale=torch.ones(3, device=device), + hc_ffn_scale=torch.ones(3, device=device), + hc_attn_base=torch.zeros(24, device=device), + hc_ffn_base=torch.zeros(24, device=device), + hc_attn_fn_broadcast=fn.view(24, mult, hidden).sum(1), + attn_norm=norm, + ffn_norm=norm, + hc_ffn_fn_bf16=fn.to(torch.bfloat16), + ) + mhc = dsv4_b12x.B12xMHCResidual( + hidden_size=hidden, + hc_mult=mult, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=20, + ) + workload = B12xWorkload( + stage="weights", + token_counts=(tokens,), + fixed_token_counts=(), + output_dtype=torch.bfloat16, + max_tokens=tokens, + max_seqs=1, + max_model_len=tokens, + ) + (unit,) = mhc.get_b12x_preparation_units(layer, workload) + request = next( + item for item in unit.requests if item.name.endswith(f".{operation}.m{tokens}") + ) + native_operation = "post_pre" if operation == "post_pre_bf16" else operation + name = f"_b12x_mhc_{native_operation}_impl" + native = getattr(_impl, name) + captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + + def observe(*args, **kwargs): + captured_call[:] = [(args, kwargs)] + return native(*args, **kwargs) + + def outputs(value): + return value if isinstance(value, tuple) else (value,) + + with PreparationSession( + device=device, autotune=False, compile_workers=0 + ) as session: + session.prepare((request,)) + call = request.prepare_call(request.plan.prepared.state) + with monkeypatch.context() as hooks: + hooks.setattr(_impl, name, observe) + call.produce() + first = outputs(call.invoke()) + args, kwargs = captured_call.pop() + reference_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} + } + pointers = [tensor.data_ptr() for tensor in first] + second = outputs(call.invoke()) + assert [tensor.data_ptr() for tensor in second] == pointers + for actual, expected in zip( + second, outputs(native(*args, **reference_kwargs)), strict=True + ): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + try: + with session.capture(), torch.cuda.graph(graph): + call.invoke() + call.invoke() + for _ in range(2): + call.produce() + for tensor in first: + tensor.fill_(float("nan")) + allocated = torch.accelerator.memory_allocated(device) + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated(device) == allocated + assert [ + tensor.data_ptr() for tensor in outputs(call.output) + ] == pointers + expected = outputs(native(*args, **reference_kwargs)) + for actual, wanted in zip(first, expected, strict=True): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + finally: + graph.reset() + + def test_b12x_mhc_keeps_plans_across_workloads(monkeypatch) -> None: from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 98c210d79767..232d84b2dc94 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -435,6 +435,24 @@ def prepare(state): dtype=torch.float32, device=device, ) + # Each candidate owns one output set. Functional allocations inside + # repeated capture samples would reserve additional graph-pool copies + # that are absent from the tuner's primed-residency measurement. + residual_out = torch.empty( + (tokens, self.hc_mult, self.hidden_size), + dtype=torch.bfloat16, + device=device, + ) + outputs = ( + {} + if operation == "post" + else { + "residual_out": residual_out, + "y_out": torch.empty_like(x), + "post_out": torch.empty_like(post), + "comb_out": torch.empty_like(comb), + } + ) def produce(): residual.normal_() @@ -455,6 +473,7 @@ def produce(): norm_weight=attn_norm.weight, norm_eps=float(attn_norm.variance_epsilon), _state=state, + **outputs, ) elif operation in ("post_pre", "post_pre_bf16"): if operation == "post_pre": @@ -488,10 +507,11 @@ def produce(): norm_weight=norm.weight, norm_eps=float(norm.variance_epsilon), _state=state, + **outputs, ) else: run = lambda: _impl._b12x_mhc_post_impl( - x, residual, post, comb, _state=state + x, residual, post, comb, out=residual_out, _state=state ) def execute(): From 7e223ba84d014aeb968c593b13468830270bc76b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:54:03 +0000 Subject: [PATCH 027/365] fix(workspace): reserve B12X scratch by model owner Keep target and DSpark draft scratch capacities independent during profile hooks and B12X preparation. Reserve shared modules in both lanes and propagate each lane capacity across microbatches before graph capture. Preserve the complete attention shape envelope and unscoped coordinator compatibility. Validation: 83 focused workspace, preparation and graph-profile tests pass. Stock DS4 TP2/DSpark K5 with a 4096-token budget admits 7.56 GiB KV instead of 6.12 GiB, retaining the 1M-token model limit. Serving throughput qualification is pending. The unchanged base reproduces eight coordinator mypy errors and the existing pickle-import lint finding; remaining hooks pass. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit ea358a87e72e6f44baf14c436d278b5f4d52b92c) --- tests/model_executor/test_b12x_warmup.py | 25 +++++ tests/v1/worker/test_workspace.py | 117 +++++++++++++++++++++ vllm/model_executor/warmup/b12x_prepare.py | 48 +++++++-- vllm/utils/b12x.py | 7 +- vllm/v1/worker/b12x_startup.py | 20 +++- vllm/v1/worker/gpu/model_runner.py | 29 +++-- vllm/v1/worker/workspace.py | 40 +++++-- 7 files changed, 259 insertions(+), 27 deletions(-) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index c54ce7c3856d..41690faa4fef 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -348,6 +348,31 @@ def test_scope_unit_calls_wraps_mapping_factories() -> None: request.prepare_call = {2: lambda state: _Call(run=lambda: _workspace_lane.get())} scoped = scope_b12x_unit_calls(_unit("u", request), 1) assert scoped.requests[0].prepare_call[2](object()).run() == 1 + assert scoped.workspace_lanes == (1,) + + +@pytest.mark.parametrize("draft_lane", [0, 1]) +def test_shared_preparation_scratch_covers_both_model_owners(draft_lane): + target, draft, shared = (torch.nn.Module() for _ in range(3)) + for module in (target, draft, shared): + module.b12x_preparation_provider = _Provider() + target.shared = shared + draft.shared = shared + units = b12x_prepare.collect_b12x_units( + _worker(target, draft=draft, draft_lane=draft_lane), _workload() + ) + ownership = b12x_prepare._request_workspace_lanes(units) + assert ownership[f"{id(target):x}"] == (0,) + assert ownership[f"{id(draft):x}"] == (draft_lane,) + assert ownership[f"{id(shared):x}"] == tuple(sorted({0, draft_lane})) + assert ( + sum( + f"{id(shared):x}" == request.name + for unit in units + for request in unit.requests + ) + == 1 + ) def test_blockscaled_holder_declares_provided_workspace_under_the_cap( diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 0eba7b71d592..1699fe90ea7f 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from typing import cast import pytest @@ -167,6 +168,122 @@ def test_workspace_lane_validation(monkeypatch) -> None: workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) +def test_profile_reserves_model_lanes_without_replicating_target_capacity( + monkeypatch, +) -> None: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + monkeypatch.setattr(workspace, "_manager", manager) + calls: list[tuple[int, int]] = [] + + class ScratchOwner(torch.nn.Module): + def __init__(self, size): + super().__init__() + self.size = size + + def reserve_profile_scratch(self): + calls.append((self.size, workspace._workspace_lane.get())) + manager.get_simultaneous(((self.size,), torch.uint8)) + + target = ScratchOwner(4096) + draft = ScratchOwner(512) + shared = ScratchOwner(256) + target.add_module("shared", shared) + draft.add_module("shared", shared) + runner = SimpleNamespace( + get_model=lambda: target, + get_draft_model=lambda: draft, + _draft_workspace_lane=1, + compilation_config=SimpleNamespace( + static_forward_context={ + "target": target, + "alias": target, + "draft": draft, + "shared": shared, + } + ), + ) + GPUModelRunner._reserve_profile_scratch(runner) + + assert calls == [(4096, 0), (512, 1), (256, 0), (256, 1)] + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + assert [buffer.numel() for buffer in buffers] == [4096, 512, 4096, 512] + pointers = [buffer.data_ptr() for buffer in buffers] + assert len(set(pointers)) == 4 + manager.lock() + manager.reserve_by_lane() + for ubatch in range(2): + active_ubatch[0] = ubatch + for lane, size in enumerate((4096, 512)): + with workspace.use_workspace_lane(lane): + (view,) = manager.get_simultaneous(((size,), torch.uint8)) + assert view.data_ptr() == pointers[ubatch * 2 + lane] + + +def test_lane_reservation_rejects_unreserved_microbatch_when_locked(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((256,), torch.uint8)) + manager.lock() + with pytest.raises(AssertionError, match="reserve_by_lane"): + manager.reserve_by_lane() + + +@pytest.mark.parametrize("scoped", [False, True]) +def test_preparation_reservation_preserves_per_lane_profile_capacity( + monkeypatch, scoped +): + from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator + + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((4096,), torch.uint8)) + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((256,), torch.uint8)) + requests = [] + for name, size in (("target", 1024), ("draft", 512), ("shared", 768)): + spec = SimpleNamespace(shape=(size,), dtype=torch.uint8) + requests.append( + SimpleNamespace( + name=name, + plan=SimpleNamespace(scratch_specs=lambda spec=spec: (spec,)), + ) + ) + job = SimpleNamespace( + advance=lambda **kwargs: SimpleNamespace(done=True, pending_compilation=False), + result=lambda: SimpleNamespace(close=lambda: None), + close=lambda: None, + ) + coordinator = B12xPreparationCoordinator( + SimpleNamespace(begin=lambda *args, **kwargs: job), + [(tuple(requests), False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=manager, + request_workspace_lanes=( + {"target": (0,), "draft": (1,), "shared": (0, 1)} if scoped else None + ), + ) + result = coordinator.advance() + assert result["done"] and result["error"] is None + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + expected = [4096, 768, 4096, 768] if scoped else [4096] * 4 + assert [buffer.numel() for buffer in buffers] == expected + assert len({buffer.data_ptr() for buffer in buffers}) == 4 + manager.lock() + manager.reserve_by_lane() + + def test_cuda_graph_capture_resources_are_scoped_to_collector() -> None: outside = object() first = object() diff --git a/vllm/model_executor/warmup/b12x_prepare.py b/vllm/model_executor/warmup/b12x_prepare.py index dafff99c8e49..6ec0d9768778 100644 --- a/vllm/model_executor/warmup/b12x_prepare.py +++ b/vllm/model_executor/warmup/b12x_prepare.py @@ -228,6 +228,7 @@ def _units_from_modules( workload: B12xWorkload, *, seen: set[int] | None = None, + module_lanes: dict[int, tuple[int, ...]] | None = None, ) -> Iterable[B12xPreparationUnit]: if seen is None: seen = set() @@ -254,7 +255,12 @@ def _units_from_modules( f"{type(provider).__qualname__} returned a non-unit " "preparation value" ) - yield unit + lanes = ( + (workload.lane,) + if module_lanes is None + else module_lanes.get(id(module), (workload.lane,)) + ) + yield replace(unit, workspace_lanes=lanes) def mark_b12x_eager_shapes(worker: Worker) -> None: @@ -328,26 +334,43 @@ def collect_b12x_units( """Collect every unit of one stage from the model, the draft, and comms.""" mark_b12x_eager_shapes(worker) units: list[B12xPreparationUnit] = [] + model = worker.get_model() draft = worker.get_draft_model() + lane = _draft_lane(worker) if draft is not None else 0 + lanes_by_module: dict[int, set[int]] = {} + for owner, owner_lane in ((model, 0), (draft, lane)): + if owner is not None: + for module in owner.modules(): + lanes_by_module.setdefault(id(module), set()).add(owner_lane) + module_lanes = {key: tuple(sorted(lanes)) for key, lanes in lanes_by_module.items()} if draft is not None: - lane = _draft_lane(worker) workload = _draft_workload(worker, workload, lane=0) # Target and draft can share the same embedding and output-head modules. seen: set[int] = set() - units.extend(_units_from_modules(worker.get_model(), workload, seen=seen)) + units.extend( + _units_from_modules(model, workload, seen=seen, module_lanes=module_lanes) + ) if draft is not None: draft_workload = replace(workload, lane=lane) from vllm.v1.worker.workspace import use_workspace_lane with use_workspace_lane(lane): - draft_units = list(_units_from_modules(draft, draft_workload, seen=seen)) + draft_units = list( + _units_from_modules( + draft, draft_workload, seen=seen, module_lanes=module_lanes + ) + ) if lane: draft_units = [scope_b12x_unit_calls(unit, lane) for unit in draft_units] units.extend(draft_units) for provider in b12x_unit_providers(): hook = getattr(provider, "get_b12x_preparation_units", None) if callable(hook): - units.extend(hook(provider, workload)) + # Communicators serve target and draft executions independently. + units.extend( + replace(unit, workspace_lanes=tuple(sorted({0, lane}))) + for unit in hook(provider, workload) + ) # The state stage also carries the weights-stage units: their prepared # plans are skipped by the session, and any plan a layer declared for a # count that only the state stage exposes is prepared before capture. @@ -447,6 +470,14 @@ def b12x_batches(units: Iterable[B12xPreparationUnit], *, autotune: bool = True) return batches +def _request_workspace_lanes(units: Iterable[B12xPreparationUnit]): + return { + request.name: unit.workspace_lanes + for unit in units + for request in unit.requests + } + + def get_b12x_session(worker: Worker): """Return the worker's preparation session, creating it on first use.""" session = getattr(worker, "_b12x_session", None) @@ -492,10 +523,12 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): from vllm.v1.worker.workspace import current_workspace_manager batches = [] + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage="state" if stage == "bind" else stage) + units = collect_b12x_units(worker, workload) batches = b12x_batches( - collect_b12x_units(worker, workload), + units, autotune=( bool(worker.vllm_config.kernel_config.enable_b12x_autotune) and os.environ.get("B12X_AUTOTUNE", "1") != "0" @@ -531,6 +564,7 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if batches else None, + request_workspace_lanes=_request_workspace_lanes(units), ) @@ -608,6 +642,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: from vllm.v1.worker.workspace import current_workspace_manager requests: tuple[PreparationRequest, ...] = () + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage=stage) units = collect_b12x_units(worker, workload) @@ -621,6 +656,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if requests else None, + request_workspace_lanes=_request_workspace_lanes(units), ) outcome = coordinator.status() while not outcome["done"]: diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 11c16cfe4227..25e6ea3d2618 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -94,6 +94,7 @@ class B12xPreparationUnit: requests: tuple[PreparationRequest, ...] stage: Literal["weights", "state"] autotune: bool = True + workspace_lanes: tuple[int, ...] = (0,) def __post_init__(self): if not isinstance(self.name, str) or not self.name: @@ -103,6 +104,10 @@ def __post_init__(self): raise ValueError("invalid native preparation stage") if type(self.autotune) is not bool: raise TypeError("autotune must be boolean") + lanes = tuple(self.workspace_lanes) + if not lanes or any(type(lane) is not int or lane < 0 for lane in lanes): + raise ValueError("workspace lanes must be nonnegative integers") + object.__setattr__(self, "workspace_lanes", tuple(sorted(set(lanes)))) requests = tuple(self.requests) names = [request.name for request in requests] if len(names) != len(set(names)): @@ -210,7 +215,7 @@ def wrap_calls(calls): ) for request in unit.requests ) - return replace(unit, requests=requests) + return replace(unit, requests=requests, workspace_lanes=(lane,)) def b12x_preparation_token_counts( diff --git a/vllm/v1/worker/b12x_startup.py b/vllm/v1/worker/b12x_startup.py index 275bb0bf79fd..49fd5f51b5d7 100644 --- a/vllm/v1/worker/b12x_startup.py +++ b/vllm/v1/worker/b12x_startup.py @@ -38,11 +38,13 @@ def __init__( world_group, process_local_only: bool = False, workspace=None, + request_workspace_lanes: dict[str, tuple[int, ...]] | None = None, ) -> None: if type(global_rank) is not int or global_rank < 0: raise ValueError("global_rank must be a nonnegative integer") self.session = session self._workspace = workspace + self._request_workspace_lanes = request_workspace_lanes self._active_requests = () self.global_rank = global_rank self.world_group = world_group @@ -275,13 +277,25 @@ def _advance_local(self) -> None: job.result().close() if self._workspace is not None: + from vllm.v1.worker.workspace import use_workspace_lane + for request in self._active_requests: specs = tuple(request.plan.scratch_specs()) if specs: - self._workspace.get_simultaneous( - *((spec.shape, spec.dtype) for spec in specs) + lanes = ( + (0,) + if self._request_workspace_lanes is None + else self._request_workspace_lanes[request.name] ) - self._workspace.reserve_all() + for lane in lanes: + with use_workspace_lane(lane): + self._workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in specs) + ) + if self._request_workspace_lanes is None: + self._workspace.reserve_all() + else: + self._workspace.reserve_by_lane() self._active_requests = () self._job = None if self._batches: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43da8b0e4646..aba519befaff 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1004,15 +1004,26 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: self.pooling_runner.dummy_pooler_run(hidden_states) def _reserve_profile_scratch(self) -> None: - seen: set[int] = set() + model_lanes: dict[int, set[int]] = {} + for model, lane in ( + (self.get_model(), 0), + (self.get_draft_model(), self._draft_workspace_lane), + ): + if model is not None: + for module in model.modules(): + model_lanes.setdefault(id(module), set()).add(lane) + seen: set[tuple[int, int]] = set() for module in self.compilation_config.static_forward_context.values(): - if id(module) in seen: - continue - seen.add(id(module)) - reserve = getattr(module, "reserve_profile_scratch", None) - if reserve is not None: - reserve() - current_workspace_manager().reserve_all() + for lane in sorted(model_lanes.get(id(module), {0})): + key = (id(module), lane) + if key in seen: + continue + seen.add(key) + reserve = getattr(module, "reserve_profile_scratch", None) + if reserve is not None: + with use_workspace_lane(lane): + reserve() + current_workspace_manager().reserve_by_lane() @torch.inference_mode() def profile_run( @@ -1054,7 +1065,7 @@ def profile_run( torch.accelerator.synchronize() del hidden_states, sample_hidden_states self._profile_deepseek_v4_attention(prepare_profile_state) - current_workspace_manager().reserve_all() + current_workspace_manager().reserve_by_lane() self.reset_encoder_cache() gc.collect() diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 81adeee47065..ddde9f8b1df5 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -229,15 +229,40 @@ def reserve_all( required_bytes, max(map(self._workspace_size_bytes, self._current_workspaces), default=0), ) + self._reserve_slots( + [required_bytes] * len(self._current_workspaces), operation="reserve_all" + ) + + def reserve_by_lane(self) -> None: + """Cover every microbatch without copying capacity between model lanes. + + A target and its drafter have separate scratch contracts. A microbatch + can execute either model, but each model always uses its assigned lane. + """ + lane_bytes = [ + max( + map( + self._workspace_size_bytes, + self._current_workspaces[lane :: self._num_lanes], + ), + default=0, + ) + for lane in range(self._num_lanes) + ] + self._reserve_slots( + lane_bytes * self._num_ubatches, operation="reserve_by_lane" + ) + + def _reserve_slots(self, sizes: list[int], *, operation: str) -> None: undersized = [ workspace_id for workspace_id, workspace in enumerate(self._current_workspaces) - if self._workspace_size_bytes(workspace) < required_bytes + if self._workspace_size_bytes(workspace) < sizes[workspace_id] ] if self._locked and undersized: raise AssertionError( - "Workspace is locked but reserve_all requires " - f"{required_bytes / _MB:.2f} MB in slot(s) {undersized}." + f"Workspace is locked but {operation} requires " + f"larger capacity in slot(s) {undersized}." ) for workspace_id in undersized: @@ -246,14 +271,13 @@ def reserve_all( del current_workspace torch.accelerator.empty_cache() self._current_workspaces[workspace_id] = torch.empty( - (required_bytes,), dtype=torch.uint8, device=self._device + (sizes[workspace_id],), dtype=torch.uint8, device=self._device ) if envs.VLLM_DEBUG_WORKSPACE and undersized: logger.info( - "[WORKSPACE DEBUG] Reserved %.2f MB in execution slots %s", - required_bytes / _MB, - undersized, + "[WORKSPACE DEBUG] Reserved execution slots (slot, MB): %s", + [(slot, sizes[slot] / _MB) for slot in undersized], ) def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: @@ -327,7 +351,7 @@ def get_caller_info() -> str: # stream-ordered, so wait for the device first. Growth happens only # before the workspace is locked, never in steady-state serving. if self._device.type == "cuda" and current_workspace is not None: - torch.cuda.synchronize(self._device) + torch.accelerator.synchronize(self._device) self._current_workspaces[workspace_id] = None del current_workspace # Release the freed segment back to CUDA so the caching From 938f7e4778593cbd5e40e907b3609d0658a0c95d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 03:29:17 +0000 Subject: [PATCH 028/365] fix(gdn): bind caller-owned prefill activations directly Keep only sequence metadata on each GDN layer. Bind projected inputs and invocation-owned outputs through the native variable-row contract instead of retaining scheduler-capacity copies per layer. Priming uses temporary activation storage; recurrent-state ownership and kernel arithmetic are unchanged. Validation coverage includes token-capacity-independent metadata storage and CUDA graph replay with 49 or 64 rows, changing worklists, immutable inputs, and allocation counters. GPU qualification is recorded separately. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 983c38e2ef0b8f3acaaa029f952ba1c2f26f8033) --- tests/kernels/mamba/test_gdn_prefill_b12x.py | 33 ++++++++- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 72 ++++++++++--------- .../layers/mamba/ops/b12x_gdn_prefill.py | 58 +++++---------- 3 files changed, 87 insertions(+), 76 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_prefill_b12x.py b/tests/kernels/mamba/test_gdn_prefill_b12x.py index 50d66d32764f..153722421c08 100644 --- a/tests/kernels/mamba/test_gdn_prefill_b12x.py +++ b/tests/kernels/mamba/test_gdn_prefill_b12x.py @@ -10,6 +10,7 @@ from vllm.model_executor.layers.mamba.ops.b12x_gdn_prefill import ( B12xGdnPrefill, + GdnPrefillStaging, prefill_capacities, ) @@ -27,6 +28,25 @@ def test_prefill_capacity_family_covers_exact_scheduler_limit(capacity, expected assert prefill_capacities(capacity) == expected +@pytest.mark.parametrize("max_tokens", (16, 6019, 32768)) +def test_prefill_staging_memory_depends_on_sequences_not_token_capacity(max_tokens): + staging = GdnPrefillStaging.allocate( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + assert staging.nbytes == ((16 + 1) + 4 * 16 + 2) * 4 + assert staging.is_compatible( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_prepared_gdn_decode_plan_is_scoped_to_the_bound_recurrent_pool( default_vllm_config, @@ -324,8 +344,10 @@ def test_prefill_pooled_state_graph_replays_changed_lengths_and_slots(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", (49, 64)) def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists( default_vllm_config, + rows, ): from b12x._lib.runtime_control import kernel_resolution_guard from b12x.preparation import PreparationSession @@ -353,7 +375,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pytest.skip("SM12x is required") torch.manual_seed(23) device = torch.device("cuda") - rows, seqs, key_heads, value_heads = 64, 3, 2, 6 + seqs, key_heads, value_heads = 3, 2, 6 width = (2 * key_heads + value_heads) * 128 layer = QwenGatedDeltaNetAttention.__new__(QwenGatedDeltaNetAttention) torch.nn.Module.__init__(layer) @@ -383,7 +405,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pool, ) layer._b12x_preparation_prefix = "test.mixed-gdn" - layer._b12x_prefill_max_tokens = rows + layer._b12x_prefill_max_tokens = 64 layer._b12x_prefill_max_seqs = seqs layer._b12x_prefill_staging = None layer._b12x_decode_plan = None @@ -446,6 +468,11 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists ), attn_metadata=attention, ) + immutable_inputs = { + name: value.clone() + for name, value in inputs.items() + if isinstance(value, torch.Tensor) and name != "core_attn_out" + } saved_conv, saved_pool = conv.clone(), pool.clone() def stage(lengths, computed, drafts, swap): @@ -593,6 +620,8 @@ def oracle(): ) assert_close("mixed state", pool, expected_pool, ratio=5e-3) torch.testing.assert_close(conv, expected_conv, rtol=0, atol=0) + for name, value in immutable_inputs.items(): + torch.testing.assert_close(inputs[name], value, rtol=0, atol=0) finally: guard.__exit__(None, None, None) finally: diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1169d9304d43..f75bc330347b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -959,16 +959,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): if api is None: raise RuntimeError("b12x GDN prefill was not initialized") recurrent_state = self.kv_cache[1] - staging = self._b12x_prefill_staging - resident_nbytes = 0 - if staging is not None and staging.is_compatible( - max_tokens=self._b12x_prefill_max_tokens, - max_seqs=self._b12x_prefill_max_seqs, - key_heads=self._b12x_local_key_heads, - value_heads=self._b12x_local_value_heads, - device=recurrent_state.device, - ): - resident_nbytes = staging.nbytes caps = api.Caps( device=current_platform.current_device(), max_tokens=capacity, @@ -979,8 +969,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): state_dtype=recurrent_state.dtype, checkpoint_export=True, null_state_index=0, - staging_key=(self._b12x_preparation_prefix, "gdn-prefill-staging"), - staging_resident_nbytes=resident_nbytes, ) return api.plan( caps, @@ -1148,7 +1136,6 @@ def _benchmark_b12x_gdn_prefill(self, state, capacity: int): def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): from b12x.preparation import PreparedCall - owners: tuple[torch.Tensor, ...] specs = tuple(state.layout.scratch_specs()) if len(specs) != 1: raise RuntimeError("b12x GDN prefill requires one scratch buffer") @@ -1159,11 +1146,30 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): # Trial and prepare factories own their scratch; the runtime path in # B12xGdnPrefill.run draws from the workspace manager instead. scratch = torch.empty(spec.shape, dtype=spec.dtype, device=device) + # Priming and candidate races use temporary inputs. Serving binds the + # caller's projection and output tensors directly, without retaining + # scheduler-capacity activation storage on every layer. + mixed_qkv = torch.empty( + ( + capacity, + (2 * self._b12x_local_key_heads + self._b12x_local_value_heads) * 128, + ), + dtype=torch.bfloat16, + device=device, + ) + a = torch.empty( + (capacity, self._b12x_local_value_heads), + dtype=torch.bfloat16, + device=device, + ) + b = torch.empty_like(a) + output = torch.empty( + (capacity, self._b12x_local_value_heads, 128), + dtype=torch.bfloat16, + device=device, + ) + owners: tuple[torch.Tensor, ...] if benchmark: - mixed_qkv = torch.empty_like(staging.mixed_qkv[:capacity]) - a = torch.empty_like(staging.a[:capacity]) - b = torch.empty_like(staging.b[:capacity]) - output = torch.empty_like(staging.output[:capacity]) cu_seqlens = torch.empty_like(staging.query_start_loc) indices = torch.full_like(staging.initial_indices, slot) final_indices = torch.full_like(staging.final_indices, slot) @@ -1187,12 +1193,6 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): num_tokens, ) else: - mixed_qkv = staging.mixed_qkv[:capacity] - a, b, output = ( - staging.a[:capacity], - staging.b[:capacity], - staging.output[:capacity], - ) cu_seqlens = staging.query_start_loc indices, final_indices = staging.initial_indices, staging.final_indices checkpoint_indices, offsets = ( @@ -2346,6 +2346,15 @@ def _forward_core( raise RuntimeError( "b12x GDN prefill cache and runtime metadata must be bound" ) + prefill_output = torch.empty( + ( + conv_output_prefill.shape[0], + self._b12x_local_value_heads, + self.head_v_dim, + ), + device=conv_output_prefill.device, + dtype=conv_output_prefill.dtype, + ) runner.run( mixed_qkv=conv_output_prefill, a=a_prefill, @@ -2355,12 +2364,10 @@ def _forward_core( has_initial_state=prefill_has_initial_state, live_counts=attn_metadata.b12x_prefill_live_counts, checkpoint=attn_metadata.prefill_checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) - core_attn_out_non_spec = runner.output[ - : conv_output_prefill.shape[0] - ].unsqueeze(0) + core_attn_out_non_spec = prefill_output.unsqueeze(0) else: initial_state = ssm_state[prefill_state_indices] initial_state[~prefill_has_initial_state, ...] = 0 @@ -2989,6 +2996,7 @@ def _forward_core_b12x_mixed( query_start_loc=metadata.query_start_loc, metadata=metadata.convolution_metadata(rows), ).transpose(0, 1) + prefill_output = torch.empty_like(core_attn_out) runner.run( mixed_qkv=convolved, a=a.index_select(0, non_spec_indices), @@ -2998,22 +3006,22 @@ def _forward_core_b12x_mixed( has_initial_state=metadata.has_initial_state, live_counts=metadata.live_counts, checkpoint=metadata.checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) self._rms_norm_gated_cuda( - runner.output[:rows], + prefill_output, output_gate.index_select(0, non_spec_indices), - runner.output[:rows], + prefill_output, ) core_attn_out.zero_() width = core_attn_out.shape[-2] * core_attn_out.shape[-1] _scatter_b12x_gdn_output[(rows, triton.cdiv(width, 256))]( - runner.output, + prefill_output, non_spec_indices, metadata.live_counts, core_attn_out, - SOURCE_STRIDE=runner.output.stride(0), + SOURCE_STRIDE=prefill_output.stride(0), OUTPUT_STRIDE=core_attn_out.stride(0), WIDTH=width, BLOCK=256, diff --git a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py index b13ad1904238..67fd3fd419bd 100644 --- a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py +++ b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py @@ -69,16 +69,16 @@ def prefill_capacities(max_tokens: int) -> tuple[int, ...]: @dataclass(frozen=True) class GdnPrefillStaging: - """Reusable capacity buffers independent of a recurrent-state generation.""" + """Sequence metadata independent of a recurrent-state generation. + + Activations and outputs remain caller-owned. Native GDN bindings accept + fewer rows than a plan's capacity, so no padded activation copies are needed. + """ max_tokens: int max_seqs: int key_heads: int value_heads: int - mixed_qkv: torch.Tensor - a: torch.Tensor - b: torch.Tensor - output: torch.Tensor query_start_loc: torch.Tensor initial_indices: torch.Tensor final_indices: torch.Tensor @@ -97,26 +97,12 @@ def allocate( value_heads: int, device: torch.device, ) -> "GdnPrefillStaging": - mixed_qkv = torch.empty( - (max_tokens, (2 * key_heads + value_heads) * 128), - dtype=torch.bfloat16, - device=device, - ) - a = torch.empty((max_tokens, value_heads), dtype=torch.bfloat16, device=device) initial_indices = torch.zeros(max_seqs, dtype=torch.int32, device=device) return cls( max_tokens=max_tokens, max_seqs=max_seqs, key_heads=key_heads, value_heads=value_heads, - mixed_qkv=mixed_qkv, - a=a, - b=torch.empty_like(a), - output=torch.empty( - (max_tokens, value_heads, 128), - dtype=torch.bfloat16, - device=device, - ), query_start_loc=torch.zeros(max_seqs + 1, dtype=torch.int32, device=device), initial_indices=initial_indices, final_indices=torch.zeros_like(initial_indices), @@ -140,7 +126,7 @@ def is_compatible( and self.max_seqs >= max_seqs and self.key_heads == key_heads and self.value_heads == value_heads - and self.mixed_qkv.device == device + and self.query_start_loc.device == device ) @property @@ -148,10 +134,6 @@ def nbytes(self) -> int: return sum( tensor.numel() * tensor.element_size() for tensor in ( - self.mixed_qkv, - self.a, - self.b, - self.output, self.query_start_loc, self.initial_indices, self.final_indices, @@ -164,10 +146,10 @@ def nbytes(self) -> int: class B12xGdnPrefill: - """Layer-held GDN capacity family with reusable staging buffers. + """Layer-held GDN capacity family with reusable sequence metadata. The layer supplies one declared ``Plan`` for every admitted capacity. - This helper owns only capacity staging buffers; the recurrent-state pool + This helper owns only metadata staging buffers; the recurrent-state pool remains a binding-time caller resource and can be invalidated without discarding the plans or these buffers. """ @@ -217,10 +199,6 @@ def __init__( ): raise ValueError("GDN staging buffers do not cover the prepared capacity") self.staging = staging - self.mixed_qkv = staging.mixed_qkv - self.a = staging.a - self.b = staging.b - self.output = staging.output self.query_start_loc = staging.query_start_loc self.initial_indices = staging.initial_indices self.final_indices = staging.final_indices @@ -259,9 +237,6 @@ def run( capacity = next(capacity for capacity in self.capacities if rows <= capacity) plan = self.plans[capacity] (scratch,) = get_b12x_scratch_buffers(plan) - self.mixed_qkv[:rows].copy_(mixed_qkv) - self.a[:rows].copy_(a) - self.b[:rows].copy_(b) _stage_metadata[(1,)]( query_start_loc, state_indices, @@ -281,18 +256,18 @@ def run( MAX_SEQS=self.max_seqs, BLOCK=triton.next_power_of_2(self.max_seqs + 1), ) - q, k, v = self.mixed_qkv[:capacity].split( + q, k, v = mixed_qkv.split( (self.key_heads * 128, self.key_heads * 128, self.value_heads * 128), dim=-1, ) binding = self.api.bind( plan, scratch=scratch, - q=q.view(capacity, self.key_heads, 128), - k=k.view(capacity, self.key_heads, 128), - v=v.view(capacity, self.value_heads, 128), - a=self.a[:capacity], - b=self.b[:capacity], + q=q.view(rows, self.key_heads, 128), + k=k.view(rows, self.key_heads, 128), + v=v.view(rows, self.value_heads, 128), + a=a, + b=b, A_log=self.A_log, dt_bias=self.dt_bias, recurrent_state=self.recurrent_state, @@ -303,13 +278,12 @@ def run( checkpoint_offsets=self.checkpoint_offsets, num_seqs=self.num_seqs, num_tokens=self.num_tokens, - output=self.output[:capacity], + output=output[:rows], ) self.api.run( binding, scale=scale, eps=eps, - max_live_tokens=capacity, + max_live_tokens=rows, max_live_seqs=self.max_seqs, ) - output[:rows].copy_(self.output[:rows]) From 2c9fa6b2574827a2139d5902f2a81601a4bffad0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:15:40 +0000 Subject: [PATCH 029/365] fix(ds4): reserve attention scratch from prepared declarations State preparation reserves the actual decode and single-pass extend plans before KV admission. Remove the independent split-envelope reservation from metadata-free profiling while retaining model-lane ownership and locked workspace bounds. Validation: metadata plans require at most 282 MiB for TP2 K5 attention with 4096 prefill rows and 48 decode rows, versus a 1633 MiB independent split-envelope estimate. Added metadata-free profile regression coverage; native and complete-model qualification pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit bd00a4acf031745194b43401e30776e2ed37fa7e) --- tests/v1/worker/test_workspace.py | 45 ++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 109 +------------------------ 2 files changed, 47 insertions(+), 107 deletions(-) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 1699fe90ea7f..d31428e706e7 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -313,3 +313,48 @@ def test_suspended_graph_resources_restore_collector_after_failure() -> None: raise ValueError("eager operation failed") assert workspace.retain_cuda_graph_capture_resource(owner) assert resources == [owner] + + +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( + monkeypatch, image_tokens +) -> None: + """Prepared state declarations own attention scratch before KV admission.""" + from b12x.attention import compressed_sparse_mla + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4 + + requested = [] + monkeypatch.setattr( + workspace, + "_manager", + SimpleNamespace(get_simultaneous=lambda *specs: requested.append(specs)), + ) + monkeypatch.setattr( + dsv4, "_require_b12x_compressed_sparse_mla", lambda: compressed_sparse_mla + ) + monkeypatch.setattr( + dsv4, "get_forward_context", lambda: SimpleNamespace(attn_metadata=None) + ) + layer = dsv4.DeepseekV4B12xAttention.__new__(dsv4.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.compress_ratio = 128 + layer.max_model_len = 1048576 + layer.max_num_batched_tokens = 4096 + layer.window_size = 128 + layer.max_image_tokens = image_tokens + layer.swa_cache_layer = SimpleNamespace(block_size=256) + layer.vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + use_dspark=lambda: True, num_speculative_tokens=5 + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4096, max_num_seqs=8), + ) + q = torch.randn((3, 32, 512), dtype=torch.bfloat16) + original_q = q.clone() + output = torch.full_like(q, float("nan")) + layer.forward_mqa(q, torch.empty(0), torch.arange(3), output) + + torch.testing.assert_close(q, original_q, rtol=0, atol=0) + assert torch.count_nonzero(output) == 0 + assert requested == [] diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 232d84b2dc94..14512c22cde3 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -4,9 +4,7 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass, replace -from functools import cache from typing import TYPE_CHECKING, Any, ClassVar, cast import torch @@ -60,7 +58,6 @@ DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.worker.workspace import ( - current_workspace_manager, retain_cuda_graph_capture_resource, ) @@ -558,24 +555,6 @@ def _c128a_profile_widths(max_width: int) -> tuple[int, ...]: return tuple(sorted(widths)) -@cache -def _max_q_chunks( - max_rows: int, - width: int, - split_chunks_for_contract: Callable[..., int], - decode_row_capacity: int | None, -) -> int: - return max( - rows - * split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - for rows in range(1, max(int(max_rows), 1) + 1) - ) - - def _cache_page_view( cache: torch.Tensor, page_size: int, @@ -1289,16 +1268,6 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: out = tensor_model_parallel_all_reduce(out) return out - def reserve_profile_scratch(self) -> None: - super().reserve_profile_scratch() - device = self.q_norm.weight.device - if device.type == "cuda": - # Reserve every layer type before compiled profile execution can - # skip its attention body. This tensor describes heads/device only. - self._reserve_profile_workspace( - torch.empty((0, self.padded_heads, _DSV4_HEAD_DIM), device=device) - ) - def _get_cache_page_view( self, cache: torch.Tensor, @@ -1312,81 +1281,6 @@ def _get_cache_page_view( self._b12x_cache_page_views[key] = view return view - def _reserve_profile_workspace(self, q: torch.Tensor) -> None: - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) - - module = _require_b12x_compressed_sparse_mla() - indexed_width = 0 - if self.compress_ratio == 4: - if self.topk_indices_buffer is not None: - indexed_width = int(self.topk_indices_buffer.shape[-1]) - elif self.indexer is not None: - indexed_width = int(self.indexer.topk_tokens) - elif self.compress_ratio > 1: - indexed_width = _c128a_topk_width( - self.max_model_len, - self.compress_ratio, - ) - - indexed_widths = ( - _c128a_profile_widths(indexed_width) - if self.compress_ratio > 4 - else (indexed_width,) - ) - swa_widths = { - int(self.window_size), - int(self.window_size) + self.max_image_tokens, - } - speculative_config = self.vllm_config.speculative_config - if speculative_config is not None and speculative_config.use_dspark(): - swa_widths.add( - get_dspark_swa_index_width( - int(self.window_size), - speculative_config.num_speculative_tokens or 0, - ) - ) - swa_width = max(swa_widths) - width = max(swa_width + indexed_width, 1) - rows = max(int(self.max_num_batched_tokens), 1) - decode_row_capacity = _get_dspark_decode_row_capacity(self.vllm_config) - max_chunks_per_row = module.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - # Split count is not monotonic in index width: a shorter prefix may - # use 12-token chunks while a longer prefix uses 64-token chunks. - # Cover every metadata width, not only the longest supported context. - max_q_chunks = max( - _max_q_chunks( - rows, - max(swa + indexed, 1), - module.split_chunks_for_contract, - decode_row_capacity, - ) - for swa in swa_widths - for indexed in indexed_widths - ) - # Allocation profiling needs geometry, not an executable declaration - # tied to live cache storage or an autotuning preparation session. - plan = plan_compressed_sparse_mla_scratch( - module.Caps( - device=q.device, - num_q_heads=int(q.shape[1]), - max_q_rows=rows, - max_width=width, - head_dim=_DSV4_HEAD_DIM, - v_head_dim=_DSV4_HEAD_DIM, - page_size=int(self.swa_cache_layer.block_size), - max_chunks_per_row=max_chunks_per_row, - max_q_chunks=max_q_chunks, - decode_row_capacity=decode_row_capacity, - ) - ) - current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) - def forward_mqa( self, q: torch.Tensor, @@ -1404,7 +1298,8 @@ def forward_mqa( attn_metadata = get_forward_context().attn_metadata if attn_metadata is None: output.zero_() - self._reserve_profile_workspace(q) + # State preparation reserves each declared operation's scratch + # before KV admission. This metadata-free pass skips attention. return assert isinstance(attn_metadata, dict) From b5a42ca91e212bc046d4a4aadd131fbba5642ecd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:20:04 +0000 Subject: [PATCH 030/365] test(ds4): retain query padding reservation with prepared scratch Verify that the B12X attention class inherits per-microbatch padded query reservation while attention scratch remains owned by prepared declarations. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 842bac6e0232a7796a01c9a3d0456e277a2ed9d9) --- tests/v1/worker/test_workspace.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index d31428e706e7..e5df32783a15 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -358,3 +358,31 @@ def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( torch.testing.assert_close(q, original_q, rtol=0, atol=0) assert torch.count_nonzero(output) == 0 assert requested == [] + + +def test_dsv4_profile_retains_padded_query_reservation() -> None: + from vllm.models.deepseek_v4.attention import DeepseekV4Attention + from vllm.models.deepseek_v4.nvidia.b12x import DeepseekV4B12xAttention + + assert ( + DeepseekV4B12xAttention.reserve_profile_scratch + is DeepseekV4Attention.reserve_profile_scratch + ) + layer = DeepseekV4B12xAttention.__new__(DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.kv_cache_torch_dtype = torch.uint8 + device = torch.device("cuda:0") + layer.q_norm = SimpleNamespace(weight=SimpleNamespace(device=device)) + layer._q_padded_scratch_num_ubatches = 2 + layer.max_num_batched_tokens = 4096 + layer.padded_heads = 32 + layer.head_dim = 512 + layer._q_padded_scratch_dtype = torch.bfloat16 + reserved = [] + layer._reserve_q_padded_scratch_buffer = lambda *args: reserved.append(args) + + layer.reserve_profile_scratch() + + assert reserved == [ + (4096, 32, 512, torch.bfloat16, device, ubatch) for ubatch in range(2) + ] From 3d9362ced812cf5dd513b36128cfc2d46e8757d8 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 06:23:19 +0000 Subject: [PATCH 031/365] test(ds4): validate declared attention workspace envelopes Exercise compressed-prefix, speculative, and multimodal scratch declarations through the preparation coordinator. Verify that both microbatch workspaces cover serving shapes after locking, without invoking the removed legacy split-envelope allocator. Import SimpleNamespace in the standalone workspace suite. Production code and serving behavior are unchanged. Validation: standalone test failure reproducer and pre-commit hooks. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 4b7b86dd057e334bc5fb953eea732ee395bfdc15) --- .../attention/test_ds4_workspace_envelope.py | 149 +++++++++++------- 1 file changed, 94 insertions(+), 55 deletions(-) diff --git a/tests/v1/attention/test_ds4_workspace_envelope.py b/tests/v1/attention/test_ds4_workspace_envelope.py index 404915c1469f..71c8a568b776 100644 --- a/tests/v1/attention/test_ds4_workspace_envelope.py +++ b/tests/v1/attention/test_ds4_workspace_envelope.py @@ -1,76 +1,115 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from math import prod from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from vllm.models.deepseek_v4.nvidia import b12x as adapter +from vllm.utils.b12x import B12xWorkload +from vllm.v1.worker import workspace as workspace_module +from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator +from vllm.v1.worker.workspace import WorkspaceManager @pytest.mark.skipif(not torch.cuda.is_available(), reason="B12X device plan required") @pytest.mark.parametrize("max_length", [786_688, 1_048_576]) @pytest.mark.parametrize("drafts", [0, 5]) -def test_profile_covers_shorter_compressed_prefixes(monkeypatch, max_length, drafts): - """Shorter prefixes can select more split chunks than the longest prefix.""" - from b12x.attention import compressed_sparse_mla as mla - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_declared_workspace_covers_shorter_compressed_prefixes( + monkeypatch, max_length, drafts, image_tokens +): + """Reserve decode-prefix and capacity-prefill plans before workspace lock. + Native plan scratch requirements are evaluated without compiling attention + kernels. The job stub completes only the coordinator's control protocol; + declaration, plan selection, and workspace reservation are real. + """ + device = torch.device("cuda", torch.accelerator.current_device_index()) spec = SimpleNamespace(use_dspark=lambda: True, num_speculative_tokens=drafts) - layer = SimpleNamespace( - compress_ratio=128, + layer = adapter.DeepseekV4B12xAttention.__new__(adapter.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.prefix = "model.layers.0.self_attn" + layer.indexer = None + layer.compress_ratio = 128 + layer.max_model_len = max_length + layer.max_num_batched_tokens = 4096 + layer.max_image_tokens = image_tokens + layer.window_size = 128 + layer.padded_heads = 32 + layer._b12x_cache_page_views = {} + layer._b12x_mla_plans = {} + layer.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=256), + speculative_config=spec if drafts else None, + scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=4096), + ) + page_bytes = adapter._DSV4_CACHE_BYTES_PER_TOKEN + layer.swa_cache_layer = SimpleNamespace( + block_size=256, + kv_cache=torch.empty((2, 256 * page_bytes), dtype=torch.uint8, device=device), + ) + layer.kv_cache = torch.empty((2, 2 * page_bytes), dtype=torch.uint8, device=device) + workload = B12xWorkload( + stage="state", + token_counts=(1, 8, 48, 128, 256, 4096), + fixed_token_counts=(1, 8, 48, 128, 256), + output_dtype=torch.bfloat16, + max_tokens=4096, + max_seqs=8, max_model_len=max_length, - max_num_batched_tokens=4096, - max_image_tokens=0, - window_size=128, - vllm_config=SimpleNamespace( - speculative_config=spec if drafts else None, - scheduler_config=SimpleNamespace( - max_num_seqs=8, max_num_batched_tokens=4096 - ), - ), - swa_cache_layer=SimpleNamespace(block_size=256), + speculative_tokens=drafts, ) - reserved = [] - monkeypatch.setattr( - adapter, - "current_workspace_manager", - lambda: SimpleNamespace( - get_simultaneous=lambda *items: reserved.append( - sum(prod(shape) * dtype.itemsize for shape, dtype in items) - ) - ), + (unit,) = layer.get_b12x_preparation_units(layer, workload) + queries = [request.plan.query for request in unit.requests] + widths = {128, 256, 512, 1024, 2048, 4096} + widths.add(6272 if max_length == 786_688 else 8192) + assert {q.indexed_width for q in queries} == widths + assert {q.query_rows for q in queries if q.mode == "decode"} == set(range(1, 257)) + assert {q.query_rows for q in queries if q.mode == "extend"} == {4096} + + job = Mock() + job.advance.return_value = SimpleNamespace( + done=True, pending_compilation=False, ready_collectives=() ) - adapter.DeepseekV4B12xAttention._reserve_profile_workspace( - layer, torch.empty((0, 32, 512), device="cuda") + session = Mock() + session.begin.return_value = job + workspace = WorkspaceManager(device, num_ubatches=2) + coordinator = B12xPreparationCoordinator( + session, + [(unit.requests, False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=workspace, ) - for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): - for index_width in (128, 256, 512, 1024, 2048, 4096): - width = 128 + index_width - capacity = 8 * (1 + drafts) if drafts else None - plan = plan_compressed_sparse_mla_scratch( - mla.Caps( - device=torch.device("cuda"), - num_q_heads=32, - max_q_rows=rows, - max_width=width, - head_dim=512, - v_head_dim=512, - page_size=256, - decode_row_capacity=capacity, - max_chunks_per_row=mla.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=capacity, - ), + outcome = coordinator.advance() + assert outcome["done"] and outcome["error"] is None, outcome + workspace.lock() + pointers = [buffer.data_ptr() for buffer in workspace._current_workspaces] + sizes = [buffer.numel() for buffer in workspace._current_workspaces] + assert sizes[0] == sizes[1] > 0 + + for ubatch in range(2): + monkeypatch.setattr( + workspace_module, "dbo_current_ubatch_id", lambda ubatch=ubatch: ubatch + ) + for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): + mode = "decode" if rows <= 256 else "extend" + swa_widths = {q.swa_width for q in queries if q.mode == mode} + for swa_width in swa_widths: + for index_width in widths: + plan = layer._b12x_mla_plan( + mode, + rows, + torch.empty((rows, swa_width), device="meta"), + torch.empty((rows, index_width), device="meta"), + ) + workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in plan.scratch_specs()) + ) + assert ( + workspace._current_workspaces[ubatch].data_ptr() == pointers[ubatch] ) - ) - required = sum( - prod(shape) * dtype.itemsize - for shape, dtype in plan.shapes_and_dtypes() - ) - assert max(reserved) >= required, (rows, width, reserved, required) From 101b033ebbda13a9c8d4925af4b4b693ace4fbee Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 09:53:36 +0000 Subject: [PATCH 032/365] fix(deepseek-v41): preserve tool namespaces and reminder messages Canonicalize explicit namespaces before Chat Completions validation can discard them. Keep definitions, named choices, assistant history and streamed responses on the same qualified function identity. Reject conflicting declarations without mutating input dictionaries, and accept the reference latest_reminder role in the V4.1 tokenizer. Retain ordinary tool names, reasoning budgets, image ordering and the existing parser. Add pinned reference prompt fixtures, API validation and streaming roundtrip coverage. Validation: 179 unit tests, pre-commit hooks, TP4 temperature-1 Chat Completions tool/history checks and image smoke checks pass. Reference contract: Qizhou Guo / DeepSeek, HF DeepSeek-V4.1-Flash dba1be0a40aa45a94ad051997016db3960a90277. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5b5dd1c055d4f9561d7bbfe221f004da73815448) --- docs/features/tool_calling.md | 42 ++ .../openai/test_tool_calls_serialization.py | 162 ++++++- tests/parser/engine/test_deepseek_v4.py | 93 ++++ .../fixtures/deepseek_v41/encoding.json | 410 ++++++++++++++++++ tests/tokenizers_/test_deepseek_v41.py | 66 +++ .../openai/chat_completion/protocol.py | 33 +- vllm/tokenizers/deepseek_v41.py | 9 +- vllm/tokenizers/deepseek_v41_encoding.py | 8 +- vllm/utils/tool_names.py | 67 +++ 9 files changed, 882 insertions(+), 8 deletions(-) create mode 100644 tests/tokenizers_/fixtures/deepseek_v41/encoding.json create mode 100644 vllm/utils/tool_names.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 9628bc431a1e..2e8178d392ad 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -338,6 +338,48 @@ Supported models: Flags: `--tool-call-parser deepseek_v31 --chat-template {see_above}` +### DeepSeek-V4.1 Models (`deepseek_v41`) + +Status: implemented. The `deepseek_v41` tokenizer and parser support spaced +DSML tool calls and namespace-qualified function names for +`deepseek-ai/DeepSeek-V4.1-Flash`. + +Use `--tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 +--reasoning-parser deepseek_v41 --enable-auto-tool-choice`. + +Chat Completions accepts a namespace on a function tool: + +```json +{ + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "lookup", + "description": "Look up an item.", + "parameters": {"type": "object", "properties": {"sku": {"type": "string"}}} + } +} +``` + +The namespace may also be a string, or appear inside `function`. Before +validation and rendering, it becomes `function.name="inventory::lookup"`. +The namespace description is prepended to the function description. +The same normalization applies to assistant tool-call history and named +`tool_choice` objects. A qualified name must agree with any explicit namespace; +conflicting declarations are rejected. Tools without namespaces are unchanged. + +Complete and streamed Chat Completions responses retain the qualified +`function.name`; they do not add a separate response `namespace` field. Dispatch +tools by that qualified name and preserve it when sending assistant history. +This extension is separate from the Responses API's `type="namespace"` schema. + +The tokenizer also accepts `{"role":"latest_reminder","content":"..."}` +between a system message and a user message, rendering the reference +`<｜latest_reminder｜>` marker for date, locale or similar context. This is a +DeepSeek-specific role, not a standard OpenAI message role. Other unknown roles +remain invalid. The encoding contract follows the +[DeepSeek V4.1 reference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/dba1be0a40aa45a94ad051997016db3960a90277/encoding/README.md). + ### OpenAI OSS Models (`openai`) Supported models: diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index cedc2575b80e..478379f5e13b 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -15,9 +15,16 @@ "ValueError: Unexpected tool call id ...". """ +import copy + import pytest +from pydantic import ValidationError -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.exceptions import VLLMValidationError def _make_tool_call(tc_id: str, name: str, args: str) -> dict: @@ -148,3 +155,156 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): # Verify after model_dump_json too _ = req.model_dump_json() assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +@pytest.mark.parametrize("iterator", [False, True]) +def test_namespaces_survive_definitions_choices_and_history( + location, namespace, iterator +): + def qualify(item, value): + target = item if location == "tool" else item["function"] + target["namespace"] = value + return item + + tools = [ + qualify({"type": "function", "function": {"name": "lookup"}}, namespace), + qualify({"type": "function", "function": {"name": "lookup"}}, "billing"), + ] + calls = [qualify(_make_tool_call("call_1", "lookup", "{}"), namespace)] + payload = { + "model": "DeepSeek-V4.1-Flash", + "tools": tools, + "tool_choice": qualify( + {"type": "function", "function": {"name": "lookup"}}, namespace + ), + "messages": [ + {"role": "user", "content": "Look up stock."}, + {"role": "assistant", "content": None, "tool_calls": calls}, + ], + } + original = copy.deepcopy(payload) + if iterator: + payload["messages"][1]["tool_calls"] = iter(calls) + request = ChatCompletionRequest.model_validate(payload) + assert [tool.function.name for tool in request.tools] == [ + "inventory::lookup", + "billing::lookup", + ] + assert request.tool_choice.function.name == "inventory::lookup" + assert request.messages[1]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + serialized = request.model_dump_json() + assert request.model_dump_json() == serialized + restored = ChatCompletionRequest.model_validate_json(serialized) + assert restored.model_dump_json() == serialized + if not iterator: + assert payload == original + assert calls == original["messages"][1]["tool_calls"] + + +def test_namespace_description_is_preserved_exactly_once(): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "inventory::lookup", + "description": "Find a SKU.", + "strict": True, + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == "Stock operations.\nFind a SKU." + assert parsed.function.strict is True + assert tool == original + assert ChatCompletionToolsParam.model_validate( + parsed.model_dump() + ).model_dump() == (parsed.model_dump()) + + +@pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) +@pytest.mark.parametrize( + "namespace, name", + [ + ("inventory", "billing::lookup"), + ("a::b", "lookup"), + ("", "lookup"), + ({}, "lookup"), + ({"name": 7}, "lookup"), + (False, "lookup"), + ("inventory", "inventory::x::lookup"), + ("inventory", "inventory::"), + ], +) +def test_invalid_namespaces_are_client_errors(namespace, name, tool_choice): + with pytest.raises(ValidationError, match="[Nn]amespace"): + ChatCompletionRequest.model_validate( + { + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "namespace": namespace, + "function": {"name": name}, + } + ], + "tool_choice": tool_choice, + } + ) + + +def test_conflicting_namespace_locations_are_rejected(): + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup", "namespace": "billing"}, + } + ) + + +def test_consistent_namespace_locations_do_not_duplicate_the_prefix(): + parsed = ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": {"name": "inventory"}, + "function": {"name": "inventory::lookup", "namespace": "inventory"}, + } + ) + assert parsed.function.name == "inventory::lookup" + + +def test_conflicting_history_namespace_is_not_silently_dropped(): + call = _make_tool_call("call_a", "inventory::lookup", "{}") + call["namespace"] = "billing" + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + _make_request([{"role": "assistant", "content": None, "tool_calls": [call]}]) + + +def test_bare_choice_cannot_select_a_namespaced_tool(): + with pytest.raises(VLLMValidationError, match="does not match"): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "lookup"}], + "tools": [ + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup"}, + } + ], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + } + ) + + +def test_plain_and_already_qualified_names_keep_their_spelling(): + for name in ("lookup", "inventory::lookup", "opaque::legacy::name"): + tool = {"type": "function", "function": {"name": name}} + assert ChatCompletionToolsParam.model_validate(tool).function.name == name diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 7cbf9791a7b0..9180a6f4b2d5 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -21,6 +21,7 @@ simulate_reasoning_streaming, simulate_tool_streaming, ) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.deepseek_v4 import ( DSML_INVOKE_END, @@ -36,10 +37,13 @@ _unwrap_wrapper_args, deepseek_v4_config, ) +from vllm.parser.deepseek_v41 import DeepSeekV41Parser +from vllm.parser.engine.adapters import make_adapters from vllm.parser.engine.registered_adapters import ( DeepSeekV4ParserReasoningAdapter, DeepSeekV4ParserToolAdapter, ) +from vllm.tokenizers.deepseek_v41_encoding import encode_messages _THINK_START_ID = 50 _THINK_END_ID = 51 @@ -48,6 +52,95 @@ _PARAM_CLOSE = "" +@pytest.mark.parametrize("chunk_size", [1, 2, 7, 17, 1000]) +def test_v41_namespaced_tools_keep_identity_in_full_stream_and_history(chunk_size): + names = ["inventory::lookup", "billing::lookup"] + request = ChatCompletionRequest( + messages=[{"role": "user", "content": "Look up A1."}], + tools=[ + { + "type": "function", + "namespace": namespace, + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": {"sku": {"type": "string"}}, + }, + }, + } + for namespace in ("inventory", "billing") + ], + ) + tokenizer = make_mock_tokenizer({"": 50, "": 51}) + _, adapter = make_adapters(DeepSeekV41Parser) + + def parser(): + return adapter( + tokenizer, request.tools, chat_template_kwargs={"thinking": False} + ) + + text = ( + "<｜DSML｜ calls>\n" + + "\n".join( + f'<｜DSML｜ invoke name="{name}">\n' + '<｜DSML｜ parameter name="sku" string="true">A1\n' + "" + for name in names + ) + + "\n" + ) + calls = parser().extract_tool_calls(text, request).tool_calls + assert [call.function.name for call in calls] == names + assert [json.loads(call.function.arguments) for call in calls] == [ + {"sku": "A1"} + ] * 2 + stream = parser() + collected: dict[int, dict[str, str]] = {} + + def collect(delta): + if delta is None: + return + for call in delta.tool_calls: + value = collected.setdefault(call.index, {"name": "", "arguments": ""}) + if call.function: + value["name"] += call.function.name or "" + value["arguments"] += call.function.arguments or "" + + previous = "" + for offset in range(0, len(text), chunk_size): + chunk = text[offset : offset + chunk_size] + collect( + stream.extract_tool_calls_streaming( + previous, + previous + chunk, + chunk, + (), + (), + (), + request, + ) + ) + previous += chunk + collect(stream.finish_streaming()) + assert [collected[i]["name"] for i in sorted(collected)] == names + assert [json.loads(collected[i]["arguments"]) for i in sorted(collected)] == ( + [{"sku": "A1"}] * 2 + ) + prompt = encode_messages( + [ + {"role": "user", "content": "Look up A1."}, + { + "role": "assistant", + "content": "", + "tool_calls": [call.model_dump(exclude_none=True) for call in calls], + }, + ], + thinking_mode="chat", + ) + assert all(f'name="{name}"' in prompt for name in names) + + def _param(name: str, is_str: str, value: str) -> str: return f"<{_PARAM_OPEN.format(name=name, is_str=is_str)}{value}{_PARAM_CLOSE}" diff --git a/tests/tokenizers_/fixtures/deepseek_v41/encoding.json b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json new file mode 100644 index 000000000000..532e842e96aa --- /dev/null +++ b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json @@ -0,0 +1,410 @@ +{ + "reference": "https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/raw/dba1be0a40aa45a94ad051997016db3960a90277/encoding/encoding.py", + "reference_author": "DeepSeek; namespace contract by Qizhou Guo", + "cases": [ + { + "name": "reminder_chat", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>You are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "reminder_thinking", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\nYou are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "mid_system_chat", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "mid_system_thinking", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "reasoning_history_chat", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "reasoning_history_thinking", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "interleaved_images_chat", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "interleaved_images_thinking", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "namespaced_tool_history_false", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "namespaced_tool_history_true", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>Consult both tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "action_task", + "messages": [ + { + "role": "user", + "content": "Find the documentation.", + "task": "action" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Find the documentation.<｜Assistant｜><｜action｜>" + }, + { + "name": "title_task", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi.", + "task": "title" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜title｜>" + } + ] +} diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index d0916fb6077b..5035a7e71431 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/tests/tokenizers_/test_deepseek_v41.py @@ -4,12 +4,14 @@ import copy import json from pathlib import Path +from typing import Any import pytest from tests.tokenizers_.test_deepseek_v4 import FakeHfTokenizer from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tokenizers.deepseek_v41 import get_deepseek_v41_tokenizer +from vllm.tokenizers.deepseek_v41_encoding import encode_messages FIXTURES = Path(__file__).parent / "fixtures" @@ -185,3 +187,67 @@ def image_part(color): def test_unsupported_media_is_a_request_error(): with pytest.raises(ValueError, match="text and image content only"): render([{"role": "user", "content": [{"type": "input_audio"}]}]) + + +_FIXTURES = json.loads( + (Path(__file__).parent / "fixtures/deepseek_v41/encoding.json").read_text() +) + + +@pytest.mark.parametrize("case", _FIXTURES["cases"], ids=lambda case: case["name"]) +def test_prompt_matches_reference_fixture(case): + original = copy.deepcopy(case["messages"]) + request = ChatCompletionRequest( + messages=case["messages"], + tools=case.get("tools"), + ) + prompt = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tools=[tool.model_dump(exclude_none=True) for tool in request.tools] + if request.tools + else None, + tokenize=False, + **case["kwargs"], + ) + assert prompt == case["expected"] + assert case["messages"] == original + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +def test_direct_encoder_and_api_preserve_the_same_tool_identity(location, namespace): + tool: dict[str, Any] = {"type": "function", "function": {"name": "lookup"}} + call: dict[str, Any] = { + "id": "call_a", + "type": "function", + "function": {"name": "lookup", "arguments": '{"sku":"A1"}'}, + } + for item in (tool, call): + (item if location == "tool" else item["function"])["namespace"] = namespace + messages = [ + {"role": "system", "content": "Use tools.", "tools": [tool]}, + {"role": "user", "content": "Find A1."}, + {"role": "assistant", "content": "", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call_a", "content": "In stock."}, + ] + original = copy.deepcopy(messages) + direct = encode_messages(messages, thinking_mode="chat") + request = ChatCompletionRequest(messages=messages) + rendered = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tokenize=False, + thinking=False, + ) + assert rendered == direct + assert '"name": "inventory::lookup"' in rendered + assert 'name="inventory::lookup"' in rendered + assert "In stock." in rendered + assert messages == original + + +def test_reminder_does_not_enable_arbitrary_message_roles(): + with pytest.raises(ValueError, match="Invalid role: SYSTEM"): + get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + [{"role": "SYSTEM", "content": "Hello"}], + tokenize=False, + ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index badd5a0344c6..f9d07659cdb8 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -51,6 +51,7 @@ ThinkingTokenBudget, ) from vllm.utils import random_uuid +from vllm.utils.tool_names import normalize_tool_namespace logger = init_logger(__name__) @@ -186,6 +187,11 @@ class ChatCompletionToolsParam(OpenAIBaseModel): function: FunctionDefinition defer_loading: bool | None = None + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + @model_validator(mode="after") def _propagate_defer_loading(self) -> "ChatCompletionToolsParam": if self.defer_loading is not None and self.function.defer_loading is None: @@ -209,6 +215,11 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): function: ChatCompletionNamedFunction type: Literal["function"] = "function" + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + class ChatCompletionRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation @@ -534,23 +545,31 @@ def check_cache_salt_support(cls, data: Any) -> Any: def _normalize_messages_before(cls, data: Any) -> Any: """Pre-process message dicts before Pydantic field validation. - Performs two normalizations in a single pass: + Normalizes tool identities and message fields in a single pass: - Converts tool_calls generators/iterators to lists so one-shot generators are not consumed during union type matching. - Renames the deprecated ``reasoning_content`` field to ``reasoning`` so downstream code only needs to check one field. + - Qualifies tool-call namespaces before union validation can drop them. """ if not isinstance(data, dict): return data messages = data.get("messages") if not isinstance(messages, list): return data + data = dict(data) + messages = [dict(msg) if isinstance(msg, dict) else msg for msg in messages] + data["messages"] = messages for msg in messages: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None and not isinstance(tool_calls, list): - msg["tool_calls"] = list(tool_calls) + if tool_calls is not None: + msg["tool_calls"] = [ + normalize_tool_namespace(call) for call in tool_calls + ] + if isinstance(msg.get("tools"), list): + msg["tools"] = [normalize_tool_namespace(tool) for tool in msg["tools"]] reasoning_content = msg.pop("reasoning_content", None) if reasoning_content is not None and msg.get("reasoning") is None: msg["reasoning"] = reasoning_content @@ -916,6 +935,14 @@ def check_tool_usage(cls, data): if not isinstance(data, dict): return data + # Named choices must be compared with qualified identities before + # Pydantic serializes function definitions and discards extra fields. + data = dict(data) + if isinstance(data.get("tools"), list): + data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + if isinstance(data.get("tool_choice"), dict): + data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) + # Reject empty tools array, matching OpenAI API behavior if data.get("tools") == []: raise VLLMValidationError( diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index 3405cbbbe7dc..68ba93e5ec8a 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -23,7 +23,14 @@ def _normalize_messages( result = [dict(message) for message in copy.deepcopy(messages)] for message in result: role = message.get("role") - if role not in ("system", "developer", "user", "assistant", "tool"): + if role not in ( + "system", + "developer", + "user", + "assistant", + "tool", + "latest_reminder", + ): raise ValueError(f"Invalid role: {role}") if role == "developer": message["role"] = "system" diff --git a/vllm/tokenizers/deepseek_v41_encoding.py b/vllm/tokenizers/deepseek_v41_encoding.py index 995115f1f7d0..89907d152bef 100644 --- a/vllm/tokenizers/deepseek_v41_encoding.py +++ b/vllm/tokenizers/deepseek_v41_encoding.py @@ -11,6 +11,8 @@ import json from typing import Any, Dict, List, Optional, Union +from vllm.utils.tool_names import normalize_tool_namespace + bos_token: str = "<｜begin▁of▁sentence｜>" eos_token: str = "<｜end▁of▁sentence｜>" thinking_start_token: str = "" @@ -55,8 +57,8 @@ def to_json(value: Any) -> str: def tools_from_openai_format(tools): - """Extract function definitions from OpenAI-format tool list.""" - return [tool["function"] for tool in tools] + """Extract function definitions with namespace-qualified identities.""" + return [normalize_tool_namespace(tool)["function"] for tool in tools] def tool_calls_from_openai_format(tool_calls): @@ -66,7 +68,7 @@ def tool_calls_from_openai_format(tool_calls): "name": tool_call["function"]["name"], "arguments": tool_call["function"]["arguments"], } - for tool_call in tool_calls + for tool_call in map(normalize_tool_namespace, tool_calls) ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py new file mode 100644 index 000000000000..02fcbc9e7c2e --- /dev/null +++ b/vllm/utils/tool_names.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Canonical function identities for namespace-bearing chat tools. + +The namespace spelling follows Qizhou Guo's DeepSeek V4.1 reference encoder: +https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/commit/dba1be0a40aa45a94ad051997016db3960a90277 +""" + +from typing import Any + + +def _namespace_parts(namespace: Any) -> tuple[str, str | None]: + description = None + if isinstance(namespace, dict): + description = namespace.get("description") + namespace = namespace.get("name") + if not isinstance(namespace, str) or not namespace or "::" in namespace: + raise ValueError("Tool namespace must be a nonempty string without '::'") + if description is not None and not isinstance(description, str): + raise ValueError("Tool namespace description must be a string") + return namespace, description + + +def normalize_tool_namespace(tool: Any) -> Any: + """Qualify explicit namespaces without mutating a definition or tool call. + + Accept a string or {name, description} namespace on the tool or function. + Consistent duplicate declarations are allowed; conflicting identities are + rejected. Remove namespace metadata after folding it into the function name + and description, so normalization is idempotent across API and tokenizer + boundaries. Inputs without namespace metadata retain their existing names. + """ + if not isinstance(tool, dict) or not isinstance(tool.get("function"), dict): + return tool + function = tool["function"] + declarations = [ + _namespace_parts(value) + for value in (tool.get("namespace"), function.get("namespace")) + if value is not None + ] + if not declarations: + return tool + namespace, description = declarations[0] + if any(name != namespace for name, _ in declarations[1:]): + raise ValueError("Conflicting tool namespaces on the tool and function") + + name = function.get("name") + if not isinstance(name, str) or not name: + raise ValueError("Namespaced function name must be a nonempty string") + if "::" in name: + prefix, name = name.split("::", 1) + if prefix != namespace: + raise ValueError("Conflicting tool namespaces in the qualified name") + if not name or "::" in name: + raise ValueError("Namespaced function name must be nonempty without '::'") + + normalized = dict(tool) + normalized.pop("namespace", None) + normalized["function"] = function = dict(function) + function.pop("namespace", None) + function["name"] = f"{namespace}::{name}" + if description: + detail = function.get("description") + if detail is not None and not isinstance(detail, str): + raise ValueError("Function description must be a string") + function["description"] = description + "\n" + (detail or "") + return normalized From 61f1bd482f52f8cf02b9c3817fbe0a0c04f5204f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 10:55:43 +0000 Subject: [PATCH 033/365] fix(tools): preserve validation paths and namespace descriptions Normalize supported tool-call iterables without preempting scalar field validation. Retain the first non-null namespace description while rejecting conflicting names. Focused regressions reproduce both review findings; 50 tool-call serialization tests pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit fb50f44a5f025c9c2eed857fbac938e8513227a7) --- .../openai/test_tool_calls_serialization.py | 37 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 5 ++- vllm/utils/tool_names.py | 7 +++- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 478379f5e13b..156d6bb140c4 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -226,6 +226,43 @@ def test_namespace_description_is_preserved_exactly_once(): ).model_dump() == (parsed.model_dump()) +@pytest.mark.parametrize("tool_calls", [1, False, 1.5]) +def test_invalid_tool_calls_shape_reports_the_field(tool_calls): + with pytest.raises(ValidationError) as error: + _make_request( + [{"role": "assistant", "content": None, "tool_calls": tool_calls}] + ) + assert any("tool_calls" in item["loc"] for item in error.value.errors()) + + +@pytest.mark.parametrize( + "outer, inner, expected", + [ + (None, "Stock operations.", "Stock operations.\nFind a SKU."), + ("Outer description.", "Inner description.", "Outer description.\nFind a SKU."), + ("", "Inner description.", "Find a SKU."), + ], +) +def test_consistent_namespaces_preserve_first_non_null_description( + outer, inner, expected +): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": outer}, + "function": { + "name": "lookup", + "namespace": {"name": "inventory", "description": inner}, + "description": "Find a SKU.", + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == expected + assert tool == original + assert ChatCompletionToolsParam.model_validate(parsed.model_dump()) == parsed + + @pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) @pytest.mark.parametrize( "namespace, name", diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index f9d07659cdb8..e5913efef6c6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -4,6 +4,7 @@ # Adapted from # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py import time +from collections.abc import Iterable from typing import Annotated, Any, ClassVar, Literal from openai.types.chat.chat_completion_audio import ( @@ -564,7 +565,9 @@ def _normalize_messages_before(cls, data: Any) -> Any: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None: + if isinstance(tool_calls, Iterable) and not isinstance( + tool_calls, (str, bytes, dict) + ): msg["tool_calls"] = [ normalize_tool_namespace(call) for call in tool_calls ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py index 02fcbc9e7c2e..1cf3f368d464 100644 --- a/vllm/utils/tool_names.py +++ b/vllm/utils/tool_names.py @@ -41,8 +41,11 @@ def normalize_tool_namespace(tool: Any) -> Any: if not declarations: return tool namespace, description = declarations[0] - if any(name != namespace for name, _ in declarations[1:]): - raise ValueError("Conflicting tool namespaces on the tool and function") + for declared_name, declared_description in declarations[1:]: + if declared_name != namespace: + raise ValueError("Conflicting tool namespaces on the tool and function") + if description is None: + description = declared_description name = function.get("name") if not isinstance(name, str) or not name: From e14c9f79512c34a65a65f8b9bacb02f2f15d1b50 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 11:13:52 +0000 Subject: [PATCH 034/365] perf(comm): index declared PCIe collective plans by their full call contract Bound lookup storage to declared shapes and avoid scans and exception construction when unsupported prefill shapes fall through to another transport. Preserve normalization weight identity, epsilon, strides and first-declaration precedence. Inspired by fernandaspets/vllm_sm120#2. Native suite: 36 passed including TP2 fused collective replay, one four-GPU case skipped. Serving throughput is not yet qualified. Co-authored-by: neolithic5452 Co-authored-by: Neo Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5e693878601f1145e41582999439080687cc0bf1) --- benchmarks/benchmark_b12x_plan_lookup.py | 81 ++++++++++++++++ .../distributed/test_b12x_pcie_all_reduce.py | 86 +++++++++++++++++ .../b12x_pcie_all_reduce.py | 93 ++++++++++++------- 3 files changed, 229 insertions(+), 31 deletions(-) create mode 100644 benchmarks/benchmark_b12x_plan_lookup.py diff --git a/benchmarks/benchmark_b12x_plan_lookup.py b/benchmarks/benchmark_b12x_plan_lookup.py new file mode 100644 index 000000000000..fc5c8fe0dcba --- /dev/null +++ b/benchmarks/benchmark_b12x_plan_lookup.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure CPU-only lookup cost for declared B12X PCIe collective shapes. + +No GPU operation is executed. Use --module-file to compare pinned source files +in the same Python/runtime environment; these timings are not serving throughput. +""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +import sys +import timeit +from pathlib import Path + +import torch + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--module-file", type=Path, required=True) + parser.add_argument("--declarations", type=int, nargs="+", default=[64, 512, 4096]) + parser.add_argument("--iterations", type=int, default=1000) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + if min(*args.declarations, args.iterations, args.repeats) < 1: + parser.error("counts must be positive") + spec = importlib.util.spec_from_file_location("pcie_lookup_probe", args.module_file) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + report = { + "source_sha256": hashlib.sha256(args.module_file.read_bytes()).hexdigest(), + "iterations": args.iterations, + "repeats": args.repeats, + "scope": "CPU metadata lookup; no collective or model execution", + "results": [], + } + for count in args.declarations: + communicator = object.__new__(module.B12xPcieAllReduce) + invocations = [ + module.B12xPcieInvocation( + name=str(rows), + operation="all_reduce", + shape=(rows, 16), + dtype=torch.bfloat16, + ) + for rows in range(1, count + 1) + ] + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + if hasattr(communicator, "_index_declared_plans"): + communicator._index_declared_plans() + for case, rows, expected in (("hit", count, True), ("miss", count + 1, False)): + source = torch.empty((rows, 16), dtype=torch.bfloat16) + + def lookup(communicator=communicator, source=source): + return communicator._has_plan_for(source) + + assert lookup() is expected + for _ in range(100): + lookup() + samples = timeit.repeat(lookup, number=args.iterations, repeat=args.repeats) + report["results"].append( + { + "declarations": count, + "case": case, + "microseconds_per_call": [ + s * 1e6 / args.iterations for s in samples + ], + "median_us": statistics.median(samples) * 1e6 / args.iterations, + } + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_b12x_pcie_all_reduce.py b/tests/distributed/test_b12x_pcie_all_reduce.py index 4a4b6a21eeda..017883b4a3e5 100644 --- a/tests/distributed/test_b12x_pcie_all_reduce.py +++ b/tests/distributed/test_b12x_pcie_all_reduce.py @@ -58,9 +58,11 @@ def _make_communicator( communicator.twoshot_max_bytes = 0 plan = object() communicator._plans = {"prepared": plan} + communicator._plan_index = {} communicator._routes = {} communicator._invocations = {} communicator._plan_for = MagicMock(return_value=plan) + communicator._lookup_plan = MagicMock(return_value=plan) return communicator, runtime @@ -323,6 +325,7 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: communicator, _ = _make_communicator() communicator._plans = {} communicator._invocations = {} + communicator._lookup_plan.return_value = None with pytest.raises( PreparationResourceUnavailableError, @@ -331,6 +334,89 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: B12xPcieAllReduce._plan_for(communicator, torch.randn(2, 4)) +def _indexed_communicator(invocations): + communicator = object.__new__(B12xPcieAllReduce) + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + communicator._index_declared_plans() + return communicator + + +def test_declared_plan_index_preserves_complete_fused_identity(): + operation = "all_reduce_fused_add_rms_norm" + weight_a = torch.ones(4) + weight_b = weight_a.clone() + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocations = [ + b12x_pcie_all_reduce.B12xPcieInvocation( + name=name, + operation=operation, + shape=(2, 4), + dtype=source.dtype, + norm_weight=weight, + epsilon=epsilon, + ) + for name, weight, epsilon in ( + ("a", weight_a, 1e-6), + ("b", weight_b, 1e-6), + ("epsilon", weight_a, 1e-5), + ) + ] + communicator = _indexed_communicator(invocations) + for invocation in invocations: + assert ( + communicator._plan_for( + source, + operation=operation, + weight=invocation.norm_weight, + epsilon=invocation.epsilon, + ) + is communicator._plans[invocation.name] + ) + assert not communicator._has_plan_for(source) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a.clone(), epsilon=1e-6 + ) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a, epsilon=1e-4 + ) + + +def test_declared_plan_index_matches_shape_dtype_stride_and_first_declaration(): + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocation = b12x_pcie_all_reduce.B12xPcieInvocation + communicator = _indexed_communicator( + [ + invocation("first", "all_reduce", (2, 4), source.dtype), + invocation( + "equivalent", "all_reduce", (2, 4), source.dtype, strides=(4, 1) + ), + invocation("strided", "all_reduce", (2, 4), source.dtype, strides=(8, 1)), + ] + ) + assert communicator._plan_for(source) is communicator._plans["first"] + strided = torch.empty((2, 8), dtype=source.dtype)[:, :4] + assert communicator._plan_for(strided) is communicator._plans["strided"] + assert not communicator._has_plan_for(source.float()) + assert not communicator._has_plan_for(source.flatten()) + communicator._invocations = {} + assert communicator._plan_for(source) is communicator._plans["first"] + communicator._index_declared_plans() + assert not communicator._has_plan_for(source) + + +def test_undeclared_plan_probe_neither_raises_nor_grows_the_index(): + communicator = _indexed_communicator([]) + for rows in range(1, 65): + source = torch.empty((rows, 4), dtype=torch.bfloat16) + assert not communicator._has_plan_for(source) + with pytest.raises( + PreparationResourceUnavailableError, match="no declared plan" + ): + communicator._plan_for(source) + assert communicator._plan_index == {} + + def test_fused_allreduce_has_an_independent_cutoff() -> None: communicator, runtime = _make_communicator( allreduce_max_bytes=16, fused_max_bytes=64 diff --git a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py index 16cf5b31fadf..e82d1510731f 100644 --- a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py @@ -248,6 +248,7 @@ def __init__( ] ] = [] self._plans: dict[str, object] = {} + self._plan_index: dict[tuple, object] = {} self._invocations: dict[str, B12xPcieInvocation] = {} self._routes: dict[str, str] = {} if len(self.global_ranks) != self.world_size: @@ -596,12 +597,14 @@ def get_b12x_preparation_units( plans: dict[str, object] = {} routes: dict[str, str] = {} collective_ranks = tuple(sorted(self.global_ranks)) - declarations = {} + declarations: dict[tuple, Any] = {} + prepare: Callable[..., Any] for invocation in invocations: route = self._route_invocation(invocation) if route is None: continue routes[invocation.name] = route + assert self._runtime is not None target = self._runtime._prepared_channel_for_stream( None, invocation.channel_id ) @@ -630,7 +633,7 @@ def get_b12x_preparation_units( plan = _oneshot_preparation.plan(query, runtime=target) declarations[key] = plan - def prepare(state, invocation=invocation, query=query): + def prepare_oneshot(state, invocation=invocation, query=query): registered = bool(query.setup["registered"]) inp = self._request_input(invocation, registered=registered) reset, produce, restore, owners = self._oneshot_prime_input( @@ -665,6 +668,8 @@ def produce_fused() -> None: restore=restore, owners=owners, ) + + prepare = prepare_oneshot elif route == "twoshot": assert self._twoshot is not None query = _twoshot_preparation.query_from_metadata( @@ -681,13 +686,15 @@ def produce_fused() -> None: plan = _twoshot_preparation.plan(query, runtime=self._twoshot) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_twoshot(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) return _twoshot_preparation.prepared_call( state, payload=inp, out=torch.empty_like(inp) ) + + prepare = prepare_twoshot else: assert self._dma is not None query = _dma_preparation.query_from_metadata( @@ -703,7 +710,7 @@ def prepare(state, invocation=invocation): plan = _dma_preparation.plan(query, runtime=self._dma) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_dma(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) @@ -711,6 +718,7 @@ def prepare(state, invocation=invocation): state, inp=inp, out=torch.empty_like(inp) ) + prepare = prepare_dma plans[invocation.name] = plan requests.append( plan.request( @@ -732,6 +740,7 @@ def prepare(state, invocation=invocation): } self._plans = plans self._routes = routes + self._index_declared_plans() if not requests: return () return ( @@ -743,6 +752,47 @@ def prepare(state, invocation=invocation): ), ) + @staticmethod + def _plan_key(operation, shape, dtype, strides, weight=None, epsilon=None): + norm = None if operation == "all_reduce" else (id(weight), epsilon) + return operation, tuple(shape), dtype, tuple(strides), norm + + def _index_declared_plans(self) -> None: + index: dict[tuple, object] = {} + for name, invocation in self._invocations.items(): + key = self._plan_key( + invocation.operation, + invocation.shape, + invocation.dtype, + invocation.strides + if invocation.strides is not None + else self._contiguous_strides(invocation.shape), + invocation.norm_weight, + invocation.epsilon, + ) + # Equivalent declarations keep the first plan, matching model order. + index.setdefault(key, self._plans[name]) + self._plan_index = index + + def _lookup_plan( + self, + inp: torch.Tensor, + *, + operation: str = "all_reduce", + weight: torch.Tensor | None = None, + epsilon: float | None = None, + ): + return self._plan_index.get( + self._plan_key( + operation, + inp.shape, + inp.dtype, + inp.stride(), + weight, + epsilon, + ) + ) + def _plan_for( self, inp: torch.Tensor, @@ -751,32 +801,16 @@ def _plan_for( weight: torch.Tensor | None = None, epsilon: float | None = None, ): - strides = tuple(inp.stride()) - for name, invocation in self._invocations.items(): - if ( - invocation.operation == operation - and tuple(inp.shape) == invocation.shape - and inp.dtype == invocation.dtype - and strides - == ( - self._contiguous_strides(invocation.shape) - if invocation.strides is None - else invocation.strides - ) - and ( - operation == "all_reduce" - or ( - invocation.norm_weight is weight - and invocation.epsilon == epsilon - ) - ) - ): - return self._plans[name] + plan = self._lookup_plan( + inp, operation=operation, weight=weight, epsilon=epsilon + ) + if plan is not None: + return plan from vllm.utils.b12x import PreparationResourceUnavailableError raise PreparationResourceUnavailableError( f"PCIe collective has no declared plan for {operation} on shape " - f"{tuple(inp.shape)} {inp.dtype} strides {strides}; declared: " + f"{tuple(inp.shape)} {inp.dtype} strides {tuple(inp.stride())}; declared: " + ", ".join( f"{invocation.operation}{invocation.shape}" for invocation in self._invocations.values() @@ -791,11 +825,7 @@ def _has_plan_for(self, inp: torch.Tensor, **kwargs) -> bool: the captured and planned serving shapes, so an undeclared shape is a prefill chunk or warm-up batch outside the graph-replayed set. """ - from vllm.utils.b12x import PreparationResourceUnavailableError - - try: - self._plan_for(inp, **kwargs) - except PreparationResourceUnavailableError: + if self._lookup_plan(inp, **kwargs) is None: logger.debug( "b12x PCIe all-reduce declines undeclared shape %s %s", tuple(inp.shape), @@ -953,6 +983,7 @@ def close(self) -> None: self._runtime.close() self._runtime = None self._plans.clear() + self._plan_index.clear() self._invocations.clear() self._routes.clear() self.disabled = True From cdddc4244bdc3044b1e94d1aae491eaea1acde5d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:36:27 +0000 Subject: [PATCH 035/365] fix(glm): invalidate physical-selection plans with cache bindings Preserve the KK prepared-state call interface and invalidate geometry-dependent declarations on cache bind/unbind. Extend the existing construction test to cover declaration lifetime. Reconciles the remaining ownership part of vLLM #789; native qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 8 ++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 5 +++++ 2 files changed, 13 insertions(+) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 843fb7e1356d..ffb61df5ec9e 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -751,9 +751,17 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( _, main = _packed_main_cache( device=torch.device("cpu"), blocks=2, layers=1, block_size=256, layer=0 ) + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() indexer.bind_main_kv_cache(main) assert indexer.indexer_op._num_q_heads == indexer._q_fp8.shape[1] == 32 + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() + indexer.unbind_main_kv_cache() + assert indexer._physical_selection_plan is None + assert indexer._index_cache is None + assert indexer._main_cache_num_blocks == 0 def test_glm53_packed_tail_accepts_nvfp4_main_record() -> None: diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index cc0311138364..3f15e9b2ce77 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -143,6 +143,7 @@ def __init__( self._expand_pooled_topk_to_physical_slots = ( b12x_sparse_mla.expand_pooled_topk_to_physical_slots ) + self._physical_selection_plan = None self.index_kpool_compress_ape = nn.Parameter( allocate_weights( torch.empty, @@ -360,6 +361,7 @@ def _index_cache_view( return index_cache, subpages_per_parent, parent_stride_pages def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: + self._physical_selection_plan = None index_cache, subpages, parent_stride_pages = self._index_cache_view(main_cache) block_size = int(main_cache.shape[1]) parent_table_width = self._max_parent_table_width( @@ -393,6 +395,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: def unbind_main_kv_cache(self) -> None: self._index_cache = None self._main_cache_num_blocks = 0 + self._physical_selection_plan = None def make_b12x_physical_selection_prepare_call( self, output: torch.Tensor, active_counts: torch.Tensor, *, state=None @@ -707,6 +710,8 @@ def forward( if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: if self._main_cache_num_blocks < 1: raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From 5c065850bd0eddee9b55f483033e2a34a1094b58 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:49:49 +0000 Subject: [PATCH 036/365] test(b12x): prepare native cache writers in graph-lifetime coverage Keep Karmic cache preparation contracts while checking that graph resources release caller activations. Preserve parser stream chunks as explicit concatenated values. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/parser/engine/test_deepseek_v4.py | 12 ++++++--- tests/v1/attention/test_b12x_v41_workspace.py | 25 ++++++++++--------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 9180a6f4b2d5..5eb369cee969 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -305,8 +305,10 @@ def test_streaming_split_next_parameter_tag_is_buffered( chunks = [ DSML_TOOL_START, f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n", - f"<{_PARAM_OPEN.format(name='location', is_str='true')}" - "Paris a<｜DSML｜parameter", + ( + f"<{_PARAM_OPEN.format(name='location', is_str='true')}" + "Paris a<｜DSML｜parameter" + ), ' name="date" string="true">tomorrow', _PARAM_CLOSE, DSML_INVOKE_END, @@ -376,8 +378,10 @@ def test_streaming_with_trailing_content(self, mock_tokenizer, mock_request): parser = DeepSeekV4Parser(mock_tokenizer) chunks = [ DSML_TOOL_START, - f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" - f"{_param('location', 'true', 'NYC')}\n", + ( + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" + f"{_param('location', 'true', 'NYC')}\n" + ), DSML_TOOL_END, "Done.", ] diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 260f386aa122..847536b72ae4 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -787,12 +787,13 @@ def metadata(page): device=device, dtype=torch.uint8, ) - from b12x.attention._shared.mla.compressed_reference import ( - pack_deepseek_v41_cache_reference, - ) - - cache[1:].copy_( - pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) + _write_cache( + kv, + cache, + torch.arange(length, device=device) + page, + page_size=page, + cache_kind=kind, + cache_format="deepseek_v41", ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -813,6 +814,12 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 + _write_index_keys( + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + page_size=main_page, + ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -836,12 +843,6 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) - attention.dsa_indexer.quantize_write_index_k_mxfp4( - layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - ) activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): From 3168d323d56579bf80bd92c9f0f6401c3b0010de Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:23:02 +0000 Subject: [PATCH 037/365] test(ds41): declare attention bounds in WO capture fixtures Prepare the output-projection fixture with the same declared attention envelope as the model. Both graph-lifetime cases pass on SM120 with the KK native wheel. Co-authored-by: OpenAI Codex --- tests/v1/attention/test_b12x_v41_workspace.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 847536b72ae4..c0c33fa5a341 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -1047,6 +1047,7 @@ def test_output_projection_capture_releases_caller_activations( ), ) layer.setup_wo_projection() + layer._declare_attention(device) workload = B12xWorkload( stage="weights", token_counts=(8, 16), From 8ff741487dfb79f359f81daee5e9f22f92444289 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:37:43 +0000 Subject: [PATCH 038/365] test(b12x): preserve Karmic preparation and precision contracts Observe caller-owned MHC outputs while requiring preparation not to retain them; retain allocation, pointer and graph replay assertions. Exercise the generic ModelOpt W4A16 method and padded native scales. Initialize the complete NVFP4 fixture. Use an IEEE FP32 reference instead of TF32 for unquantized projections; existing numerical tolerances are unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../kernels/test_b12x_linear.py | 52 ++++++++++++++----- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 6a1a540400e3..5a2e6d0a7db7 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -493,7 +493,8 @@ def pack(weight: torch.Tensor, weight_scale: torch.Tensor): assert len(calls) == 1 weight, weight_scale = calls[0] assert weight.shape == (48, 128) - assert weight_scale.shape == (48, 4) + # Preserve the checkpoint's padded row layout for native packed scales. + assert weight_scale.shape == (64, 4) assert weight.dtype == torch.float8_e4m3fn assert weight_scale.dtype == torch.uint8 assert layer.weight.numel() == 0 @@ -933,7 +934,14 @@ def test_b12x_fp4_processes_scale_and_preserves_loader( if kernel_cls is B12xNvFp4LinearKernel: layer.weight = torch.empty((48, 64), dtype=torch.uint8) layer.weight_global_scale = torch.tensor(0.5) - packed = object() + layer.input_global_scale_inv = torch.tensor(1.0) + + @dataclass(frozen=True) + class PackedWeight: + in_features: int + values: torch.Tensor + + packed = PackedWeight(0, layer.weight) def pack_weight(weight, scale, *, recipe, global_scale): assert weight.data_ptr() == layer.weight.data_ptr() @@ -948,13 +956,16 @@ def pack_weight(weight, scale, *, recipe, global_scale): lambda: types.SimpleNamespace(pack_weight=pack_weight), ) kernel = object.__new__(kernel_cls) + if kernel_cls is B12xNvFp4LinearKernel: + kernel.config = NvFp4LinearLayerConfig(use_a16=False) kernel.process_weights_after_loading(layer) assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() assert layer.weight_scale.weight_loader is weight_loader if kernel_cls is B12xNvFp4LinearKernel: - assert layer.b12x_nvfp4_packed_weight is packed + assert layer.b12x_nvfp4_packed_weight.in_features == 128 + assert layer.b12x_nvfp4_packed_weight.values is packed.values def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: @@ -1055,7 +1066,7 @@ def test_b12x_w4a16_modelopt_vision_width_preserves_bf16_activations(monkeypatch from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, - ModelOptNvFp4W4A16LinearMethod, + build_linear_method, ) from vllm.v1.worker.workspace import ( current_workspace_manager, @@ -1077,11 +1088,13 @@ def reject_activation_quantization(*args, **kwargs): n, k = 1152, 4304 config = VllmConfig(kernel_config=KernelConfig(linear_backend="b12x")) with set_current_vllm_config(config), torch.no_grad(): - method = ModelOptNvFp4W4A16LinearMethod( + method = build_linear_method( ModelOptNvFp4Config( quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True, - ) + ), + "W4A16_NVFP4", + "vision_tower.mlp", ) layer = torch.nn.Module() with torch.device(device): @@ -1932,10 +1945,13 @@ def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( name = f"_b12x_mhc_{native_operation}_impl" native = getattr(_impl, name) captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + captured_outputs = [] def observe(*args, **kwargs): captured_call[:] = [(args, kwargs)] - return native(*args, **kwargs) + result = native(*args, **kwargs) + captured_outputs[:] = outputs(result) + return result def outputs(value): return value if isinstance(value, tuple) else (value,) @@ -1948,7 +1964,10 @@ def outputs(value): with monkeypatch.context() as hooks: hooks.setattr(_impl, name, observe) call.produce() - first = outputs(call.invoke()) + assert call.invoke() is None + first = tuple(captured_outputs) + assert call.invoke() is None + second = tuple(captured_outputs) args, kwargs = captured_call.pop() reference_kwargs = { key: value @@ -1956,7 +1975,6 @@ def outputs(value): if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} } pointers = [tensor.data_ptr() for tensor in first] - second = outputs(call.invoke()) assert [tensor.data_ptr() for tensor in second] == pointers for actual, expected in zip( second, outputs(native(*args, **reference_kwargs)), strict=True @@ -1975,9 +1993,9 @@ def outputs(value): graph.replay() torch.accelerator.synchronize() assert torch.accelerator.memory_allocated(device) == allocated - assert [ - tensor.data_ptr() for tensor in outputs(call.output) - ] == pointers + # Preparation must not pin full-prefill results in the session. + assert call.output is None + assert [tensor.data_ptr() for tensor in first] == pointers expected = outputs(native(*args, **reference_kwargs)) for actual, wanted in zip(first, expected, strict=True): assert torch.isfinite(actual).all() @@ -2722,7 +2740,7 @@ def expected(rows): assert set(plans) == {4, capacity} actual = run(11) assert plans[11].prepared is not None - assert plans[11].selection.source == "fixed" + assert plans[11].selection.source == "default" torch.testing.assert_close(actual, expected(11), rtol=0.02, atol=0.125) session.freeze() with kernel_resolution_guard("FP8 prepared exact-M execution"): @@ -2746,11 +2764,17 @@ def expected(rows): @pytest.mark.parametrize("output_dtype", [torch.bfloat16, torch.float32]) @torch.no_grad() -def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay(output_dtype): +def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay( + output_dtype, request +): if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: pytest.skip("native b12x projection requires SM12x") from vllm.models.deepseek_v4_1.b12x_layers import B12xLinearMethod + # A TF32 reference rounds FP32 inputs, unlike the native SIMT projection. + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = torch.nn.Module() layer.weight = torch.nn.Parameter( From 120c01e3ca574449785cde8c60add3550f07cdc5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 22:04:32 +0000 Subject: [PATCH 039/365] test(ds41): model native FP8 prefill arithmetic in CED oracle Compare compact replay with an independent BF16 QK / FP8 PV reference and assert the complete causal candidate set. Retain numerical tolerances, poisoned discarded rows, padded outputs, caller scratch and CUDA graph checks. Production kernels and precision are unchanged. Validation: 333 native tests passed, three multi-GPU tests separately passed; all three CED geometries pass at unchanged tolerances. Reference adapted from B12X commit 4af78b86771d8491ec7b9d27168fdfbc1d67a773. Existing mypy and CUDA-API hook findings are unchanged. Co-authored-by: Luke Alonso Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_v41_workspace.py | 129 ++++++++++++++---- 1 file changed, 106 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index c0c33fa5a341..66ea9c7bd3e6 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -701,6 +701,73 @@ def test_mhc_fixed_capacity_buckets_preserve_decode_policy( assert bool(torch.isfinite(y).all()) +def _v41_fp8_operands(cache, kind): + """Decode native cache rows into E4M3 operands and per-64 scales. + + This independent operand reference follows B12X's + tests/_reference/v41_fp8.py. SWA pairs its E8M0 scales; indexed NVFP4 + uses a power-of-two scale that bounds four adjacent scale groups. + """ + if kind == "swa": + rows = cache.reshape(-1, 528) + codes = rows[:, :512].contiguous().view(torch.float8_e4m3fn).double() + original = rows[:, 512:528].double() + exponent = original.reshape(-1, 8, 2).amax(-1).clamp_min(1) - 127 + values = codes * torch.exp2(original - 127).repeat_interleave(32, dim=1) + else: + rows = cache.reshape(-1, 288) + packed = rows[:, :256].long() + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(1) + lut = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], + dtype=torch.float64, + device=cache.device, + ) + original = rows[:, 256:288].contiguous().view(torch.float8_e4m3fn).double() + values = lut[codes] * original.repeat_interleave(16, dim=1) + bound = original.reshape(-1, 8, 4).amax(-1) * (6 / 448) + exponent = torch.ceil(torch.log2(bound.clamp_min(2.0**-126))) + scales = torch.exp2(exponent) + codes = (values / scales.repeat_interleave(64, dim=1)).float() + return codes.to(torch.float8_e4m3fn).float(), scales.float() + + +def _v41_fp8_prefill_reference(q, values, scales, valid, sink): + """Model BF16 Q/K and per-64-key FP8 probability/value products. + + Prefill computes Q/K with BF16 operands. Each output group independently + quantizes scaled probabilities to E4M3 before its product with FP8 V; + partial numerators remain FP32 until the final LSE merge. + """ + keys = (values * scales.repeat_interleave(64, dim=-1)).bfloat16().float() + logits = torch.einsum("rhd,rkd->rhk", q.float(), keys) * 512**-0.5 + logits.masked_fill_(~valid[:, None], -torch.inf) + partials, lses = [], [] + for first in range(0, values.shape[1], 64): + local = logits[:, :, first : first + 64] + maximum = local.amax(-1) + p = torch.exp(local - maximum[:, :, None]) + p = torch.where(valid[:, None, first : first + 64], p, 0) + denominator = p.sum(-1) + groups = [] + for group in range(8): + weighted = p * scales[:, None, first : first + 64, group] + pscale = weighted.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 + pq = (weighted / pscale).to(torch.float8_e4m3fn).float() + vq = values[:, first : first + 64, group * 64 : (group + 1) * 64] + groups.append(torch.einsum("rhk,rkd->rhd", pq, vq) * pscale) + partials.append(torch.cat(groups, -1) / denominator.clamp_min(1e-30)[..., None]) + lses.append( + torch.where(denominator > 0, maximum + denominator.log(), -torch.inf) + ) + lse = torch.stack(lses, -1) + maximum = torch.maximum(lse.amax(-1), sink.float()[None]) + weights = torch.where(torch.isfinite(lse), torch.exp(lse - maximum[..., None]), 0) + denominator = weights.sum(-1) + torch.exp(sink.float()[None] - maximum) + output = (torch.stack(partials, -2) * weights[..., None]).sum(-2) + return (output / denominator.clamp_min(1e-30)[..., None]).bfloat16().float() + + @pytest.mark.parametrize("main_page,swa_page", [(64, 32), (128, 64), (256, 128)]) @pytest.mark.parametrize( "is_decode,rows,live_rows", @@ -1561,14 +1628,14 @@ def test_ced_compact_attention_bounded_oracle_and_frozen_replay( ): from contextlib import ExitStack - from b12x.attention._shared.mla.compressed_reference import ( - unpack_deepseek_v41_cache_reference, - ) from b12x.preparation import PreparationSession from vllm.utils.b12x import B12xWorkload attention, manager, workspace = native_workspace + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) torch.manual_seed(713) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = _layer(attention, 20, swa_page=swa_page) @@ -1633,17 +1700,19 @@ def metadata(page): ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache - swa_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + swa_values, swa_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) # NaN FP8 payloads in every old SWA row: reading below the # replay boundary cannot accidentally look like a valid zero page. cache[1 : 1 + boundary // page].fill_(0x7F) else: layer.kv_cache = cache - main_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + main_values, main_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) layer.indexer.k_cache.kv_cache = torch.zeros( ( length // main_page + 1, @@ -1683,21 +1752,35 @@ def run(): ) def oracle(live): - values = torch.cat((swa_values[boundary:], main_values)) - logical = torch.cat( + selected = layer.topk_indices_buffer[:live].long() + # Candidate order affects per-tile FP8 rounding. Verify its complete + # causal set independently before using that order in the numeric oracle. + columns = torch.arange(selected.shape[1], device=device)[None] + expected_ids = columns.expand_as(selected).clone() + expected_ids.masked_fill_(columns > positions[:live, None], length) + actual_ids = selected.masked_fill(selected < 0, length).sort(dim=1).values + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + swa_logical = torch.arange(boundary, length, device=device)[None] + swa_valid = swa_logical <= positions[:live, None] + main_valid = (selected >= 0) & (selected <= positions[:live, None]) + valid = torch.cat((swa_valid, main_valid), dim=1) + values = torch.cat( ( - torch.arange(boundary, length, device=device), - torch.arange(length, device=device), - ) - ) - logits = torch.einsum("rhd,kd->rhk", q[:live].float(), values) * 512**-0.5 - logits.masked_fill_( - logical[None, None] > positions[:live, None, None], -torch.inf - ) - logits = torch.cat( - (logits, layer.attn_sink[None, :, None].expand(live, -1, -1)), -1 + swa_values[None, boundary:].expand(live, -1, -1), + main_values[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + scales = torch.cat( + ( + swa_scales[None, boundary:].expand(live, -1, -1), + main_scales[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + return _v41_fp8_prefill_reference( + q[:live], values, scales, valid, layer.attn_sink ) - return torch.einsum("rhk,kd->rhd", logits.softmax(-1)[..., :-1], values) workload = B12xWorkload( stage="state", @@ -1744,9 +1827,9 @@ def oracle(live): visible.copy_((positions + 1).clamp_min(0).int()) starts[1] = live q.normal_() - expected = oracle(live) # Eager execution also catches live-count specialization leaks. run() + expected = oracle(live) torch.testing.assert_close( out[:live].float(), expected, rtol=0.04, atol=0.025 ) From faf2e328a8fba4fd5610e307f5689cf75140983e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:11:18 +0000 Subject: [PATCH 040/365] fix(dflash): initialize fused MXFP8 projection with generic ModelOpt Recognize MXFP8 through its QuantSpec weight key and give the fused context K/V projection an independent kernel lifecycle. Preserve input/output dtypes and bias metadata without changing checkpoint precision or draft sampling. Keep BF16 loading unchanged and adapt the draft-loader fixture to the shared pipeline-safe loader helper. Validation: 34 DFlash tests pass. Native B12X fused projection passes exact representable-input checks at 1/8/32 rows for quantized and A16 activations, including changed-input CUDA graph replay with zero additional allocated bytes. Applicable pre-commit hooks pass. Whole-model MXFP8 DFlash qualification remains pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/spec_decode/test_dflash2.py | 100 +++++++++++++++++++-- vllm/model_executor/models/qwen3_dflash.py | 45 ++++++---- 2 files changed, 120 insertions(+), 25 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 315f28235456..09dccc5054e9 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -14,6 +14,8 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): + from vllm.config import LoadConfig + draft_load_config = object() draft_model_config = SimpleNamespace(hf_config=SimpleNamespace()) speculative_config = SimpleNamespace( @@ -25,6 +27,7 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): vllm_config = SimpleNamespace( attention_config=SimpleNamespace(), cache_config=SimpleNamespace(), + load_config=LoadConfig(load_format="fastsafetensors"), speculative_config=speculative_config, ) loaded = SimpleNamespace(model=SimpleNamespace()) @@ -42,17 +45,16 @@ def fake_get_model(**kwargs): monkeypatch.setattr(dflash_utils, "replace", fake_replace) monkeypatch.setattr(dflash_utils, "get_model", fake_get_model) monkeypatch.setattr( - dflash_utils, - "get_pp_group", + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.compilation.backends.set_model_tag", - lambda _tag: nullcontext(), + "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group", + lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.model_executor.models.qwen3_dflash.dflash_target_rope_is_neox_style", - lambda _model: None, + "vllm.compilation.backends.set_model_tag", + lambda _tag: nullcontext(), ) monkeypatch.setattr( "vllm.model_executor.models.qwen3_dflash.dflash_has_any_non_causal", @@ -61,6 +63,8 @@ def fake_get_model(**kwargs): assert dflash_utils.load_dflash_model(SimpleNamespace(), vllm_config) is loaded assert captured["load_config"] is draft_load_config + assert captured["vllm_config"].load_config.load_format == "auto" + assert vllm_config.load_config.load_format == "fastsafetensors" def test_dflash_reset_attn_releases_cache_layout_state(): @@ -158,7 +162,7 @@ def test_dflash_context_projection_rejects_mixed_quantization(mxfp8_layer: int): from torch import nn from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, + build_linear_method, ) from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model @@ -166,8 +170,8 @@ class UnreadableWeight: def __getitem__(self, _key): raise AssertionError("weights must not be read before validation") - mxfp8_method = object.__new__(ModelOptMxFp8LinearMethod) - methods = [None, None] + mxfp8_method = build_linear_method(None, "MXFP8", "") + methods = [build_linear_method(None, "FP8", ""), None] methods[mxfp8_layer] = mxfp8_method layers_attn = [ SimpleNamespace( @@ -186,6 +190,84 @@ def __getitem__(self, _key): model._build_context_kv_buffers(layers_attn, has_bias=False) +@pytest.mark.parametrize("mxfp8", [False, True]) +@pytest.mark.parametrize("has_bias", [False, True]) +def test_fused_context_projection_owns_its_linear_method(monkeypatch, mxfp8, has_bias): + """Fused K/V packing must not reconfigure the query projection's kernel.""" + from torch import nn + + import vllm.model_executor.layers.quantization.modelopt as modelopt + import vllm.model_executor.parameter as parameter + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + kernels = [] + + def select_kernel(spec, layer, runtime_dtypes, **kwargs): + kernel = SimpleNamespace( + process_weights_after_loading=lambda layer: None, + input_quant_key=lambda: None, + ) + kernels.append((kernel, layer.output_size_per_partition)) + assert runtime_dtypes.input_dtype == torch.bfloat16 + assert runtime_dtypes.out_dtype == torch.bfloat16 + assert runtime_dtypes.marlin_input_dtype == torch.float16 + assert layer.has_bias is has_bias + return kernel + + monkeypatch.setattr(modelopt, "select_linear_kernel", select_kernel) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + method = modelopt.build_linear_method(None, "MXFP8", "") if mxfp8 else None + original_kernel = object() + if method is not None: + method.kernel = original_kernel + method.input_dtype = method.out_dtype = torch.bfloat16 + method.marlin_input_dtype = torch.float16 + layers = [] + for index in range(2): + projection = nn.Module() + projection.quant_method = method + weight = torch.full((192, 64), index + 1, dtype=torch.bfloat16) + if mxfp8: + weight = weight.to(torch.float8_e4m3fn) + projection.weight_scale = nn.Parameter( + torch.full((192, 2), 127, dtype=torch.uint8), requires_grad=False + ) + projection.weight = nn.Parameter(weight, requires_grad=False) + if has_bias: + projection.bias = nn.Parameter(torch.full((192,), float(index + 1))) + attention = SimpleNamespace( + qkv_proj=projection, + q_size=128, + k_norm=SimpleNamespace(weight=torch.ones(32)), + ) + layers.append(SimpleNamespace(self_attn=attention)) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.layers = layers + model.hidden_norm = SimpleNamespace(weight=torch.ones(64, dtype=torch.bfloat16)) + model._fused_kv_linear = nn.Module() + model._fused_kv_quant_method = None + model._build_context_kv_buffers([layer.self_attn for layer in layers], has_bias) + expected = model._fused_kv_weight.clone() + model.process_weights_after_loading() + if mxfp8: + assert method is not None + assert model._fused_kv_quant_method is not method + assert method.kernel is original_kernel + assert kernels == [(model._fused_kv_quant_method.kernel, 128)] + assert model._fused_kv_linear.weight_block_size == [1, 32] + torch.testing.assert_close( + model._fused_kv_linear.weight.float(), expected.float() + ) + assert model._fused_kv_weight is None + assert model._fused_kv_weight_scale is None + else: + assert not kernels + assert model._fused_kv_quant_method is None + torch.testing.assert_close(model._fused_kv_weight, expected) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3c7c301814ff..0ea09746a4bb 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -27,6 +27,8 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -511,13 +513,11 @@ def _build_context_kv_buffers( layers_attn: list[nn.Module], has_bias: bool, ) -> None: - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_methods = [a.qkv_proj.quant_method for a in layers_attn] uses_mxfp8 = [ - isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods + isinstance(method, ModelOptLinearMethod) + and method.spec.weight == kMxfp8Static + for method in quant_methods ] if any(uses_mxfp8) and not all(uses_mxfp8): raise ValueError( @@ -593,12 +593,11 @@ def _build_fused_kv_buffers(self) -> None: def process_weights_after_loading(self) -> None: """Pack the serialized MXFP8 context projection for its GEMM backend.""" - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_method = self.layers[0].self_attn.qkv_proj.quant_method - if not isinstance(quant_method, ModelOptMxFp8LinearMethod): + if not ( + isinstance(quant_method, ModelOptLinearMethod) + and quant_method.spec.weight == kMxfp8Static + ): return if self._fused_kv_weight is None or self._fused_kv_weight_scale is None: raise RuntimeError( @@ -607,9 +606,23 @@ def process_weights_after_loading(self) -> None: ) output_size, input_size = self._fused_kv_weight.shape - self._fused_kv_linear.input_size_per_partition = input_size - self._fused_kv_linear.output_size_per_partition = output_size - self._fused_kv_linear.logical_widths = [output_size] + # Fused context K/V has an independent weight shape and kernel + # lifecycle; do not repack through the query projection's method. + fused_method = ModelOptLinearMethod( + quant_method.spec, quant_method.ctx, quant_method.fmt + ) + fused_method.input_dtype = quant_method.input_dtype + fused_method.out_dtype = quant_method.out_dtype + fused_method.marlin_input_dtype = quant_method.marlin_input_dtype + self._fused_kv_linear.has_bias = self._fused_kv_bias is not None + fused_method.create_weights( + self._fused_kv_linear, + input_size_per_partition=input_size, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=self.hidden_norm.weight.dtype, + ) self._fused_kv_linear.register_parameter( "weight", nn.Parameter(self._fused_kv_weight, requires_grad=False) ) @@ -617,13 +630,13 @@ def process_weights_after_loading(self) -> None: "weight_scale", nn.Parameter(self._fused_kv_weight_scale, requires_grad=False), ) - quant_method.process_weights_after_loading(self._fused_kv_linear) - self._fused_kv_quant_method = quant_method + fused_method.process_weights_after_loading(self._fused_kv_linear) + self._fused_kv_quant_method = fused_method self._fused_kv_weight = None self._fused_kv_weight_scale = None logger.info_once( "Using %s for the fused DFlash context K/V projection.", - type(quant_method.kernel).__name__, + type(fused_method.kernel).__name__, ) def _project_context_kv( From f1ab68fd149ccafec001ac0c078367f0faa58fa3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:21:42 +0000 Subject: [PATCH 041/365] fix(v2): register native sampling warmup contracts Register the existing top-k/top-p compile-key declarations on sampling ranks. Seeded and processed-logprob requests can select native filtering even when unseeded warmup uses FlashInfer. This covers singleton split-row kernels omitted by the bounded 8/16-row warmup; sampling arithmetic and the production JIT monitor default are unchanged. Validation: six GPU warmup cases pass for vocabularies 257/129280 and row limits 1/7/128. Every reachable row count and top-k/top-p subset runs with a hook that rejects post-warmup compilation. Applicable pre-commit hooks pass. The DS4.1 strict-monitor reproducer requires a rebuilt-image rerun. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/sample/test_topk_topp_sampler.py | 24 +++++++++++++++++++---- vllm/v1/worker/gpu/model_runner.py | 5 +++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 4d4740524de0..cbf1d7655c9e 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from torch import Generator @@ -13,6 +15,7 @@ apply_top_k_top_p_probs, apply_top_k_top_p_pytorch, random_sample, + register_top_k_top_p_warmups, warmup_top_k_top_p, ) from vllm.v1.sample.sampler import Sampler @@ -74,21 +77,34 @@ def test_apply_top_k_top_p_probs_matches_processed_logits( reason="requires GPU Triton compilation", ) @pytest.mark.parametrize("vocab_size", [257, 129280]) +@pytest.mark.parametrize("max_num_rows", [1, 7, 128]) def test_filter_warmup_covers_speculative_rows_and_constraint_subsets( - monkeypatch: pytest.MonkeyPatch, vocab_size: int + monkeypatch: pytest.MonkeyPatch, vocab_size: int, max_num_rows: int ): - """Four K5 requests must not compile top-p-only filtering after warmup.""" + """Seeded filtering must cover singleton, split-row and large-batch keys.""" from triton import knobs + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry + device = torch.device(DEVICE_TYPE) - warmup_top_k_top_p(vocab_size, 24, device) + registry = JitWarmupRegistry( + SimpleNamespace( + model_config=SimpleNamespace(get_vocab_size=lambda: vocab_size), + scheduler_config=SimpleNamespace(max_num_seqs=max_num_rows), + num_speculative_tokens=0, + ) + ) + with registry.activate(): + register_top_k_top_p_warmups() + registry.warmup() + warmup_top_k_top_p(vocab_size, max_num_rows, device) torch.accelerator.synchronize() def unexpected_compile(**kwargs): pytest.fail("Filtering compiled a kernel after bounded sampler warmup") monkeypatch.setattr(knobs.runtime, "jit_post_compile_hook", unexpected_compile) - for num_rows in (8, 9, 15, 16, 17, 24): + for num_rows in range(1, max_num_rows + 1): logits = torch.zeros(num_rows, vocab_size, device=device) top_k = torch.full((num_rows,), 20, dtype=torch.int32, device=device) top_p = torch.full((num_rows,), 0.95, device=device) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index aba519befaff..5b57f9e55146 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,6 +78,7 @@ ModelRunnerOutput, RoutedExpertsTensors, ) +from vllm.v1.sample.ops.topk_topp_sampler import register_top_k_top_p_warmups from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.watermarking.spec_decode import ( @@ -472,6 +473,10 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: + # Seeded and processed-logprob requests use native filtering even + # when unseeded warmup requests select FlashInfer sampling. + with self.jit_warmup_registry.activate(): + register_top_k_top_p_warmups() sampler_kwargs: dict[str, Any] = { "max_num_reqs": self.max_num_reqs, "vocab_size": self.vocab_size, From 8dd07c4c8bf40f8234a4470061a3767e8818d41c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:41:37 +0000 Subject: [PATCH 042/365] fix(chat): normalize namespaced tool iterables before matching Materialize Python tool collections before qualifying namespace metadata and validating named choices. Lists, tuples and generators retain the same serialized function identities without mutating definitions. JSON request behavior is unchanged. Validation: tuple inputs reproduce a false tool-choice mismatch before the fix. The namespace/history suite passes for lists, tuples and generators; combined CPU suites pass 90 tests with 43 GPU-only skips. Applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/entrypoints/openai/test_tool_calls_serialization.py | 7 +++++-- vllm/entrypoints/openai/chat_completion/protocol.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 156d6bb140c4..c62bff2d1ab6 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -160,8 +160,9 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): @pytest.mark.parametrize("location", ["tool", "function"]) @pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) @pytest.mark.parametrize("iterator", [False, True]) +@pytest.mark.parametrize("tools_collection", [list, tuple, iter]) def test_namespaces_survive_definitions_choices_and_history( - location, namespace, iterator + location, namespace, iterator, tools_collection ): def qualify(item, value): target = item if location == "tool" else item["function"] @@ -185,6 +186,7 @@ def qualify(item, value): ], } original = copy.deepcopy(payload) + payload["tools"] = tools_collection(tools) if iterator: payload["messages"][1]["tool_calls"] = iter(calls) request = ChatCompletionRequest.model_validate(payload) @@ -200,8 +202,9 @@ def qualify(item, value): assert request.model_dump_json() == serialized restored = ChatCompletionRequest.model_validate_json(serialized) assert restored.model_dump_json() == serialized - if not iterator: + if not iterator and tools_collection is list: assert payload == original + assert tools == original["tools"] assert calls == original["messages"][1]["tool_calls"] diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index e5913efef6c6..3607ec221b6a 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -941,8 +941,11 @@ def check_tool_usage(cls, data): # Named choices must be compared with qualified identities before # Pydantic serializes function definitions and discards extra fields. data = dict(data) - if isinstance(data.get("tools"), list): - data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + tools = data.get("tools") + if isinstance(tools, Iterable) and not isinstance( + tools, (str, bytes, bytearray, dict) + ): + data["tools"] = [normalize_tool_namespace(tool) for tool in tools] if isinstance(data.get("tool_choice"), dict): data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) From c082782f34606bca03852d2f39555320a5d1897d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:42:16 +0000 Subject: [PATCH 043/365] fix(glm): reject undeclared selection before recurrent cache updates Check physical-selection resources before dispatching pooled-cache writes on pure DCP1 decode. An unbound main cache or missing selection plan now fails without advancing request state. Supported selection arithmetic and the hot-path checks are unchanged. Validation: both missing-resource cases reproduce an early cache-write attempt before the fix and pass afterwards. All 67 pooled-indexer tests pass on stock SM120, including changed-input replay and prepared selection; applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 60 +++++++++++++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 17 ++++-- 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index ffb61df5ec9e..5b6d3ca555bc 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -689,6 +689,66 @@ def test_glm53_physical_selection_provider_is_explicit() -> None: ) +@pytest.mark.parametrize( + "main_blocks, error", + [(0, "main cache is not bound"), (1, "physical selection has not been declared")], +) +def test_glm53_missing_physical_plan_does_not_mutate_recurrent_cache( + monkeypatch, main_blocks, error +) -> None: + from vllm.models.glm5next.nvidia import pooled_indexer as module + + hidden = torch.zeros((1, 128)) + metadata = SimpleNamespace( + num_reqs=1, + num_actual_tokens=1, + num_decode_tokens=1, + num_decodes=1, + max_query_len=1, + block_table=torch.zeros((1, 1), dtype=torch.int32), + selector_state_slot_ids=torch.zeros(1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + indexer = SimpleNamespace( + max_tokens=1, + wq_b=lambda value: (torch.zeros((1, 32 * 128)), None), + wk=lambda value: (hidden, None), + k_norm=lambda value: value, + index_kpool_compress_gate=torch.zeros((128, 128)), + index_kpool_compress_ape=torch.zeros((4, 128)), + _project_head_weights=lambda value: torch.zeros((1, 32)), + _q_fp8=torch.empty((1, 32, 128), dtype=torch.float8_e4m3fn), + _q_scale=torch.empty((1, 32)), + main_layer_name="attention", + _index_cache=torch.zeros((1, 64, 132), dtype=torch.uint8), + _tail=torch.zeros((1, 2, 4, 128)), + _state_slots=Glm5NextPooledIndexer._state_slots, + _parent_table_width=1, + _parent_stride_pages=1, + block_size=256, + dcp_world_size=1, + _emit_physical_selection=True, + _main_cache_num_blocks=main_blocks, + _physical_selection_plan=None, + ) + monkeypatch.setattr(module, "fwht128_quant_fp8", lambda *args, **kwargs: None) + monkeypatch.setattr( + module, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={"attention": metadata}), + ) + + def reject_cache_update(*args, **kwargs): + pytest.fail("Missing physical-selection resources must not update the cache") + + monkeypatch.setattr(module, "update_decode_pools", reject_cache_update) + with pytest.raises(RuntimeError, match=error): + Glm5NextPooledIndexer.forward( + indexer, hidden, hidden, torch.zeros(1, dtype=torch.int64), None + ) + + def _packed_main_cache( *, device: torch.device, diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 3f15e9b2ce77..63ee0d63546f 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -568,6 +568,16 @@ def forward( f"actual={actual_table_width}, required={self._parent_table_width}" ) + decode_only = decode_rows == live_rows + emit_physical_selection = ( + decode_only and self.dcp_world_size == 1 and self._emit_physical_selection + ) + if emit_physical_selection: + if self._main_cache_num_blocks < 1: + raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") + update_decode_pools( index_cache, self._tail, @@ -586,7 +596,6 @@ def forward( ) parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] seq_lens = self._pool_seq_lens[:live_rows] - decode_only = decode_rows == live_rows decode_table = self._decode_block_table[:decode_rows] if decode_only: prepare_c4_decode_metadata( @@ -707,11 +716,7 @@ def forward( output = self.topk_indices_buffer[:rows] if live_rows < rows: output[live_rows:].fill_(-1) - if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: - if self._main_cache_num_blocks < 1: - raise RuntimeError("GLM selector main cache is not bound") - if self._physical_selection_plan is None: - raise RuntimeError("GLM physical selection has not been declared") + if emit_physical_selection: self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From e754e457c35b4bb9b56e8b87188ac40b099dd2e7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:52:25 +0000 Subject: [PATCH 044/365] Preserve qualified tool identities in tuple message histories Normalize tuple histories accepted by the Python request API before Pydantic union validation. Copy message dictionaries so caller-owned histories remain unchanged. JSON/list requests retain their existing behavior. Invalid scalar and mapping collections remain rejected. Validation: the tuple reproducer loses inventory::lookup before this fix. All 72 tool-call serialization tests and applicable pre-commit hooks pass with the fix; no model math or sampling behavior changes. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../openai/test_tool_calls_serialization.py | 30 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 9 +++++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index c62bff2d1ab6..ec223ed2f70d 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -229,6 +229,36 @@ def test_namespace_description_is_preserved_exactly_once(): ).model_dump() == (parsed.model_dump()) +@pytest.mark.parametrize("collection", [list, tuple]) +def test_iterable_history_preserves_namespaced_tool_identity(collection): + call = _make_tool_call("call_history", "lookup", "{}") + call["namespace"] = "inventory" + messages = [{"role": "assistant", "content": None, "tool_calls": [call]}] + original = copy.deepcopy(messages) + request = ChatCompletionRequest.model_validate( + {"model": "test-model", "messages": collection(messages)} + ) + assert request.messages[0]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + serialized = request.model_dump_json() + restored = ChatCompletionRequest.model_validate_json(serialized) + assert restored.messages[0]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + assert request.model_dump_json() == serialized + assert messages == original + + +@pytest.mark.parametrize("messages", ["text", b"text", bytearray(b"text"), {}]) +def test_invalid_history_collection_remains_a_client_error(messages): + with pytest.raises(ValidationError) as error: + ChatCompletionRequest.model_validate( + {"model": "test-model", "messages": messages} + ) + assert any(item["loc"] == ("messages",) for item in error.value.errors()) + + @pytest.mark.parametrize("tool_calls", [1, False, 1.5]) def test_invalid_tool_calls_shape_reports_the_field(tool_calls): with pytest.raises(ValidationError) as error: diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 3607ec221b6a..89ecef017748 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -552,11 +552,18 @@ def _normalize_messages_before(cls, data: Any) -> Any: - Renames the deprecated ``reasoning_content`` field to ``reasoning`` so downstream code only needs to check one field. - Qualifies tool-call namespaces before union validation can drop them. + + Args: + data: Unvalidated request payload, including Python message sequences. + + Returns: + A copied payload with normalized list/tuple message histories, or + the unchanged input when it is not a supported payload shape. """ if not isinstance(data, dict): return data messages = data.get("messages") - if not isinstance(messages, list): + if not isinstance(messages, (list, tuple)): return data data = dict(data) messages = [dict(msg) if isinstance(msg, dict) else msg for msg in messages] From 3e9d78044375b2bc078d9a864cb5d5e5a9bd01dc Mon Sep 17 00:00:00 2001 From: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:10:50 -0700 Subject: [PATCH 045/365] Cut transient memory in NVFP4 Marlin scale-factor computation `_nvfp4_compute_scale_factor` materialized an FP32 copy of the whole scale tensor, a boolean mask, and a gathered copy of the positive entries just to find the largest scale. For fused MoE parameters the helper runs on the scales of every expert at once, so on a 288-expert model (GLM-5.3-Flash NVFP4) one call touched about 300M elements and allocated several GiB of transient device memory during weight loading. On unified-memory machines with little headroom at load time (DGX Spark GB10, see #50925) that spike is enough to fail the load. The factor only depends on the largest scale, so reduce the input tensor directly. The result is identical: a max commutes with the exact multiplication by 2^7, non-positive scales never win the max when a positive one exists, and an all-non-positive tensor still yields 1.0. An empty tensor also yields 1.0 as before. A NaN scale now raises instead of being silently skipped. Measured on CPU with a bf16 tensor of 37.7M elements (one eighth of the production fused w13 scales): peak RSS growth 1011 MiB before, 3 MiB after, same factor (32768.0). The same shape on CUDA is covered by the new test, which asserts no tensor-sized temporaries. Signed-off-by: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Co-Authored-By: Claude Fable 5.1 (cherry picked from commit 941bfa989ca050aba3627e455ab8bd70e45a49ed) --- .../test_nvfp4_marlin_scale_factor.py | 100 ++++++++++++++++++ .../quantization/utils/marlin_utils_fp4.py | 20 ++-- 2 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py new file mode 100644 index 000000000000..f99eb06317b6 --- /dev/null +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the NVFP4 Marlin power-of-2 scale factor helper. + +The helper must return the same factor as the previous mask-and-gather +implementation while allocating no tensor-sized temporaries. +""" + +import pytest +import torch + +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + _nvfp4_compute_scale_factor, +) + +DTYPES = [torch.float16, torch.bfloat16] +SHAPES = [(64,), (128, 256), (8, 64, 128)] +UPPER = 448 * (2**7) + + +def _reference_scale_factor( + marlin_scales: torch.Tensor, a_dtype: torch.dtype | None = None +) -> float: + """The previous implementation, kept verbatim as the oracle.""" + if a_dtype is not None and a_dtype == torch.half: + return 1.0 + ws_float = marlin_scales.float() * (2**7) + nonzero_mask = ws_float > 0 + if nonzero_mask.any(): + max_val = ws_float[nonzero_mask].max() + if max_val < UPPER: + sf = (UPPER / max_val).log2().floor().exp2() + return sf.item() + return 1.0 + + +def _cases(shape: tuple[int, ...], dtype: torch.dtype) -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(0) + small = torch.rand(shape, generator=gen) * 0.01 + with_zeros = small.clone() + with_zeros.flatten()[::3] = 0.0 + with_negatives = small.clone() + with_negatives.flatten()[::5] *= -1.0 + normalized = torch.rand(shape, generator=gen) * 400.0 + 48.0 + return { + "small": small.to(dtype), + "with_zeros": with_zeros.to(dtype), + "with_negatives": with_negatives.to(dtype), + "all_zeros": torch.zeros(shape, dtype=dtype), + "all_negative": (-small - 0.001).to(dtype), + "already_normalized": normalized.to(dtype), + "single_large": torch.full(shape, 0.5, dtype=dtype), + } + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("shape", SHAPES) +@pytest.mark.parametrize("a_dtype", [None, torch.float16, torch.bfloat16]) +def test_scale_factor_matches_reference( + dtype: torch.dtype, shape: tuple[int, ...], a_dtype: torch.dtype | None +): + for name, scales in _cases(shape, dtype).items(): + expected = _reference_scale_factor(scales, a_dtype) + actual = _nvfp4_compute_scale_factor(scales, a_dtype) + assert actual == expected, name + assert actual >= 1.0 + assert actual == 2.0 ** round(torch.tensor(actual).log2().item()) + + +def test_scale_factor_rescales_into_range(): + scales = torch.full((32, 64), 2**-10, dtype=torch.bfloat16) + sf = _nvfp4_compute_scale_factor(scales) + rescaled = scales.float().max() * (2**7) * sf + assert 2.0 <= rescaled < UPPER + + +def test_scale_factor_empty_tensor(): + assert _nvfp4_compute_scale_factor(torch.empty(0, dtype=torch.bfloat16)) == 1.0 + + +def test_scale_factor_rejects_nan(): + scales = torch.rand(16, 16, dtype=torch.bfloat16) + scales[3, 3] = float("nan") + with pytest.raises(ValueError, match="NaN"): + _nvfp4_compute_scale_factor(scales) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_scale_factor_allocates_no_tensor_temporaries(): + device = torch.device("cuda") + scales = torch.rand((288, 512, 256), dtype=torch.bfloat16, device=device) * 0.01 + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + baseline = torch.cuda.memory_allocated(device) + sf = _nvfp4_compute_scale_factor(scales) + torch.cuda.synchronize(device) + transient = torch.cuda.max_memory_allocated(device) - baseline + assert sf > 1.0 + # The previous implementation needed ~4.5x the tensor size here. + assert transient < scales.numel() * scales.element_size() // 64 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 3b6990101ed7..1e5ad5660186 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -47,14 +47,20 @@ def _nvfp4_compute_scale_factor( # no rescaling is applied here if active dtype is half. if a_dtype is not None and a_dtype == torch.half: return 1.0 + if marlin_scales.numel() == 0: + return 1.0 - ws_float = marlin_scales.float() * (2**7) - nonzero_mask = ws_float > 0 - if nonzero_mask.any(): - max_val = ws_float[nonzero_mask].max() - if max_val < 448 * (2**7): - sf = (448 * (2**7) / max_val).log2().floor().exp2() - return sf.item() + # The factor depends only on the largest scale, so reduce the input + # directly. Materializing an FP32 copy, a boolean mask, and a gathered + # copy of every scale costs up to 9x the tensor size in transient device + # memory; for fused MoE parameters (all experts at once) that is several + # GiB during weight loading. + max_val = marlin_scales.max().float() * (2**7) + if torch.isnan(max_val): + raise ValueError("NVFP4 Marlin scales contain NaN") + if 0 < max_val < 448 * (2**7): + sf = (448 * (2**7) / max_val).log2().floor().exp2() + return sf.item() return 1.0 From 9a7206c69fc09bbb1b6e78862d5613b0afc150ac Mon Sep 17 00:00:00 2001 From: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:38:18 -0700 Subject: [PATCH 046/365] Shorten the NVFP4 Marlin scale-factor comment Keep the helper comment to two lines, per AGENTS.md guidance on brevity. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> (cherry picked from commit 78722557d8eeab820555e04147256d847b0a8756) --- .../layers/quantization/utils/marlin_utils_fp4.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 1e5ad5660186..90d833bbc0c3 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -50,11 +50,8 @@ def _nvfp4_compute_scale_factor( if marlin_scales.numel() == 0: return 1.0 - # The factor depends only on the largest scale, so reduce the input - # directly. Materializing an FP32 copy, a boolean mask, and a gathered - # copy of every scale costs up to 9x the tensor size in transient device - # memory; for fused MoE parameters (all experts at once) that is several - # GiB during weight loading. + # Reduce the input directly: the FP32 copy, mask, and gathered copy of the + # old path cost several times the tensor size in transient memory. max_val = marlin_scales.max().float() * (2**7) if torch.isnan(max_val): raise ValueError("NVFP4 Marlin scales contain NaN") From 25fec51de35363e7f565fec39237dc3db9adc59b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:31:00 +0000 Subject: [PATCH 047/365] test(marlin): qualify bounded NVFP4 scale reduction on CUDA Cover finite E4M3 scale values in contiguous and strided GPU tensors, preserve exact scale-factor parity, and record scalar-reduction peak memory. Use portable accelerator memory APIs. The imported upstream implementation and author attribution remain intact. Validation: 26 native tests pass on RTX PRO 6000; mask-and-gather reference uses 1,243,946,496 transient bytes versus 2,560 for the reduction on 37,748,736 BF16 scales. Applicable pre-commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../test_nvfp4_marlin_scale_factor.py | 36 ++++++++++++++----- .../quantization/utils/marlin_utils_fp4.py | 4 +-- 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py index f99eb06317b6..f2521c35f518 100644 --- a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the NVFP4 Marlin power-of-2 scale factor helper. -The helper must return the same factor as the previous mask-and-gather +The helper must return the same factor as the mask-and-gather implementation while allocating no tensor-sized temporaries. """ @@ -21,7 +21,7 @@ def _reference_scale_factor( marlin_scales: torch.Tensor, a_dtype: torch.dtype | None = None ) -> float: - """The previous implementation, kept verbatim as the oracle.""" + """Reference reduction over positive scales in FP32.""" if a_dtype is not None and a_dtype == torch.half: return 1.0 ws_float = marlin_scales.float() * (2**7) @@ -86,15 +86,33 @@ def test_scale_factor_rejects_nan(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_scale_factor_allocates_no_tensor_temporaries(): +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_scale_factor_matches_fp8_scale_domain_on_cuda(dtype, contiguous): + """All nonnegative finite E4M3 scales retain their factor after reduction.""" + values = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn) + scales = values.to(dtype=dtype, device="cuda").repeat(4, 1).T + if contiguous: + scales = scales.contiguous() + for multiplier in (1.0, 2**-7, 2**-10): + sample = scales * multiplier + assert _nvfp4_compute_scale_factor(sample, torch.bfloat16) == ( + _reference_scale_factor(sample, torch.bfloat16) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_scale_factor_allocates_no_tensor_temporaries(record_property): device = torch.device("cuda") scales = torch.rand((288, 512, 256), dtype=torch.bfloat16, device=device) * 0.01 - torch.cuda.synchronize(device) - torch.cuda.reset_peak_memory_stats(device) - baseline = torch.cuda.memory_allocated(device) + torch.accelerator.synchronize(device) + torch.accelerator.reset_peak_memory_stats(device) + baseline = torch.accelerator.memory_allocated(device) sf = _nvfp4_compute_scale_factor(scales) - torch.cuda.synchronize(device) - transient = torch.cuda.max_memory_allocated(device) - baseline + torch.accelerator.synchronize(device) + transient = torch.accelerator.max_memory_allocated(device) - baseline + record_property("transient_bytes", transient) + record_property("input_bytes", scales.numel() * scales.element_size()) assert sf > 1.0 - # The previous implementation needed ~4.5x the tensor size here. + # Scalar reduction must not allocate a scale-tensor-sized temporary. assert transient < scales.numel() * scales.element_size() // 64 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 90d833bbc0c3..863f4a6d6ac7 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -50,8 +50,8 @@ def _nvfp4_compute_scale_factor( if marlin_scales.numel() == 0: return 1.0 - # Reduce the input directly: the FP32 copy, mask, and gathered copy of the - # old path cost several times the tensor size in transient memory. + # Reduce before conversion to keep temporary storage independent of the + # number of experts. Power-of-two multiplication preserves the maximum. max_val = marlin_scales.max().float() * (2**7) if torch.isnan(max_val): raise ValueError("NVFP4 Marlin scales contain NaN") From 4a8904910ab41785dc66353924a3dcb6fa5a4b32 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:52:24 +0000 Subject: [PATCH 048/365] Bound NVFP4 Marlin expert-scale preparation memory Reduce one expert at a time to select the shared scale factor, then write each converted expert directly into the output bank. Avoid retaining a full half-precision bank and concatenating a second full output. Finite scale bytes, shared global factors, input ownership, gated and ungated padding remain unchanged; inference kernels are untouched. Validation: 35 CPU/CUDA tests pass on SM120, including bitwise whole-bank reference parity and a memory regression. For 27 MiB output scales, peak additional memory is 28,775,936 bytes versus 75,503,616 bytes before this change. All applicable pre-commit hooks pass. Whole-model TP2 qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../test_nvfp4_marlin_scale_factor.py | 137 ++++++++++++++++++ .../quantization/utils/marlin_utils_fp4.py | 42 ++++-- 2 files changed, 164 insertions(+), 15 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py index f2521c35f518..a1fa6726c07b 100644 --- a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -6,11 +6,20 @@ implementation while allocating no tensor-sized temporaries. """ +from types import SimpleNamespace + import pytest import torch +from vllm.model_executor.layers.quantization.utils import marlin_utils_fp4 +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + marlin_permute_scales, +) from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( _nvfp4_compute_scale_factor, + nvfp4_marlin_process_global_scale, + nvfp4_marlin_process_scales, + prepare_nvfp4_moe_layer_for_marlin, ) DTYPES = [torch.float16, torch.bfloat16] @@ -116,3 +125,131 @@ def test_scale_factor_allocates_no_tensor_temporaries(record_property): assert sf > 1.0 # Scalar reduction must not allocate a scale-tensor-sized temporary. assert transient < scales.numel() * scales.element_size() // 64 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("intermediate", [64, 96]) +@pytest.mark.parametrize("gated", [False, True]) +def test_moe_scales_match_whole_bank_reference(dtype, intermediate, gated): + """Per-expert conversion preserves scale bytes, padding and global factors.""" + experts, hidden = 4, 128 + shards = 2 if gated else 1 + padded = (intermediate + 63) // 64 * 64 + layer = SimpleNamespace( + num_experts=experts, + hidden_size=hidden, + intermediate_size_per_partition=intermediate, + params_dtype=dtype, + ) + + def scales(shape): + count = 1 + for extent in shape: + count *= extent + # Include zero and finite positive E4M3 scales at several exponents. + values = ( + (torch.arange(count, device="cuda") % 120) + .to(torch.uint8) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + factors = torch.tensor([0.125, 0.25, 0.5, 1.0], device="cuda") + return (values.float() * factors[:, None, None]).to(torch.float8_e4m3fn) + + s13 = scales((experts, shards * intermediate, hidden // 16)) + s2 = scales((experts, hidden, intermediate // 16)) + g = torch.arange(1, experts + 1, dtype=torch.float32, device="cuda") + inputs = (s13.view(torch.uint8).clone(), s2.view(torch.uint8).clone()) + result = prepare_nvfp4_moe_layer_for_marlin( + layer, + torch.zeros( + (experts, shards * intermediate, hidden // 2), + dtype=torch.uint8, + device="cuda", + ), + s13, + g, + torch.zeros( + (experts, hidden, intermediate // 2), dtype=torch.uint8, device="cuda" + ), + s2, + g, + gated, + ) + for source, actual, actual_global, is_gate in ( + (s13, result[1], result[2], True), + (s2, result[4], result[5], False), + ): + source = source.to(dtype) + if is_gate: + source = source.reshape(experts, shards, intermediate, -1) + source = torch.nn.functional.pad(source, (0, 0, 0, padded - intermediate)) + source = source.reshape(experts, shards * padded, -1) + size_k, size_n = hidden, shards * padded + else: + source = torch.nn.functional.pad(source, (0, (padded - intermediate) // 16)) + size_k, size_n = padded, hidden + factor = _reference_scale_factor(source, dtype) + expected = [] + for expert in source: + permuted = marlin_permute_scales(expert.T, size_k, size_n, 16) + converted, _ = nvfp4_marlin_process_scales(permuted, factor, dtype) + expected.append(converted) + torch.testing.assert_close( + actual.view(torch.uint8), + torch.stack(expected).view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual_global, + nvfp4_marlin_process_global_scale(g, dtype) / factor, + rtol=0, + atol=0, + ) + torch.testing.assert_close(s13.view(torch.uint8), inputs[0], rtol=0, atol=0) + torch.testing.assert_close(s2.view(torch.uint8), inputs[1], rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_moe_scale_conversion_avoids_bank_temporaries(monkeypatch, record_property): + """Additional storage contains output scales and one expert's conversion.""" + experts, hidden, intermediate = 288, 1024, 512 + layer = SimpleNamespace( + num_experts=experts, + hidden_size=hidden, + intermediate_size_per_partition=intermediate, + params_dtype=torch.bfloat16, + ) + # Weight packing is independently tested; isolate scale preparation's peak. + monkeypatch.setattr(marlin_utils_fp4, "_repack_marlin_experts", lambda w, *a: w) + weight = torch.zeros((), dtype=torch.uint8, device="cuda") + s13 = torch.ones( + (experts, 2 * intermediate, hidden // 16), + dtype=torch.float8_e4m3fn, + device="cuda", + ) + s2 = torch.ones( + (experts, hidden, intermediate // 16), dtype=torch.float8_e4m3fn, device="cuda" + ) + global_scale = torch.ones(experts, dtype=torch.float32, device="cuda") + torch.accelerator.synchronize() + torch.accelerator.reset_peak_memory_stats() + before = torch.accelerator.memory_allocated() + result = prepare_nvfp4_moe_layer_for_marlin( + layer, + weight.expand(experts, 2 * intermediate, hidden // 2), + s13, + global_scale, + weight.expand(experts, hidden, intermediate // 2), + s2, + global_scale, + True, + ) + torch.accelerator.synchronize() + extra = torch.accelerator.max_memory_allocated() - before + output_bytes = sum(t.numel() * t.element_size() for t in (result[1], result[4])) + record_property("peak_extra_bytes", extra) + record_property("output_scale_bytes", output_bytes) + assert extra < output_bytes + 4 * 1024 * 1024 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 863f4a6d6ac7..a94823434cbf 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -379,9 +379,10 @@ def pad_w13(x: torch.Tensor) -> torch.Tensor: tensor to padded_N rows.""" if padded_N == N: return x - x = x.view(E, num_shards, N, x.size(-1)) + experts = x.size(0) + x = x.view(experts, num_shards, N, x.size(-1)) x = torch.nn.functional.pad(x, (0, 0, 0, padded_N - N)) - return x.reshape(E, num_shards * padded_N, -1) + return x.reshape(experts, num_shards * padded_N, -1) def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: """Zero-pad the packed N (last) dim of a (E, K, N / packing) @@ -423,24 +424,29 @@ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: def permute_scales( scales: torch.Tensor, g_scales: torch.Tensor, name: str ) -> tuple[torch.Tensor, torch.Tensor]: - scales = scales.to(param_dtype) - - tensor_list = [] if "w13" in name: - scales = pad_w13(scales) size_n, size_k = padded_N * num_shards, K else: - scales = pad_w2(scales, packing=GROUP_SIZE) size_n, size_k = K, padded_N - # All experts share one global_scale, so compute the max - # scale_factor across all experts first, then apply uniformly. - combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype) + # A shared factor must cover all experts. Convert one expert at a time + # so loading never retains a half-precision copy of the entire bank. + combined_scale_factor = 1.0 + if param_dtype != torch.half: + maxima = torch.empty(E, device=scales.device, dtype=param_dtype) + for i in range(E): + torch.amax(scales[i].to(param_dtype), out=maxima[i]) + combined_scale_factor = _nvfp4_compute_scale_factor(maxima, param_dtype) + output = None for i in range(E): - scale = scales[i].T + scale = scales[i : i + 1].to(param_dtype) + if "w13" in name: + scale = pad_w13(scale) + else: + scale = pad_w2(scale, packing=GROUP_SIZE) marlin_scales = marlin_permute_scales( - s=scale, + s=scale[0].T, size_k=size_k, size_n=size_n, group_size=GROUP_SIZE, @@ -449,12 +455,18 @@ def permute_scales( marlin_scales, _ = nvfp4_marlin_process_scales( marlin_scales, scale_factor=combined_scale_factor, a_dtype=param_dtype ) - tensor_list.append(marlin_scales) + if output is None: + output = torch.empty( + (E, *marlin_scales.shape), + device=marlin_scales.device, + dtype=marlin_scales.dtype, + ) + output[i].copy_(marlin_scales) - scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + assert output is not None g_scales = nvfp4_marlin_process_global_scale(g_scales, param_dtype) g_scales = g_scales / combined_scale_factor - return scales, g_scales + return output, g_scales w13_scale, w13_scale_2 = permute_scales(w13_scale, w13_scale_2, "w13") w2_scale, w2_scale_2 = permute_scales(w2_scale, w2_scale_2, "w2") From 59aa0765cdb341e80b6048773f1a001d9bafc1ed Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:16:11 +0000 Subject: [PATCH 049/365] Bound GLM MTP vocabulary-head scale reduction memory Reduce and sanitize FP32 blocks without materializing a full vocabulary-sized copy. Preserve exact scale values, packed NVFP4 bytes, and the target BF16 head. CUDA parity and allocation tests pass; full TP2 startup is validated separately. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_model.py | 76 +++++++++++++++++++ vllm/models/glm5next/nvidia/mtp_draft_head.py | 15 +++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 39e845f70a04..2491b3b9f645 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -308,6 +308,82 @@ def test_glm5next_mtp_draft_head_rejects_unknown_mode(monkeypatch) -> None: mtp_draft_head.configured_draft_head_mode() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("kind", ["finite", "zeros", "nonfinite", "bf16_bits"]) +def test_glm5next_mtp_head_scale_preserves_fp32_reduction(kind: str) -> None: + if kind == "bf16_bits": + weight = torch.arange(65536, dtype=torch.int32, device="cuda").to(torch.int16) + weight = weight.view(torch.bfloat16).view(256, 256) + else: + weight = torch.linspace(-8, 6, 512, device="cuda", dtype=torch.bfloat16) + weight = weight.view(16, 32) + if kind == "zeros": + weight.zero_() + elif kind == "nonfinite": + weight[0, :3] = torch.tensor( + [float("nan"), float("inf"), -float("inf")], device="cuda" + ) + original_bits = weight.view(torch.int16).clone() + expected = 2688.0 / weight.float().abs().nan_to_num().max() + + actual = mtp_draft_head._nvfp4_weight_global_scale(weight) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(weight.view(torch.int16), original_bits, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_glm5next_mtp_head_scale_bounds_conversion_memory(record_property) -> None: + weight = torch.ones((1024, 16384), dtype=torch.bfloat16, device="cuda") + weight[-1, -1] = -8 + mtp_draft_head._nvfp4_weight_global_scale(weight[:1]) + torch.accelerator.synchronize() + torch.accelerator.reset_peak_memory_stats() + allocated = torch.accelerator.memory_allocated() + + actual = mtp_draft_head._nvfp4_weight_global_scale(weight) + + torch.accelerator.synchronize() + peak = torch.accelerator.max_memory_allocated() - allocated + record_property("peak_additional_bytes", peak) + torch.testing.assert_close( + actual, torch.tensor(336.0, device="cuda"), rtol=0, atol=0 + ) + assert weight[-1, -1].item() == -8 + assert peak < 17 * 1024 * 1024 + + +@pytest.mark.skipif( + not current_platform.is_device_capability(120), reason="requires SM120" +) +def test_glm5next_mtp_head_quantization_preserves_packed_weights() -> None: + import flashinfer + + source = torch.nn.Linear(256, 256, bias=False, dtype=torch.bfloat16).cuda() + source.shard_indices = SimpleNamespace() + original = source.weight.detach().clone() + scale = 2688.0 / original.float().abs().nan_to_num().max() + weight, scales = flashinfer.nvfp4_quantize( + original, + scale, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=False, + backend="cuda", + ) + expected = flashinfer.prepare_bf16_fp4_weights( + weight.view(torch.uint8), + scales, + scale.reciprocal().reshape(1), + backend="cute-dsl", + ) + + actual = mtp_draft_head.QuantizedDraftHead(source, "nvfp4") + + for observed, reference in zip(actual.buffers(), expected, strict=True): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + torch.testing.assert_close(source.weight, original, rtol=0, atol=0) + + def test_glm5next_mtp_prepares_configured_draft_head(monkeypatch) -> None: source_head = torch.nn.Linear(4, 8, bias=False) quantized_head = torch.nn.Linear(4, 8, bias=False) diff --git a/vllm/models/glm5next/nvidia/mtp_draft_head.py b/vllm/models/glm5next/nvidia/mtp_draft_head.py index ae662622e9eb..03a4f1f34937 100644 --- a/vllm/models/glm5next/nvidia/mtp_draft_head.py +++ b/vllm/models/glm5next/nvidia/mtp_draft_head.py @@ -30,6 +30,17 @@ _SUPPORTED_MODES = frozenset(("bf16", "nvfp4")) +def _nvfp4_weight_global_scale(weight: torch.Tensor) -> torch.Tensor: + """Compute the BF16 head's FP32 scale with at most 16 MiB of conversion scratch.""" + # Each float() result owns its storage; sanitization must not alter the + # BF16 target head. Chunk maxima preserve the full FP32 reduction exactly. + maxima = [ + chunk.float().abs_().nan_to_num_().amax() + for chunk in weight.view(-1).split(4 * 1024 * 1024) + ] + return _NVFP4_GLOBAL_MAX / torch.stack(maxima).amax() + + def configured_draft_head_mode() -> str: """Return the validated GLM MTP draft-head storage mode.""" mode = os.getenv("VLLM_GLM53_MTP_DRAFT_HEAD", "bf16").strip().lower() @@ -91,9 +102,7 @@ def __init__(self, source_head: nn.Module, mode: str) -> None: ) import flashinfer - weight_global_scale = ( - _NVFP4_GLOBAL_MAX / weight.float().abs().nan_to_num().max() - ) + weight_global_scale = _nvfp4_weight_global_scale(weight) weight_fp4, weight_sf = flashinfer.nvfp4_quantize( weight, weight_global_scale, From b47c3d049a0b9d0a6ba893c7a18e5c4fa008167a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 11 Sep 2026 23:46:40 +0000 Subject: [PATCH 050/365] Share GLM pooled-indexer temporary storage across sequential target layers Keep query, scale, score and block-table buffers under one model-owned scratch module. Preserve full token-row capacity, per-layer recurrent tails and independent MTP storage. Rebinding equal table geometry preserves graph addresses. Five focused scratch/capacity tests pass, including live-input CUDA graph replay and independent recurrent/draft ownership. Existing model and attention suites are recorded in local qualification artifacts. Co-authored-by: Codex (cherry picked from commit d5b19d25eeedec5a1e127531d39938e11de364e2) Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 154 +++++++++++++++++- vllm/models/glm5next/nvidia/attention.py | 4 +- vllm/models/glm5next/nvidia/model.py | 11 +- vllm/models/glm5next/nvidia/pooled_indexer.py | 152 ++++++++--------- 4 files changed, 236 insertions(+), 85 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 5b6d3ca555bc..21248d24b3b2 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -10,6 +10,7 @@ from torch import nn from vllm.models.deepseek_v4.nvidia.b12x_indexer import _flatten_index_cache +from vllm.models.glm5next.nvidia import pooled_indexer as indexer_module from vllm.models.glm5next.nvidia.ops import glm_kpool from vllm.models.glm5next.nvidia.ops.glm_kpool import ( expand_c4_block_table, @@ -19,7 +20,10 @@ prepare_c4_decode_metadata, update_decode_pools, ) -from vllm.models.glm5next.nvidia.pooled_indexer import Glm5NextPooledIndexer +from vllm.models.glm5next.nvidia.pooled_indexer import ( + Glm5NextIndexerScratch, + Glm5NextPooledIndexer, +) from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( @@ -659,7 +663,7 @@ def test_glm53_physical_selection_provider_is_explicit() -> None: indexer.dcp_world_size = 1 indexer._emit_physical_selection = True indexer.topk_indices_buffer = torch.empty((8, 2051), dtype=torch.int32) - indexer._physical_active_counts = torch.empty(8, dtype=torch.int32) + indexer.scratch = Glm5NextIndexerScratch(8, 1, torch.device("cpu")) selected = indexer.get_b12x_physical_selection( num_tokens=3, @@ -718,8 +722,7 @@ def test_glm53_missing_physical_plan_does_not_mutate_recurrent_cache( index_kpool_compress_gate=torch.zeros((128, 128)), index_kpool_compress_ape=torch.zeros((4, 128)), _project_head_weights=lambda value: torch.zeros((1, 32)), - _q_fp8=torch.empty((1, 32, 128), dtype=torch.float8_e4m3fn), - _q_scale=torch.empty((1, 32)), + scratch=Glm5NextIndexerScratch(1, 1, torch.device("cpu")), main_layer_name="attention", _index_cache=torch.zeros((1, 64, 132), dtype=torch.uint8), _tail=torch.zeros((1, 2, 4, 128)), @@ -815,7 +818,7 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( indexer._physical_selection_plan = object() indexer.bind_main_kv_cache(main) - assert indexer.indexer_op._num_q_heads == indexer._q_fp8.shape[1] == 32 + assert indexer.indexer_op._num_q_heads == indexer.scratch.q_fp8.shape[1] == 32 assert indexer._physical_selection_plan is None indexer._physical_selection_plan = object() indexer.unbind_main_kv_cache() @@ -862,14 +865,19 @@ def publish(*args, **kwargs): indexer.indexer_op = SimpleNamespace( max_model_len=4096 // 4, set_b12x_index_cache=publish ) + indexer.scratch = Glm5NextIndexerScratch( + indexer.max_tokens, indexer.max_seqs, device + ) indexer.bind_main_kv_cache(main) - assert indexer._decode_block_table.shape[0] == indexer.max_tokens - assert indexer._decode_block_table.shape[0] > indexer.max_seqs * (5 + 1) + assert indexer.scratch.decode_block_table.shape[0] == indexer.max_tokens + assert indexer.scratch.decode_block_table.shape[0] > indexer.max_seqs * (5 + 1) assert indexer.indexer_op.max_model_len == 4096 // 4 assert published["num_q_heads"] == 32 - assert published["max_page_table_width"] == indexer._pool_block_table.shape[1] + assert ( + published["max_page_table_width"] == indexer.scratch.pool_block_table.shape[1] + ) def test_b12x_c4_declaration_matches_glm_query_and_page_table_layout() -> None: @@ -1025,7 +1033,8 @@ def test_glm53_physical_selection_prepares_before_resolution_freeze( indexer.max_tokens = 32 indexer.prefix = "model.layers.3.indexer" indexer.topk_indices_buffer = output - indexer._physical_active_counts = counts + indexer.scratch = Glm5NextIndexerScratch(32, 1, device) + indexer.scratch.physical_active_counts = counts request = indexer.get_b12x_physical_selection_preparation_request() session = PreparationSession(device=device, autotune=autotune, compile_workers=0) session.prepare((request,)) @@ -1072,6 +1081,133 @@ def test_glm53_selector_capacity_tracks_auto_fit_max_model_len() -> None: assert indexer.indexer_op.max_model_len == 497 +def test_glm53_indexer_scratch_rebind_preserves_same_geometry_storage() -> None: + scratch = Glm5NextIndexerScratch(128, 4, torch.device("cpu")) + scratch.bind_tables(32, torch.device("cpu")) + pointers = {name: buffer.data_ptr() for name, buffer in scratch.named_buffers()} + scratch.decode_block_table.fill_(53) + scratch.bind_tables(32, torch.device("cpu")) + assert { + name: buffer.data_ptr() for name, buffer in scratch.named_buffers() + } == pointers + assert torch.all(scratch.decode_block_table == 53) + scratch.bind_tables(16, torch.device("cpu")) + assert scratch.decode_block_table.shape == (128, 16) + assert scratch.pool_block_table.shape == (4, 16) + assert scratch.q_fp8.data_ptr() == pointers["q_fp8"] + assert not scratch.state_dict() + + +def test_glm53_shared_scratch_keeps_layer_checkpoints_and_draft_selection_private( + monkeypatch, +) -> None: + monkeypatch.setattr( + indexer_module, "ReplicatedLinear", lambda *a, **k: nn.Identity() + ) + monkeypatch.setattr( + indexer_module, "B12xC4SparseIndexer", lambda *a, **k: SimpleNamespace() + ) + monkeypatch.setattr( + indexer_module, + "get_b12x_sparse_mla", + lambda: SimpleNamespace(expand_pooled_topk_to_physical_slots=lambda *a: None), + ) + config = SimpleNamespace( + index_topk=2048, + index_n_heads=32, + index_head_dim=128, + index_kpool=4, + qk_rope_head_dim=0, + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=128, max_num_seqs=4), + model_config=SimpleNamespace(max_model_len=4096), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, cp_kv_cache_interleave_size=4 + ), + speculative_config=SimpleNamespace(num_speculative_tokens=3), + ) + scratch = Glm5NextIndexerScratch(128, 4, torch.device("cpu")) + target_ids = torch.zeros((128, 2051), dtype=torch.int32) + target_pools = torch.zeros((128, 512), dtype=torch.int32) + + def make(name, shared): + return Glm5NextPooledIndexer( + vllm_config, + config, + 8, + 8, + None, + SimpleNamespace(block_size=2048), + target_ids if shared else torch.full_like(target_ids, 42), + target_pools if shared else torch.full_like(target_pools, 42), + main_layer_name=name, + prefix=name, + scratch=scratch if shared else None, + ) + + first, second, draft = ( + make("target.0", True), + make("target.1", True), + make("mtp", False), + ) + assert first.scratch is second.scratch + assert draft.scratch is not scratch + for indexer, value in ((first, 1), (second, 2), (draft, 3)): + indexer._tail.fill_(value) + indexer.snapshot_speculative_interval_starts() + scratch.q_fp8.view(torch.uint8).fill_(255) + target_ids.fill_(-1) + first._tail.zero_() + first.restore_speculative_interval_starts() + for indexer, value in ((first, 1), (second, 2), (draft, 3)): + assert torch.all(indexer._tail == value) + assert torch.all(indexer._tail_snapshot == value) + assert torch.all(draft.topk_indices_buffer == 42) + assert torch.all(draft.pool_topk_indices_buffer == 42) + + +def test_glm53_shared_indexer_scratch_graph_consumes_each_live_query() -> None: + device = _require_glm_gpu() + scratch = Glm5NextIndexerScratch(8, 1, device) + inputs = [ + torch.randn((8, 32, 128), device=device, dtype=torch.bfloat16) for _ in range(2) + ] + outputs = [torch.empty_like(scratch.q_fp8) for _ in range(2)] + scales = [torch.empty_like(scratch.q_scale) for _ in range(2)] + + def run() -> None: + for query, output, scale in zip(inputs, outputs, scales): + glm_kpool.fwht128_quant_fp8( + query.view(-1, 128), + scratch.q_fp8.view(-1, 128), + scratch.q_scale.view(-1), + ) + output.copy_(scratch.q_fp8) + scale.copy_(scratch.q_scale) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for _ in range(3): + for query in inputs: + query.normal_() + scratch.q_fp8.view(torch.uint8).fill_(255) + scratch.q_scale.fill_(float("nan")) + graph.replay() + for query, output, scale in zip(inputs, outputs, scales): + reference_q = torch.empty_like(output) + reference_scale = torch.empty_like(scale) + glm_kpool.fwht128_quant_fp8( + query.view(-1, 128), + reference_q.view(-1, 128), + reference_scale.view(-1), + ) + torch.testing.assert_close(output, reference_q, rtol=0, atol=0) + torch.testing.assert_close(scale, reference_scale, rtol=0, atol=0) + + def test_glm53_parent_table_width_tracks_dcp_sharding() -> None: max_model_len = 524288 block_size = 2304 diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py index f3a093a14ff5..fe8714b7c0e6 100644 --- a/vllm/models/glm5next/nvidia/attention.py +++ b/vllm/models/glm5next/nvidia/attention.py @@ -22,7 +22,7 @@ from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.registry import AttentionBackendEnum -from .pooled_indexer import Glm5NextPooledIndexer +from .pooled_indexer import Glm5NextIndexerScratch, Glm5NextPooledIndexer def _select_sparse_backend( @@ -63,6 +63,7 @@ def __init__( skip_rope: bool | None = False, is_mtp_layer: bool = False, attn_backend: type | None = None, + indexer_scratch: Glm5NextIndexerScratch | None = None, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -171,6 +172,7 @@ def __init__( prefix=f"{prefix}.indexer", # MTP compacts and reuses request-relative selections. emit_physical_selection=not is_mtp_layer, + scratch=indexer_scratch, ) else: self.indexer = None diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index a1f5adecf675..8107a0d9d8f1 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -96,7 +96,7 @@ from . import l2_prefetch as _l2pf from .attention import Glm5NextMLAAttention from .kda import Glm5NextLinearAttention -from .pooled_indexer import Glm5NextPooledIndexer +from .pooled_indexer import Glm5NextIndexerScratch, Glm5NextPooledIndexer logger = init_logger(__name__) @@ -333,6 +333,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None = None, pool_topk_indices_buffer: torch.Tensor | None = None, is_mtp_layer: bool = False, + indexer_scratch: Glm5NextIndexerScratch | None = None, **kwargs, ) -> None: super().__init__() @@ -389,6 +390,7 @@ def __init__( pool_topk_indices_buffer=pool_topk_indices_buffer, skip_rope=getattr(config, "mla_nope", False), is_mtp_layer=is_mtp_layer, + indexer_scratch=indexer_scratch, ) # MTP layers sit past the base model's hidden layers (layer_idx >= @@ -960,10 +962,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dtype=torch.int32, device=self.device, ) + indexer_scratch = Glm5NextIndexerScratch( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.scheduler_config.max_num_seqs, + torch.device(self.device), + ) else: # Full-MLA config (no kpool sparse indexer): no topk buffer. topk_indices_buffer = None pool_topk_indices_buffer = None + indexer_scratch = None if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -983,6 +991,7 @@ def get_layer(prefix: str): prefix=prefix, topk_indices_buffer=topk_indices_buffer, pool_topk_indices_buffer=pool_topk_indices_buffer, + indexer_scratch=indexer_scratch, ) self.start_layer, self.end_layer, self.layers = make_layers( diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 63ee0d63546f..9e94d04e6455 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -51,6 +51,55 @@ _INDEX_PAGE_BYTES = _INDEX_PAGE_SIZE * _INDEX_CACHE_WIDTH +class Glm5NextIndexerScratch(nn.Module): + """Temporary storage for sequential pooled-indexer calls on one model stream. + + A model owns one instance; its draft model owns another. KV views, recurrent + tails and selected token IDs are not scratch and remain separately owned. + Bind tables before graph capture; calls must finish consuming the scratch + before another layer overwrites it. + """ + + q_fp8: torch.Tensor + q_scale: torch.Tensor + pool_scores: torch.Tensor + pool_seq_lens: torch.Tensor + physical_active_counts: torch.Tensor + pool_block_table: torch.Tensor + decode_block_table: torch.Tensor + + def __init__(self, max_tokens: int, max_seqs: int, device: torch.device) -> None: + super().__init__() + self.max_tokens = max_tokens + self.max_seqs = max_seqs + for name, shape, dtype in ( + ("q_fp8", (max_tokens, _INDEX_HEADS, _INDEX_HEAD_DIM), torch.float8_e4m3fn), + ("q_scale", (max_tokens, _INDEX_HEADS), torch.float32), + ("pool_scores", (max_tokens, _POOL_TOPK), torch.float32), + ("pool_seq_lens", (max_tokens,), torch.int32), + ("physical_active_counts", (max_tokens,), torch.int32), + ("pool_block_table", (max_seqs, 1), torch.int32), + ("decode_block_table", (max_tokens, 1), torch.int32), + ): + self.register_buffer( + name, torch.empty(shape, dtype=dtype, device=device), persistent=False + ) + + def bind_tables(self, width: int, device: torch.device) -> None: + if self.pool_block_table.shape[1] == width and self.q_fp8.device == device: + return + if self.q_fp8.device != device: + raise ValueError("GLM indexer scratch and KV cache must share a device") + self.pool_block_table = torch.empty( + (self.max_seqs, width), dtype=torch.int32, device=device + ) + # Short prefills can use decode metadata, so rows need the token budget, + # not merely max_seqs times the speculative width. + self.decode_block_table = torch.empty( + (self.max_tokens, width), dtype=torch.int32, device=device + ) + + class Glm5NextPooledIndexer(nn.Module): """Produce GLM C4 pools, select them with b12x, and expand token IDs.""" @@ -68,6 +117,7 @@ def __init__( main_layer_name: str, prefix: str, emit_physical_selection: bool = True, + scratch: Glm5NextIndexerScratch | None = None, ) -> None: super().__init__() if cache_config is None: @@ -132,6 +182,17 @@ def __init__( raise ValueError("GLM pool selection buffer has the wrong row capacity") device = topk_indices_buffer.device + self.scratch = ( + scratch + if scratch is not None + else Glm5NextIndexerScratch(self.max_tokens, self.max_seqs, device) + ) + if ( + self.scratch.max_tokens != self.max_tokens + or self.scratch.max_seqs != self.max_seqs + or self.scratch.q_fp8.device != device + ): + raise ValueError("GLM indexer scratch has incompatible capacity or device") b12x_sparse_mla = get_b12x_sparse_mla() if b12x_sparse_mla is None or not hasattr( b12x_sparse_mla, "expand_pooled_topk_to_physical_slots" @@ -219,57 +280,6 @@ def __init__( torch.empty_like(self._tail), persistent=False, ) - self.register_buffer( - "_q_fp8", - torch.empty( - (self.max_tokens, _INDEX_HEADS, _INDEX_HEAD_DIM), - dtype=torch.float8_e4m3fn, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_physical_active_counts", - torch.empty(self.max_tokens, dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_q_scale", - torch.empty( - (self.max_tokens, _INDEX_HEADS), - dtype=torch.float32, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_pool_seq_lens", - torch.empty(self.max_tokens, dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_pool_scores", - torch.empty( - (self.max_tokens, _POOL_TOPK), - dtype=torch.float32, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_pool_block_table", - torch.empty((self.max_seqs, 1), dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_decode_block_table", - torch.empty( - (self.max_tokens, 1), - dtype=torch.int32, - device=device, - ), - persistent=False, - ) self._weights_proj_fp32: torch.Tensor | None = None self._index_cache: torch.Tensor | None = None self._parent_table_width = 0 @@ -370,15 +380,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: self.dcp_world_size, ) pool_table_width = parent_table_width * subpages - device = main_cache.device - self._pool_block_table = torch.empty( - (self.max_seqs, pool_table_width), dtype=torch.int32, device=device - ) - self._decode_block_table = torch.empty( - (self.max_tokens, pool_table_width), - dtype=torch.int32, - device=device, - ) + self.scratch.bind_tables(pool_table_width, main_cache.device) self._index_cache = index_cache self._parent_table_width = parent_table_width self._subpages_per_parent = subpages @@ -450,7 +452,7 @@ def get_b12x_physical_selection_preparation_request(self): def prepare(state): return self.make_b12x_physical_selection_prepare_call( torch.empty_like(self.topk_indices_buffer), - torch.empty_like(self._physical_active_counts), + torch.empty_like(self.scratch.physical_active_counts), state=state, ) @@ -485,7 +487,7 @@ def get_b12x_physical_selection( return None return ( self.topk_indices_buffer[:num_tokens], - self._physical_active_counts[:num_tokens], + self.scratch.physical_active_counts[:num_tokens], ) def _project_head_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -522,8 +524,8 @@ def forward( gate = F.linear(hidden_states, self.index_kpool_compress_gate) weights = self._project_head_weights(hidden_states) - q_fp8 = self._q_fp8[:rows] - q_scale = self._q_scale[:rows] + q_fp8 = self.scratch.q_fp8[:rows] + q_scale = self.scratch.q_scale[:rows] fwht128_quant_fp8( query.contiguous().view(-1, _INDEX_HEAD_DIM), q_fp8.view(-1, _INDEX_HEAD_DIM), @@ -595,8 +597,8 @@ def forward( parent_stride_pages=self._parent_stride_pages, ) parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] - seq_lens = self._pool_seq_lens[:live_rows] - decode_table = self._decode_block_table[:decode_rows] + seq_lens = self.scratch.pool_seq_lens[:live_rows] + decode_table = self.scratch.decode_block_table[:decode_rows] if decode_only: prepare_c4_decode_metadata( parent_table, @@ -613,7 +615,7 @@ def forward( else: expand_c4_block_table( parent_table, - self._pool_block_table, + self.scratch.pool_block_table, rows=num_reqs, subpages_per_parent=self._subpages_per_parent, parent_stride_pages=self._parent_stride_pages, @@ -626,12 +628,14 @@ def forward( pool_interleave=self.pool_interleave, ) pool_ids = self.pool_topk_indices_buffer[:rows] - pool_scores = self._pool_scores[:rows] if self.dcp_world_size > 1 else None + pool_scores = ( + self.scratch.pool_scores[:rows] if self.dcp_world_size > 1 else None + ) if decode_rows: if not decode_only: gather_c4_block_table_rows( - self._pool_block_table, + self.scratch.pool_block_table, main_metadata.req_id_per_token[:decode_rows], decode_table, ) @@ -671,19 +675,19 @@ def forward( active_pages = self._active_index_page_count( int(request_seq_lens_cpu[local_request]) ) - if active_pages > int(self._pool_block_table.shape[1]): + if active_pages > int(self.scratch.pool_block_table.shape[1]): raise RuntimeError( "GLM selector visible C4 pages exceed table capacity: " f"active={active_pages}, " - f"capacity={int(self._pool_block_table.shape[1])}" + f"capacity={int(self.scratch.pool_block_table.shape[1])}" ) - request_table = self._pool_block_table[ + request_table = self.scratch.pool_block_table[ request : request + 1, :active_pages ] else: # DCP pool ownership is interleaved across ranks, so a global # sequence length does not define a contiguous local prefix. - request_table = self._pool_block_table[request : request + 1] + request_table = self.scratch.pool_block_table[request : request + 1] shared_table = request_table.expand(int(query_len), -1) self.indexer_op.run_paged_topk( q=q_fp8[row_start:row_end], @@ -723,7 +727,7 @@ def forward( main_metadata.req_id_per_token[:live_rows], main_metadata.block_table, output[:live_rows], - self._physical_active_counts[:live_rows], + self.scratch.physical_active_counts[:live_rows], pool_size=_POOL_SIZE, block_size=self.block_size, block_stride_rows=self.block_size, From a311102cd031b02b5fac9cfa3ed2dc2af01c767a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:26:17 +0000 Subject: [PATCH 051/365] Reuse GLM indexer storage across serial target and MTP execution Port the storage-sharing part of PR #739 commit 7f4aecc66e857d093e012a2c57d03711bd23be4a to Karmic loading. Require one draft layer, matching geometry, PP1 and no overlapped microbatches. Keep parameters, recurrent tails and KV views independent. Eleven CPU/CUDA ownership tests pass, including live-input graph replay. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_indexer_storage.py | 194 ++++++++++++++++++ vllm/models/glm5next/nvidia/mtp.py | 71 +++++++ vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 9 +- 3 files changed, 273 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_glm5next_indexer_storage.py diff --git a/tests/models/test_glm5next_indexer_storage.py b/tests/models/test_glm5next_indexer_storage.py new file mode 100644 index 000000000000..9b872bffb211 --- /dev/null +++ b/tests/models/test_glm5next_indexer_storage.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ownership checks for serial GLM target/MTP indexer storage.""" + +import weakref +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.glm5next.nvidia.mtp import Glm5NextMTP +from vllm.models.glm5next.nvidia.pooled_indexer import ( + Glm5NextIndexerScratch, + Glm5NextPooledIndexer, +) + + +def make_indexer(device="cpu", shared=None): + indexer = Glm5NextPooledIndexer.__new__(Glm5NextPooledIndexer) + nn.Module.__init__(indexer) + for name, value in dict( + max_tokens=16, + max_seqs=4, + max_model_len=1048576, + block_size=2048, + dcp_world_size=2, + dcp_rank=0, + pool_interleave=1, + ).items(): + setattr(indexer, name, value) + indexer.scratch = ( + shared.scratch + if shared + else Glm5NextIndexerScratch(16, 4, torch.device(device)) + ) + for name, width in ( + ("topk_indices_buffer", 2051), + ("pool_topk_indices_buffer", 512), + ): + value = ( + getattr(shared, name) + if shared + else torch.empty((16, width), dtype=torch.int32, device=device) + ) + setattr(indexer, name, value) + indexer.indexer_op = nn.Module() + indexer.indexer_op.topk_indices_buffer = indexer.pool_topk_indices_buffer + indexer._index_cache = torch.ones(32, dtype=torch.uint8, device=device) + indexer.register_buffer("_tail", torch.ones(32, device=device), persistent=False) + indexer.register_parameter("weight", nn.Parameter(torch.ones(4, device=device))) + return indexer + + +def make_pair(device="cpu"): + target = nn.Module() + first = make_indexer(device) + target.layers = nn.ModuleList([first, make_indexer(device, shared=first)]) + draft = Glm5NextMTP.__new__(Glm5NextMTP) + nn.Module.__init__(draft) + draft.model = nn.Module() + draft.model.num_mtp_layers = 1 + draft.model.indexer = make_indexer(device) + draft.model.attention = nn.Module() + draft.model.attention.topk_indices_buffer = draft.model.indexer.topk_indices_buffer + draft.model.attention.impl = SimpleNamespace( + topk_indices_buffer=draft.model.indexer.topk_indices_buffer + ) + return target, draft + + +def test_share_indexer_storage_preserves_roles_and_persistent_state(): + target, draft = make_pair() + source, indexer = target.layers[0], draft.model.indexer + scratch_ref = weakref.ref(indexer.scratch) + token_ref = weakref.ref(indexer.topk_indices_buffer) + pool_ref = weakref.ref(indexer.pool_topk_indices_buffer) + cache, tail, weight = indexer._index_cache, indexer._tail, indexer.weight + assert draft.share_target_indexer_storage(target) + assert scratch_ref() is token_ref() is pool_ref() is None + assert indexer.scratch is source.scratch + assert indexer.topk_indices_buffer is source.topk_indices_buffer + assert indexer.pool_topk_indices_buffer is source.pool_topk_indices_buffer + assert indexer.indexer_op.topk_indices_buffer is source.pool_topk_indices_buffer + assert draft.model.attention.topk_indices_buffer is source.topk_indices_buffer + assert draft.model.attention.impl.topk_indices_buffer is source.topk_indices_buffer + assert indexer._index_cache is cache + assert indexer._tail is tail + assert indexer.weight is weight + source.scratch.bind_tables(2048, torch.device("cpu")) + pointer = source.scratch.decode_block_table.data_ptr() + indexer.scratch.bind_tables(2048, torch.device("cpu")) + assert indexer.scratch.decode_block_table.data_ptr() == pointer + assert draft.share_target_indexer_storage(target) # Idempotent ownership. + + +@pytest.mark.parametrize( + "difference", + [ + "layers", + "geometry", + "capacity", + "dtype", + "scratch", + "target_storage", + ], +) +def test_incompatible_indexer_storage_remains_independent(difference): + target, draft = make_pair() + indexer = draft.model.indexer + if difference == "layers": + draft.model.num_mtp_layers = 2 + elif difference == "geometry": + indexer.dcp_world_size = 1 + elif difference == "capacity": + indexer.topk_indices_buffer = indexer.topk_indices_buffer[:8] + elif difference == "dtype": + indexer.pool_topk_indices_buffer = indexer.pool_topk_indices_buffer.long() + elif difference == "scratch": + indexer.scratch.bind_tables(2, torch.device("cpu")) + else: + target.layers[1].scratch = Glm5NextIndexerScratch(16, 4, torch.device("cpu")) + scratch, tokens, pools = ( + indexer.scratch, + indexer.topk_indices_buffer, + indexer.pool_topk_indices_buffer, + ) + assert not draft.share_target_indexer_storage(target) + assert indexer.scratch is scratch + assert indexer.topk_indices_buffer is tokens + assert indexer.pool_topk_indices_buffer is pools + + +@pytest.mark.parametrize( + "pp,dbo,expected", [(1, False, True), (2, False, False), (1, True, False)] +) +def test_loader_shares_only_serial_indexers(monkeypatch, pp, dbo, expected): + from vllm.v1.worker.gpu.spec_decode.eagle import utils + + inner, draft = make_pair() + target = nn.Module() + target.model = inner + config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=None, + moe_backend=None, + kv_cache_dtype=None, + attention_backend=None, + ), + parallel_config=SimpleNamespace(enable_dbo=dbo), + load_config=SimpleNamespace(load_format="auto"), + ) + monkeypatch.setattr(utils, "_make_eagle_draft_vllm_config", lambda value: value) + monkeypatch.setattr(utils, "get_model", lambda **_: draft) + monkeypatch.setattr(utils, "get_pp_group", lambda: SimpleNamespace(world_size=pp)) + monkeypatch.setattr(utils, "get_pp_safe_draft_load_config", lambda value: value) + assert utils.load_eagle_model(target, config) is draft + assert (draft.model.indexer.scratch is inner.layers[0].scratch) is expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph ownership") +def test_shared_indexer_storage_serial_graph_replay(): + target, draft = make_pair("cuda") + assert draft.share_target_indexer_storage(target) + source, indexer = target.layers[0], draft.model.indexer + value = torch.ones((), device="cuda") + + def run(owner, factor): + owner.scratch.pool_scores.copy_(value * factor) + owner.topk_indices_buffer.fill_(factor) + owner.pool_topk_indices_buffer.fill_(factor + 1) + return ( + owner.scratch.pool_scores.clone(), + owner.topk_indices_buffer.clone(), + owner.pool_topk_indices_buffer.clone(), + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run(source, 2) + run(indexer, 3) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + target_result = run(source, 2) + draft_result = run(indexer, 3) + for scalar in (1, 5, 9): + value.fill_(scalar) + graph.replay() + for result, factor in ((target_result, 2), (draft_result, 3)): + assert torch.equal(result[0], torch.full_like(result[0], scalar * factor)) + assert torch.equal(result[1], torch.full_like(result[1], factor)) + assert torch.equal(result[2], torch.full_like(result[2], factor + 1)) diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index 3bf518b5d636..423a3416f4b6 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -336,6 +336,77 @@ def prepare_draft_lm_head(self, source_head: nn.Module) -> None: """Create a draft-only quantized copy of the shared target head.""" self.model.prepare_draft_lm_head(source_head) + def share_target_indexer_storage(self, target: nn.Module) -> bool: + """Reuse temporary indexer storage under serial target/MTP execution. + + The V2 loader calls this before profiling or graph capture, without + pipeline parallelism or overlapped microbatches. One MTP layer consumes + its selections before target execution resumes. KV views, pool tails + and model parameters retain independent ownership. + """ + if self.model.num_mtp_layers != 1: + return False + targets = [ + module + for module in target.modules() + if isinstance(module, Glm5NextPooledIndexer) + ] + drafts = [ + module + for module in self.model.modules() + if isinstance(module, Glm5NextPooledIndexer) + ] + if not targets or len(drafts) != 1: + return False + source, draft = targets[0], drafts[0] + geometry = ( + "max_tokens", + "max_seqs", + "max_model_len", + "block_size", + "dcp_world_size", + "dcp_rank", + "pool_interleave", + ) + if any(getattr(source, key) != getattr(draft, key) for key in geometry): + return False + names = ("topk_indices_buffer", "pool_topk_indices_buffer") + for module in targets: + if module.scratch is not source.scratch or any( + getattr(module, name) is not getattr(source, name) for name in names + ): + return False + for name in names: + src, dst = getattr(source, name), getattr(draft, name) + if (src.shape, src.dtype, src.device) != (dst.shape, dst.dtype, dst.device): + return False + if any( + (src.shape, src.dtype, src.device) != (dst.shape, dst.dtype, dst.device) + for (_, src), (_, dst) in zip( + source.scratch.named_buffers(), + draft.scratch.named_buffers(), + strict=True, + ) + ): + return False + + replacements = { + id(getattr(draft, name)): getattr(source, name) for name in names + } + # The native indexer holds 512 pool IDs; MLA consumes 2051 token IDs. + # Rebind by identity, not by attribute name, including non-Module impls. + for module in self.model.modules(): + for owner in (module, getattr(module, "impl", None)): + if owner is None: + continue + for name in names: + value = getattr(owner, name, None) + replacement = replacements.get(id(value)) + if replacement is not None: + setattr(owner, name, replacement) + draft.scratch = source.scratch + return True + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 3d42451a166a..89ed39d14f2b 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -182,7 +182,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod # every module in the draft that holds a buffer reference so that # the per-layer indexer and sparse-attention backends all point to # the target's buffer. - if hasattr(target_inner, "topk_indices_buffer"): + share_indexer_storage = getattr(eagle_model, "share_target_indexer_storage", None) + if share_indexer_storage is not None: + if ( + get_pp_group().world_size == 1 + and not vllm_config.parallel_config.enable_dbo + ): + share_indexer_storage(target_inner) + elif hasattr(target_inner, "topk_indices_buffer"): target_buffer = target_inner.topk_indices_buffer if target_buffer is not None: for _, module in draft_inner.named_modules(): From adc85c66df33e94f4c8e6eb704d1e7d3cfbbfb6b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:37:31 +0000 Subject: [PATCH 052/365] Bound GLM vision projection memory without partitioning attention Port the vision and workspace portions of 7f4aecc66e857 and the rotary-bounded profiling canvas from ab51db9b4801a551e5c5fc4a6e0b4e20fb37bbf4 to the common GLM model module. Preserve the Karmic processor and execution-lane contracts. Token-local projections use bounded chunks; image attention retains complete Q/K/V and sequence boundaries. Validation: 30 focused tests pass. Native 24-layer BF16 vision on TP2 produces bit-identical output on both ranks; additional peak allocation falls from 662164480 to 221254656 bytes, or 179295232 bytes when borrowing preallocated worker storage. Whole-model serving validation is separate. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_model.py | 162 +++++++++++++++++++ tests/v1/worker/test_workspace.py | 4 + vllm/models/glm5next/common/multimodal.py | 184 +++++++++++++++++++++- vllm/v1/worker/workspace.py | 17 ++ 4 files changed, 360 insertions(+), 7 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 2491b3b9f645..8dd7d4c8dca8 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -5,6 +5,7 @@ import math import sys import types +import weakref from dataclasses import dataclass from types import SimpleNamespace @@ -2415,6 +2416,167 @@ def fake_from_pretrained(model, **kwargs): ] +def test_glm5next_vision_attention_releases_projection_and_norm_inputs(monkeypatch): + from vllm.models.glm5next.common import multimodal + + owners = {} + + def project(value): + packed = value.repeat(1, 1, 3) + owners["packed"] = weakref.ref(packed) + return packed, None + + def normalize(q, k, *args): + owners["q_norm_input"] = weakref.ref(q) + owners["k_norm_input"] = weakref.ref(k) + return q.clone(), k.clone() + + def attention(query, key, value, **kwargs): + assert all(owner() is None for owner in owners.values()) + owners.update(q=weakref.ref(query), k=weakref.ref(key), v=weakref.ref(value)) + return value.clone() + + def output_project(value): + assert all(owner() is None for owner in owners.values()) + return value, None + + layer = SimpleNamespace( + head_dim=2, + num_attention_heads_per_partition=2, + hidden_size_per_attention_head=2, + qkv=project, + q_norm=SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + k_norm=SimpleNamespace(weight=torch.ones(2)), + apply_rotary_emb=lambda value, cos, sin: value.clone(), + attn=attention, + proj=output_project, + ) + layer.split_qkv = lambda value: multimodal.Glm5NextVisionAttention.split_qkv( + layer, value + ) + layer._project_qkv = lambda *args: multimodal.Glm5NextVisionAttention._project_qkv( + layer, *args + ) + monkeypatch.setattr(multimodal, "fused_q_kv_rmsnorm", normalize) + x = torch.arange(28, dtype=torch.float32).view(7, 1, 4) + original = x.clone() + output = multimodal.Glm5NextVisionAttention.forward( + layer, x, torch.tensor([0, 7]), torch.ones(7, 1), torch.zeros(7, 1), 7 + ) + torch.testing.assert_close(output, original, rtol=0, atol=0) + torch.testing.assert_close(x, original, rtol=0, atol=0) + + +@pytest.mark.parametrize("reserved_bytes", [0, 256, 4096]) +def test_glm5next_chunked_projection_keeps_whole_image_attention( + monkeypatch, reserved_bytes +): + from vllm.models.glm5next.common import multimodal + from vllm.v1.worker import workspace + + manager = workspace.WorkspaceManager(torch.device("cpu")) + manager.reserve_all(((reserved_bytes,), torch.uint8)) + manager.lock() + monkeypatch.setattr(workspace, "_manager", manager) + monkeypatch.setattr(multimodal, "_VISION_PROJECTION_CHUNK_ROWS", 4) + x = torch.randn(11, 1, 8, generator=torch.Generator().manual_seed(72)) + cu_seqlens = torch.tensor([0, 3, 11]) + cos, sin = torch.ones(11, 2), torch.zeros(11, 2) + projected_rows = [] + attention_calls = [] + + def project(value, cos, sin): + projected_rows.append(len(value)) + assert cos.shape[0] == sin.shape[0] == len(value) + value = value.reshape(1, len(value), 2, 4) + return value, value * 0.5, value * 2 + + def attention(query, key, value, **metadata): + assert metadata["cu_seqlens"] is cu_seqlens + attention_calls.append(query.shape[1]) + outputs = [] + for start, stop in zip(cu_seqlens[:-1], cu_seqlens[1:]): + q, k, v = (t[:, start:stop].transpose(1, 2) for t in (query, key, value)) + outputs.append( + torch.nn.functional.scaled_dot_product_attention(q, k, v).transpose( + 1, 2 + ) + ) + return torch.cat(outputs, dim=1) + + layer = SimpleNamespace( + num_attention_heads_per_partition=2, + hidden_size_per_attention_head=4, + _project_qkv=project, + _borrow_qkv_workspace=True, + attn=attention, + ) + norm = torch.nn.RMSNorm(8) + q, k, v = project(norm(x), cos, sin) + expected = attention(q, k, v, cu_seqlens=cu_seqlens).reshape_as(x) + projected_rows.clear() + attention_calls.clear() + actual = multimodal.Glm5NextVisionAttention.context_with_chunked_projection( + layer, x, norm, cu_seqlens, cos, sin, 8 + ) + torch.testing.assert_close(actual, expected) + assert projected_rows == [4, 4, 3] + assert attention_calls == [11] + + +@pytest.mark.parametrize( + "token_budget,expand,expected_tokens", + [ + (8000, 1, 8000), + (8192, 1, 8192), + (8000, 2, 8000), + (8191, 1, 8190), + (8191, 2, 8188), + ], +) +def test_glm5next_image_budget_covers_rectangular_canvases( + token_budget, expand, expected_tokens +): + from vllm.models.glm5next.common.multimodal import Glm5NextProcessingInfo + + processor = SimpleNamespace( + patch_size=14, merge_size=2, temporal_patch_size=2, patch_expand_factor=expand + ) + info = Glm5NextProcessingInfo.__new__(Glm5NextProcessingInfo) + info.get_hf_processor = lambda **kwargs: SimpleNamespace(image_processor=processor) + info.get_hf_config = lambda: SimpleNamespace( + vision_config=SimpleNamespace( + patch_size=14, spatial_merge_size=2, temporal_patch_size=2 + ) + ) + info._get_image_max_pixels = lambda: token_budget * 2 * 28**2 + canvas = info.get_image_size_with_most_features() + assert canvas.width % (28 * expand) == 0 + assert canvas.height % (28 * expand) == 0 + assert canvas.width // processor.patch_size <= 8192 + assert canvas.height // processor.patch_size <= 8192 + assert canvas.width * canvas.height // 28**2 == expected_tokens + assert info.get_max_image_tokens() == expected_tokens + assert ( + info.get_num_image_tokens(image_width=3082, image_height=2048) <= token_budget + ) + if token_budget == 8000 and expand == 1: + assert (canvas.width, canvas.height) == (2800, 2240) + assert canvas.width * canvas.height // 28**2 > 7957 + + +def test_glm5next_image_pixel_budget_respects_token_override(): + from vllm.models.glm5next.common.multimodal import Glm5NextProcessingInfo + + processor = SimpleNamespace(patch_size=14, merge_size=2, temporal_patch_size=2) + info = Glm5NextProcessingInfo.__new__(Glm5NextProcessingInfo) + info.get_hf_processor = lambda **kwargs: SimpleNamespace(image_processor=processor) + info.ctx = SimpleNamespace( + get_merged_mm_kwargs=lambda kwargs: {"max_image_tokens": 2048} + ) + assert info._get_image_max_pixels() == 2048 * 2 * 28**2 + + def test_glm5next_processor_counts_video_only_tokens() -> None: class FakeVideoProcessor: merge_size = 2 diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index e5df32783a15..341c9359a446 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -50,12 +50,16 @@ def test_workspace_lanes_do_not_alias_and_restore_context(monkeypatch) -> None: ) assert manager._current_workspaces == [None, None, None, None] + assert manager.available_bytes() == 0 (target,) = manager.get_simultaneous(((512,), torch.uint8)) + assert manager.available_bytes() == 512 with workspace.use_workspace_lane(1): (draft,) = manager.get_simultaneous(((256,), torch.uint8)) (draft_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + assert manager.available_bytes() == 256 (target_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + assert manager.available_bytes() == 512 assert manager._current_workspaces[0].numel() == 512 # type: ignore[union-attr] assert manager._current_workspaces[1].numel() == 256 # type: ignore[union-attr] diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py index aa2db6da23a4..a0fe790a4993 100644 --- a/vllm/models/glm5next/common/multimodal.py +++ b/vllm/models/glm5next/common/multimodal.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from functools import cached_property, partial +from math import isqrt, prod import numpy as np import torch @@ -24,6 +25,7 @@ MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope @@ -44,6 +46,13 @@ from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.multimodal.parse import ImageSize, MultiModalDataItems from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +_VISION_PROJECTION_CHUNK_ROWS = 4096 +_VISION_ROPE_MAX_POSITION = 8192 class Glm5NextVisionPatchEmbed(nn.Module): @@ -169,6 +178,10 @@ def __init__( prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + self._borrow_qkv_workspace = ( + isinstance(self.qkv.quant_method, UnquantizedLinearMethod) + and self.attn.attn_backend == AttentionBackendEnum.FLASH_ATTN + ) def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]: seq_len, bs, _ = qkv.shape @@ -190,6 +203,25 @@ def forward( rotary_pos_emb_sin: torch.Tensor, max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: + q, k, v = self._project_qkv(x, rotary_pos_emb_cos, rotary_pos_emb_sin) + context_layer = self.attn( + query=q, + key=k, + value=v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + del q, k, v + context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() + output, _ = self.proj(context_layer) + return output + + def _project_qkv( + self, + x: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor | None, + rotary_pos_emb_sin: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: x, _ = self.qkv(x) q, k, v = self.split_qkv(x) @@ -204,10 +236,14 @@ def forward( self.k_norm.weight, self.q_norm.variance_epsilon, ) + del q_flat, k_flat q = q.view(q_shape) k = k.view(k_shape) q, k, v = (rearrange(t, "s b ... -> b s ...").contiguous() for t in (q, k, v)) + # Contiguous attention inputs own their storage; the packed projection + # and norm input copies must not stay live through attention/output GEMM. + del x if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: qk_concat = torch.cat([q, k], dim=0) qk_rotated = self.apply_rotary_emb( @@ -216,18 +252,64 @@ def forward( rotary_pos_emb_sin, ) q, k = torch.chunk(qk_rotated, 2, dim=0) + del qk_concat, qk_rotated - context_layer = self.attn( + return q, k, v + + def context_with_chunked_projection( + self, + x: torch.Tensor, + input_norm: nn.Module, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None, + ) -> torch.Tensor: + """Bound projection scratch without partitioning image attention. + + Normalization, QKV projection and rotary embedding are token-local. + Attention still consumes the complete Q/K/V and original image/frame + sequence boundaries. No cache or tensor dtype changes are involved. + """ + rows, batch_size, _ = x.shape + shape = ( + batch_size, + rows, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ) + manager = ( + current_workspace_manager() + if self._borrow_qkv_workspace and is_workspace_manager_initialized() + else None + ) + qkv_bytes = 3 * ((prod(shape) * x.element_size() + 255) // 256) * 256 + if manager is not None and manager.available_bytes() >= qkv_bytes: + # The BF16 projections, rotary operation and FlashAttention do not + # borrow worker scratch. These Q/K/V views expire before the output + # projection; the text model executes after the encoder returns. + q, k, v = manager.get_simultaneous(*((shape, x.dtype),) * 3) + else: + q, k, v = (x.new_empty(shape) for _ in range(3)) + for start in range(0, rows, _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, rows) + cq, ck, cv = self._project_qkv( + input_norm(x[start:stop]), + rotary_pos_emb_cos[start:stop], + rotary_pos_emb_sin[start:stop], + ) + q[:, start:stop].copy_(cq) + k[:, start:stop].copy_(ck) + v[:, start:stop].copy_(cv) + del cq, ck, cv + context = self.attn( query=q, key=k, value=v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, ) - context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() - - output, _ = self.proj(context_layer) - return output + return rearrange(context, "b s h d -> s b (h d)").contiguous() class Glm5NextVisionBlock(nn.Module): @@ -270,6 +352,24 @@ def forward( rotary_pos_emb_sin: torch.Tensor, max_seqlen: int | None = None, ) -> torch.Tensor: + if x.shape[0] > _VISION_PROJECTION_CHUNK_ROWS: + context = self.attn.context_with_chunked_projection( + x, + self.norm1, + cu_seqlens, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + max_seqlen, + ) + for start in range(0, x.shape[0], _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, x.shape[0]) + x_attn, _ = self.attn.proj(context[start:stop]) + normalized, residual = self.norm2(x[start:stop], residual=x_attn) + # Each token's residual is consumed before its owned input + # storage is reused. Attention has already read the full image. + x[start:stop].copy_(residual + self.mlp(normalized)) + del normalized, residual, x_attn + return x x_attn = self.attn( self.norm1(x), cu_seqlens=cu_seqlens, @@ -330,6 +430,7 @@ def forward(self, x: torch.Tensor): x = self.extra_activation_func(self.post_projection_norm(x)) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) + del gate_up x, _ = self.down_proj(x) return x @@ -390,7 +491,7 @@ def __init__( head_dim = self.hidden_size // self.num_heads self.rotary_pos_emb = get_rope( head_size=head_dim, - max_position=8192, + max_position=_VISION_ROPE_MAX_POSITION, is_neox_style=True, rope_parameters={"partial_rotary_factor": 0.5}, ) @@ -599,6 +700,38 @@ def forward( max_seqlen=max_seqlen, ) + # The GLM merger preserves the number of elements in each spatial + # group. Reuse the owned transformer output only after that group's + # normalization, convolution and merger have consumed its values. + if ( + x.shape[0] > _VISION_PROJECTION_CHUNK_ROWS + and self.out_hidden_size == self.spatial_merge_size**2 * x.shape[-1] + ): + hidden_size = x.shape[-1] + for start in range(0, x.shape[0], _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, x.shape[0]) + x[start:stop].copy_(self.post_layernorm(x[start:stop])) + x = x.view(-1, self.out_hidden_size) + groups_per_chunk = max( + 1, _VISION_PROJECTION_CHUNK_ROWS // self.spatial_merge_size**2 + ) + for start in range(0, x.shape[0], groups_per_chunk): + stop = min(start + groups_per_chunk, x.shape[0]) + patches = ( + x[start:stop] + .view( + -1, + self.spatial_merge_size, + self.spatial_merge_size, + hidden_size, + ) + .permute(0, 3, 1, 2) + ) + merged = self.downsample(patches).view(-1, self.out_hidden_size) + x[start:stop].copy_(self.merger(merged)) + del patches, merged + return x + # adapter x = self.post_layernorm(x) x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1]) @@ -652,7 +785,44 @@ def _get_image_max_pixels(self) -> int: mm_kwargs = self.ctx.get_merged_mm_kwargs({}) if (override := mm_kwargs.get("max_pixels")) is not None: return int(override) - return self._processor_pixel_budget(self.get_hf_processor().image_processor)[1] + processor = self.get_hf_processor().image_processor + if (tokens := mm_kwargs.get("max_image_tokens")) is not None: + return ( + int(tokens) + * processor.temporal_patch_size + * (processor.patch_size * processor.merge_size) ** 2 + ) + return self._processor_pixel_budget(processor)[1] + + def get_image_size_with_most_features(self) -> ImageSize: + """Maximize aligned canvas area within the token and vision RoPE limits. + + A square canvas can underfill the budget: 89 by 89 merged patches + produce 7921 features, while a valid 80 by 100 canvas produces 8000. + Encoder admission and profiling must cover the rectangular maximum. + """ + processor = self.get_hf_processor().image_processor + factor = ( + processor.patch_size * processor.merge_size * processor.patch_expand_factor + ) + cells = self._get_image_max_pixels() // ( + processor.temporal_patch_size * factor**2 + ) + if cells < 1: + raise ValueError("Image budget must fit at least one aligned canvas cell") + max_side = _VISION_ROPE_MAX_POSITION // ( + processor.merge_size * processor.patch_expand_factor + ) + if max_side < 1: + raise ValueError("One aligned canvas cell exceeds the vision RoPE grid") + width = height = 1 + # Include every feasible shorter side. Maximum area prevents memory + # underprofiling; equal-area ties prefer the least elongated canvas. + for candidate_height in range(1, min(isqrt(cells), max_side) + 1): + candidate_width = min(cells // candidate_height, max_side) + if candidate_width * candidate_height >= width * height: + width, height = candidate_width, candidate_height + return ImageSize(width=width * factor, height=height * factor) def _get_video_max_pixels(self) -> int: mm_kwargs = self.ctx.get_merged_mm_kwargs({}) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index ddde9f8b1df5..b44225340491 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -178,6 +178,23 @@ def is_locked(self) -> bool: """Check if workspace is locked.""" return self._locked + def available_bytes(self) -> int: + """Capacity of the active execution slot, without allocating storage. + + A borrower must ensure that no nested or concurrent operation consumes + this slot while its views are live. + """ + ubatch_id = dbo_current_ubatch_id() + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + return self._workspace_size_bytes( + self._current_workspaces[ubatch_id * self._num_lanes + lane] + ) + def get_simultaneous( self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] ) -> list[torch.Tensor]: From d9021bc5cfd576d2dccc423138dcb136586ae81c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 11 Sep 2026 23:47:14 +0000 Subject: [PATCH 053/365] Gather GLM DCP packed cache in place in caller-owned receive storage Pack the local contribution into the rank slice defined by the active padded token count and retain the complete receive capacity. Remove the redundant local workspace reservation without changing record bytes, cache precision or the full-CKV cutover. Validate rank-offset aliasing, zero padding, unused tails and native record widths; two-GPU PyNCCL, Torch distributed and coordinator transports preserve bytes exactly. Co-authored-by: Codex (cherry picked from commit 0fb58627fa8fc9f4ead0070adab0700c4befd623) Signed-off-by: Martin Vit --- tests/distributed/test_pynccl.py | 54 +++++++++++++++++++ .../v1/attention/test_b12x_sparse_mla_api.py | 44 +++++++++------ .../attention/backends/mla/b12x_mla_sparse.py | 22 +++----- 3 files changed, 90 insertions(+), 30 deletions(-) diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 91dccd579869..80d180f4db25 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from types import SimpleNamespace import multiprocess as mp import numpy as np @@ -200,6 +201,59 @@ def test_pynccl_all_gather(): distributed_run(all_gather_worker_fn, 2) +@worker_fn_wrapper +def ckv_inplace_all_gather_worker_fn(): + from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( + _dcp_all_gather_current_stream, + ) + + world = get_world_group() + comm = PyNcclCommunicator(world.cpu_group, device=world.device) + rank, size = comm.rank, comm.world_size + with ensure_current_vllm_config(): + ensure_model_parallel_initialized(size, 1) + tp = get_tp_group() + storage = torch.empty((size * 8192, 528), dtype=torch.uint8, device=world.device) + for transport in ("pynccl", "torch", "coordinator"): + group = SimpleNamespace( + world_size=size, + rank_in_group=rank, + device_communicator=SimpleNamespace( + pynccl_comm=comm if transport == "pynccl" else None + ), + device_group=world.device_group if transport == "torch" else None, + all_gather=tp.all_gather, + ) + for padded in (1024, 2048, 8192, 1024): + storage.fill_(255) + expected_parts = [] + for source_rank in range(size): + part = ( + ( + torch.arange(padded * 528, device=world.device).remainder(251) + + source_rank + ) + .to(torch.uint8) + .view(padded, 528) + ) + part[padded - source_rank * 4 - 1 :].zero_() + expected_parts.append(part) + send = storage[rank * padded : (rank + 1) * padded] + send.copy_(expected_parts[rank]) + receive = storage[: size * padded] + _dcp_all_gather_current_stream(group, send.view(-1), receive.view(-1)) + torch.accelerator.synchronize() + assert torch.equal(receive, torch.cat(expected_parts)) + assert torch.all(storage[size * padded :] == 255) + + +@pytest.mark.parametrize("world_size", [2, 4]) +def test_ckv_inplace_all_gather_preserves_padded_rank_records(world_size): + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Need {world_size} GPUs") + distributed_run(ckv_inplace_all_gather_worker_fn, world_size) + + @worker_fn_wrapper def cuda_communicator_all_gather_dim_worker_fn(): with ensure_current_vllm_config(): diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 534f971261d3..a7865d140610 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -803,16 +803,18 @@ def test_b12x_glm5_next_full_ckv_workspaces_follow_cache_format( impl._cache_record_bytes = record_bytes specs = impl._workspace_specs(input_num_heads=8, include_ckv=True) - assert specs[-2:] == ( - ((128, record_bytes), torch.uint8), - ((512, record_bytes), torch.uint8), - ) + assert len(specs) == 3 + assert specs[-1] == ((512, record_bytes), torch.uint8) @pytest.mark.parametrize("record_bytes", [528, 304]) +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("padded_tokens", [2, 4]) def test_b12x_glm5_next_full_ckv_gather_preserves_native_records( monkeypatch: pytest.MonkeyPatch, record_bytes: int, + rank: int, + padded_tokens: int, ) -> None: impl = object.__new__(B12xMLASparseImpl) impl._kernel_page_size = 2 @@ -826,37 +828,48 @@ def test_b12x_glm5_next_full_ckv_gather_preserves_native_records( .to(torch.uint8) .view(2, 2, record_bytes) ) - local_buffer = torch.full((4, record_bytes), 255, dtype=torch.uint8) - gathered_buffer = torch.empty((8, record_bytes), dtype=torch.uint8) + gathered_buffer = torch.full((8, record_bytes), 255, dtype=torch.uint8) + local_tokens = padded_tokens - 1 metadata = SimpleNamespace( - num_actual_tokens=3, - dcp_local_total_tokens=3, - dcp_padded_total_tokens=4, - dcp_local_cu_seq_lens=torch.tensor([0, 3], dtype=torch.int32), + num_actual_tokens=local_tokens, + dcp_local_total_tokens=local_tokens, + dcp_padded_total_tokens=padded_tokens, + dcp_local_cu_seq_lens=torch.tensor([0, local_tokens], dtype=torch.int32), block_table=torch.tensor([[0, 1]], dtype=torch.int32), num_reqs=1, ) def fake_cp_gather_cache(**kwargs: Any) -> None: - kwargs["dst"].copy_(kwargs["src_cache"].view(-1, record_bytes)[:3]) + kwargs["dst"].copy_(kwargs["src_cache"].view(-1, record_bytes)[:local_tokens]) def fake_all_gather(_group: Any, src: torch.Tensor, dst: torch.Tensor) -> None: + assert src.data_ptr() == dst.data_ptr() + rank * padded_tokens * record_bytes dst.copy_(src.repeat(2)) monkeypatch.setattr(b12x_mla_sparse.ops, "cp_gather_cache", fake_cp_gather_cache) monkeypatch.setattr( b12x_mla_sparse, "_dcp_all_gather_current_stream", fake_all_gather ) - monkeypatch.setattr(b12x_mla_sparse, "get_dcp_group", lambda: object()) + monkeypatch.setattr( + b12x_mla_sparse, + "get_dcp_group", + lambda: SimpleNamespace(rank_in_group=rank, world_size=2), + ) - gathered = impl._gather_full_ckv(kv_cache, metadata, local_buffer, gathered_buffer) + gathered = impl._gather_full_ckv(kv_cache, metadata, gathered_buffer) expected_rank = torch.cat( - (kv_cache.view(-1, record_bytes)[:3], torch.zeros((1, record_bytes))), + ( + kv_cache.view(-1, record_bytes)[:local_tokens], + torch.zeros((1, record_bytes), dtype=torch.uint8), + ), dim=0, ) assert gathered.shape == (4, 2, record_bytes) - assert torch.equal(gathered.view(-1, record_bytes), expected_rank.repeat(2, 1)) + assert torch.equal( + gathered.view(-1, record_bytes)[: 2 * padded_tokens], expected_rank.repeat(2, 1) + ) + assert torch.all(gathered.view(-1, record_bytes)[2 * padded_tokens :] == 255) def test_b12x_glm5_next_full_ckv_gather_rejects_wrong_record_width() -> None: @@ -872,7 +885,6 @@ def test_b12x_glm5_next_full_ckv_gather_rejects_wrong_record_width() -> None: impl._gather_full_ckv( torch.empty((2, 2, 528), dtype=torch.uint8), metadata, - torch.empty((4, 304), dtype=torch.uint8), torch.empty((8, 304), dtype=torch.uint8), ) diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index fe264953c2dc..00dbc252d9ae 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1592,10 +1592,6 @@ def _workspace_specs( ) scratch_spec = ((self._scratch_nbytes,), torch.uint8) ckv_specs = ( - ( - (self._ckv_local_capacity, self._cache_record_bytes), - torch.uint8, - ), ( ( self.dcp_world_size * self._ckv_local_capacity, @@ -1804,7 +1800,6 @@ def _gather_full_ckv( self, kv_cache: torch.Tensor, attn_metadata: B12xMLASparseMetadata, - local_buffer: torch.Tensor, gathered_buffer: torch.Tensor, ) -> torch.Tensor: if not self.uses_full_ckv_dcp(attn_metadata, attn_metadata.num_actual_tokens): @@ -1819,22 +1814,22 @@ def _gather_full_ckv( f"{self._cache_record_bytes}-byte records; " f"got shape={tuple(kv_cache.shape)}, stride={kv_cache.stride()}" ) - expected_local_shape = ( - self._ckv_local_capacity, - self._cache_record_bytes, - ) expected_gathered_shape = ( self.dcp_world_size * self._ckv_local_capacity, self._cache_record_bytes, ) - if tuple(local_buffer.shape) != expected_local_shape: - raise RuntimeError("CKV local workspace has an invalid shape") if tuple(gathered_buffer.shape) != expected_gathered_shape: raise RuntimeError("CKV gathered workspace has an invalid shape") assert attn_metadata.dcp_local_cu_seq_lens is not None local_tokens = attn_metadata.dcp_local_total_tokens padded_tokens = attn_metadata.dcp_padded_total_tokens + group = get_dcp_group() + # NCCL in-place all-gather reads each rank's contribution from its + # receive slice. The stride is this call's padded count, not capacity. + local_buffer = gathered_buffer.narrow( + 0, group.rank_in_group * padded_tokens, padded_tokens + ) if local_tokens: ops.cp_gather_cache( src_cache=kv_cache, @@ -1846,7 +1841,7 @@ def _gather_full_ckv( if local_tokens < padded_tokens: local_buffer[local_tokens:padded_tokens].zero_() _dcp_all_gather_current_stream( - get_dcp_group(), + group, local_buffer[:padded_tokens].view(-1), gathered_buffer[: self.dcp_world_size * padded_tokens].view(-1), ) @@ -1918,11 +1913,10 @@ def forward_mqa( topk_indices = self.topk_indices_buffer[:num_tokens] kv_cache_for_run = kv_c_and_k_pe_cache if use_ckv_gather: - local_buffer, gathered_buffer = workspaces[2:] + gathered_buffer = workspaces[2] kv_cache_for_run = self._gather_full_ckv( kv_c_and_k_pe_cache, attn_metadata, - local_buffer, gathered_buffer, ) assert attn_metadata.ckv_selected_indices is not None From 7c7c9868a3370faee7dd907b638b4b9db29debaa Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 12 Sep 2026 01:29:01 +0000 Subject: [PATCH 054/365] Reuse caller-owned sparse attention scratch for GLM DCP query gathering Stage rank-major query shards in kernel scratch and gather in place before copying into the established final query buffer. The execution stream orders the copy before attention reuses scratch. Preserve direct-transport selection and decode collectives. Avoid concatenation for zero-width RoPE and release the intermediate projection after its input owner has been established. No weights, cache precision, attention arithmetic or selected-token capacity change. Validation: ten focused dispatch, layout and reservation tests; two-GPU PyNCCL/Torch bitwise parity and live-input CUDA graph replay; formatting, lint and typing hooks. The broader suite has 176 passes and one metadata pointer-sharing failure reproduced with the unmodified control files; it is retained separately from this query-memory change. Serving capacity qualification is recorded in the deployment artifacts. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 2e6ff184416ab800abeb41176d33a2b7aa551354) --- .../test_dcp_direct_a2a_lse_reduce.py | 37 ++++++++++ tests/distributed/test_pynccl.py | 69 ++++++++++++++++++ .../v1/attention/test_b12x_sparse_mla_api.py | 72 +++++++++++++++++++ .../layers/attention/mla_attention.py | 10 ++- .../attention/backends/mla/b12x_mla_sparse.py | 44 ++++++++++++ vllm/v1/attention/ops/dcp.py | 8 ++- 6 files changed, 238 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index 2db898805cbf..569df823b28a 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -524,6 +524,43 @@ def test_dcp_workspace_covers_parallel_drafting(): assert dcp.get_dcp_workspace_max_num_tokens(config) == 28 +@pytest.mark.parametrize("direct_capacity", [None, 4]) +@pytest.mark.parametrize("rows", [2, 8]) +def test_mla_query_scratch_fallback_preserves_direct_transport_priority( + monkeypatch, direct_capacity, rows +): + group = MagicMock(world_size=2) + direct = ( + None if direct_capacity is None else MagicMock(max_num_tokens=direct_capacity) + ) + fallback = MagicMock(return_value=torch.empty((rows, 4, 8))) + monkeypatch.setattr(dcp, "get_dcp_group", lambda: group) + monkeypatch.setattr(dcp, "get_direct_dcp_q_gather_workspace", lambda *args: direct) + manager = dcp.MLADCPManager( + vllm_config=_manager_config(dcp_comm_backend="ag_rs"), + device=torch.device("cpu"), + num_heads=2, + query_head_dim=8, + output_head_dim=4, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + is_lse_base_on_e=True, + use_pcp=False, + query_gather_fallback=fallback, + ) + query = torch.empty((rows, 2, 8), dtype=torch.bfloat16) + actual = manager.query_gather(query) + if direct is not None and rows <= direct_capacity: + direct.gather.assert_called_once_with(query) + fallback.assert_not_called() + assert actual is direct.gather.return_value + else: + fallback.assert_called_once_with(query) + assert actual is fallback.return_value + group.all_gather.assert_not_called() + + def test_mla_dcp_manager_selects_pcp_combine(monkeypatch): import vllm.v1.attention.ops.dcp as dcp_manager diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 80d180f4db25..1e4e2ee71078 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -254,6 +254,75 @@ def test_ckv_inplace_all_gather_preserves_padded_rank_records(world_size): distributed_run(ckv_inplace_all_gather_worker_fn, world_size) +@worker_fn_wrapper +def glm_dcp_query_scratch_worker_fn(): + from vllm.v1.attention.backends.mla import b12x_mla_sparse + from vllm.v1.worker.workspace import WorkspaceManager + + world = get_world_group() + comm = PyNcclCommunicator(world.cpu_group, device=world.device) + rank, size = comm.rank, comm.world_size + group = SimpleNamespace( + world_size=size, + rank_in_group=rank, + device_communicator=SimpleNamespace(pynccl_comm=comm), + device_group=world.device_group, + ) + manager = WorkspaceManager(world.device) + impl = object.__new__(b12x_mla_sparse.B12xMLASparseImpl) + impl._is_glm_next = True + impl._decode_max_rows = 16 + impl._max_tokens = 512 + impl._input_num_heads = size * 4 + impl._q_head_dim = 512 + impl._scratch_nbytes = 512 * impl._input_num_heads * 512 * 2 + impl._decode_plan = object() + impl._ckv_local_capacity = 0 + impl._cache_record_bytes = 528 + impl.dcp_world_size = size + b12x_mla_sparse.current_workspace_manager = lambda: manager + b12x_mla_sparse.get_dcp_group = lambda: group + q_buffer, scratch = impl._borrow_workspaces() + manager.lock() + for transport in ("pynccl", "torch"): + group.device_communicator.pynccl_comm = comm if transport == "pynccl" else None + for rows in (32, 512, 64): + base = torch.arange(rows * 4 * 512, device=world.device).remainder(127) + parts = [ + (base + source).to(torch.bfloat16).view(4, rows, 512).transpose(0, 1) + for source in range(size) + ] + actual = impl.gather_dcp_query(parts[rank]) + assert actual.data_ptr() == q_buffer.data_ptr() + expected = torch.cat(parts, dim=1) + assert torch.equal(actual, expected) + scratch.fill_(255) + assert torch.equal(actual, expected) + + group.device_communicator.pynccl_comm = comm + local = torch.empty((4, 64, 512), dtype=torch.bfloat16, device=world.device) + local.fill_(rank) + query = local.transpose(0, 1) + impl.gather_dcp_query(query) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = impl.gather_dcp_query(query) + for iteration in range(3): + local.fill_(rank + iteration * 8) + graph.replay() + expected = torch.cat( + [torch.full_like(query, source + iteration * 8) for source in range(size)], + dim=1, + ) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Need two GPUs") +def test_glm_dcp_query_scratch_preserves_values_and_graph_inputs(): + distributed_run(glm_dcp_query_scratch_worker_fn, 2) + + @worker_fn_wrapper def cuda_communicator_all_gather_dim_worker_fn(): with ensure_current_vllm_config(): diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index a7865d140610..09155066f9f2 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -807,6 +807,78 @@ def test_b12x_glm5_next_full_ckv_workspaces_follow_cache_format( assert specs[-1] == ((512, record_bytes), torch.uint8) +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("transposed", [False, True]) +def test_glm_dcp_prefill_query_reuses_scratch_without_aliasing_output( + monkeypatch, world_size: int, transposed: bool +) -> None: + from vllm.v1.worker.workspace import WorkspaceManager + + manager = WorkspaceManager(torch.device("cpu")) + impl = object.__new__(B12xMLASparseImpl) + impl._is_glm_next = True + impl._decode_max_rows = 4 + impl._max_tokens = 32 + impl._input_num_heads = world_size * 4 + impl._q_head_dim = 8 + impl._scratch_nbytes = 32 * impl._input_num_heads * 8 * 2 + impl._decode_plan = object() + impl._ckv_local_capacity = 0 + impl._cache_record_bytes = 528 + impl.dcp_world_size = world_size + monkeypatch.setattr(b12x_mla_sparse, "current_workspace_manager", lambda: manager) + q_buffer, scratch = impl._borrow_workspaces() + manager.lock() + for rank in range(world_size): + group = SimpleNamespace(world_size=world_size, rank_in_group=rank) + monkeypatch.setattr(b12x_mla_sparse, "get_dcp_group", lambda group=group: group) + for rows in (19, 7, 23): + local_queries = [] + for source in range(world_size): + values = (torch.arange(rows * 4 * 8) % 127 + source).to(torch.bfloat16) + local_queries.append( + values.view(4, rows, 8).transpose(0, 1) + if transposed + else values.view(rows, 4, 8) + ) + + def gather( + _group, send, receive, rank=rank, rows=rows, local_queries=local_queries + ): + assert send.data_ptr() == receive.data_ptr() + rank * send.nbytes + assert torch.equal(send.view(rows, 4, 8), local_queries[rank]) + receive.view(world_size, rows, 4, 8).copy_(torch.stack(local_queries)) + + monkeypatch.setattr( + b12x_mla_sparse, "_dcp_all_gather_current_stream", gather + ) + actual = impl.gather_dcp_query(local_queries[rank]) + assert actual.data_ptr() == q_buffer.data_ptr() + expected = torch.cat(local_queries, dim=1) + assert torch.equal(actual, expected) + scratch.fill_(255) + assert torch.equal(actual, expected) + + +def test_glm_dcp_decode_keeps_the_transport_collective(monkeypatch) -> None: + impl = object.__new__(B12xMLASparseImpl) + impl._is_glm_next = True + impl._decode_max_rows = 16 + query = torch.empty((16, 4, 8), dtype=torch.bfloat16) + expected = torch.empty((16, 8, 8), dtype=torch.bfloat16) + calls = [] + + def gather(value, dim): + calls.append((value, dim)) + return expected + + monkeypatch.setattr( + b12x_mla_sparse, "get_dcp_group", lambda: SimpleNamespace(all_gather=gather) + ) + assert impl.gather_dcp_query(query) is expected + assert calls == [(query, 1)] + + @pytest.mark.parametrize("record_bytes", [528, 304]) @pytest.mark.parametrize("rank", [0, 1]) @pytest.mark.parametrize("padded_tokens", [2, 4]) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index d5128e2a8d1b..c9ec73b57f84 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -861,6 +861,7 @@ def __init__( padded_num_heads=self.q_pad_num_heads, is_lse_base_on_e=self.impl.lse_base_on_e, use_pcp=self.use_pcp, + query_gather_fallback=getattr(self.impl, "gather_dcp_query", None), ) self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled() @@ -1329,6 +1330,9 @@ def forward_impl( ) else: mqa_q = (mqa_ql_nope, mqa_q_pe) + # The input tuple owns the projection until query preparation + # finishes; no extra reference must retain it through attention. + del mqa_ql_nope # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused) if self.impl.dcp_world_size > 1: assert self.dcp_manager is not None @@ -1340,7 +1344,11 @@ def forward_impl( else: if isinstance(mqa_q, tuple): # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) - mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = ( + mqa_q[0] + if self.qk_rope_head_dim == 0 + else torch.cat(mqa_q, dim=-1) + ) if not qrep_decode and not full_ckv_dcp: assert self.dcp_manager.query_gather is not None mqa_q = self.dcp_manager.query_gather(mqa_q) diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 00dbc252d9ae..054243f8303c 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1364,6 +1364,16 @@ def _set_kernel_page_size(self, kernel_page_size: int) -> None: self._scratch_nbytes = max( int(plan.scratch_specs()[0].nbytes) for plan in self._plans.values() ) + if self._is_glm_next and self.dcp_world_size > 1: + # Query gathering finishes before attention consumes scratch. Its + # rank-major receive storage can therefore reuse the kernel arena. + self._scratch_nbytes = max( + self._scratch_nbytes, + self._max_tokens + * self._input_num_heads + * self._q_head_dim + * torch.bfloat16.itemsize, + ) self._ckv_local_capacity = _round_up_ckv_rank_tokens( self._ckv_capacity_tokens, page_size=kernel_page_size, @@ -1633,6 +1643,40 @@ def supports_fused_mla_query_output( and self._q_head_dim == 576 ) + def gather_dcp_query(self, query: torch.Tensor) -> torch.Tensor: + """Gather GLM prefill heads into caller-owned attention query storage. + + The rank-major receive view borrows kernel scratch only until the + layout copy completes on the execution stream. The attention call + rebinds that scratch while retaining the final query at the same address. + Decode and transport-specific all-gathers keep their established path. + """ + group = get_dcp_group() + if not self._is_glm_next or query.shape[0] <= self._decode_max_rows: + return group.all_gather(query, dim=1) + communicator = getattr(group, "device_communicator", None) + b12x_comm = getattr(communicator, "b12x_ar_comm", None) + if ( + b12x_comm is not None + and not b12x_comm.disabled + and getattr(b12x_comm, "should_all_gather", None) is not None + and b12x_comm.should_all_gather(query, 1) + ): + return group.all_gather(query, dim=1) + + rows, local_heads, head_dim = query.shape + q_buffer, scratch = self._borrow_workspaces()[:2] + output = q_buffer[:rows] + receive = scratch[: output.numel() * output.element_size()].view(query.dtype) + receive = receive.view(self.dcp_world_size, rows, local_heads, head_dim) + local = receive[group.rank_in_group] + local.copy_(query) + _dcp_all_gather_current_stream(group, local.view(-1), receive.view(-1)) + output.view(rows, self.dcp_world_size, local_heads, head_dim).copy_( + receive.movedim(0, 1) + ) + return output + def get_fused_mla_query_output( self, num_tokens: int, diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index d46e37215d32..6313f43e571b 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -1467,6 +1467,7 @@ def __init__( padded_num_heads: int | None, is_lse_base_on_e: bool, use_pcp: bool, + query_gather_fallback: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> None: parallel_config = vllm_config.parallel_config self.group = get_dcp_group() @@ -1475,6 +1476,7 @@ def __init__( self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config) self.use_a2a = parallel_config.dcp_comm_backend == "a2a" self.padded_num_heads = padded_num_heads + self._query_gather_fallback = query_gather_fallback self.combine = self._init_combine( num_heads, @@ -1595,7 +1597,11 @@ def _init_query_gather( return self._gather_query def _gather_query(self, query: torch.Tensor) -> torch.Tensor: - query = self.group.all_gather(query, dim=1) + query = ( + self.group.all_gather(query, dim=1) + if self._query_gather_fallback is None + else self._query_gather_fallback(query) + ) if self.padded_num_heads is not None: query = reserve_query_head_storage(query, self.padded_num_heads) return query From 7285bd3731c37775c3d48caaf39a9e5f3d82cb4c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 12 Sep 2026 01:41:35 +0000 Subject: [PATCH 055/365] perf(attention): reserve full-CKV scratch for local query heads Use the full-CKV plan layout when its cache receive buffer is live. Global-head query gathering and decode split-K intermediates are not concurrent with that plan. Other paths retain their shared maximum reservation. Thirteen focused tests pass, including real two-rank collectives and CUDA graph replay; serving capacity is not yet qualified. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 0114a09bee647b24662cfc99fc61ce42394429f1) --- tests/v1/attention/test_b12x_sparse_mla_api.py | 6 ++++++ vllm/v1/attention/backends/mla/b12x_mla_sparse.py | 10 +++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 09155066f9f2..b10d5da158ad 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -801,9 +801,15 @@ def test_b12x_glm5_next_full_ckv_workspaces_follow_cache_format( impl._ckv_local_capacity = 128 impl.dcp_world_size = 4 impl._cache_record_bytes = record_bytes + impl._plans = { + ("ckv_extend", 32): SimpleNamespace( + scratch_specs=lambda: (SimpleNamespace(nbytes=8),) + ) + } specs = impl._workspace_specs(input_num_heads=8, include_ckv=True) assert len(specs) == 3 + assert specs[1] == ((8,), torch.uint8) assert specs[-1] == ((512, record_bytes), torch.uint8) diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 054243f8303c..6eb5dc44e7d0 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1600,7 +1600,15 @@ def _workspace_specs( (self._max_tokens, input_num_heads, self._q_head_dim), torch.bfloat16, ) - scratch_spec = ((self._scratch_nbytes,), torch.uint8) + # Full-CKV attention has only local query heads and does not gather + # queries. Its cache receive buffer need not coexist with global-head + # query-gather scratch or the decode plan's split-K intermediates. + scratch_nbytes = ( + int(self._plan(("ckv_extend", self._max_tokens)).scratch_specs()[0].nbytes) + if include_ckv + else self._scratch_nbytes + ) + scratch_spec = ((scratch_nbytes,), torch.uint8) ckv_specs = ( ( ( From dfdbaeaff5a08a7d85b247d181fe93154f7c2684 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 12 Sep 2026 01:54:25 +0000 Subject: [PATCH 056/365] perf(attention): reuse GLM query storage for DCP output reduction Pack corrected prefill output into consumed caller-owned query storage and reduce into its rank slice with the documented in-place NCCL contract. Preserve head-major SM120 value-projection inputs, decode dispatch, symmetric-memory transports and PCP/A2A selection. Twenty-nine focused tests pass, including bitwise two-rank LSE/output parity and live-input CUDA graph replay. Broader A2A failures are retained for control verification; serving memory/performance qualification remains pending. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 317afc7c4bf4378ece58d1d4052afaddf089f44e) --- .../test_dcp_direct_a2a_lse_reduce.py | 27 +++++++ tests/distributed/test_pynccl.py | 63 ++++++++++++++- .../v1/attention/test_b12x_sparse_mla_api.py | 77 +++++++++++++++++++ .../layers/attention/mla_attention.py | 3 + .../attention/backends/mla/b12x_mla_sparse.py | 30 ++++++++ vllm/v1/attention/ops/dcp.py | 18 ++++- 6 files changed, 216 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index 569df823b28a..9612d6969e1a 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -583,6 +583,33 @@ def test_mla_dcp_manager_selects_pcp_combine(monkeypatch): assert manager.query_gather is None +@pytest.mark.parametrize( + "backend,use_pcp", [("ag_rs", False), ("ag_rs", True), ("a2a", False)] +) +def test_mla_output_workspace_hook_only_applies_to_head_reduce_scatter( + monkeypatch, backend, use_pcp +): + monkeypatch.setattr(dcp, "get_dcp_group", lambda: MagicMock(world_size=2)) + monkeypatch.setattr(dcp, "get_direct_dcp_q_gather_workspace", lambda *args: None) + monkeypatch.setattr(dcp, "get_direct_dcp_a2a_workspace", lambda *args: None) + callback = MagicMock() + manager = dcp.MLADCPManager( + vllm_config=_manager_config(dcp_comm_backend=backend), + device=torch.device("cpu"), + num_heads=2, + query_head_dim=8, + output_head_dim=4, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + is_lse_base_on_e=True, + use_pcp=use_pcp, + output_reduce_scatter=callback, + ) + expected = callback if backend == "ag_rs" and not use_pcp else None + assert manager.combine.keywords.get("output_reduce_scatter") is expected + + def test_dcp_chunk_workspace_alignment_covers_interleave(): from vllm.model_executor.layers.attention.mla_attention import ( align_mla_chunked_context_workspace_size, diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 1e4e2ee71078..0cf5eefbec12 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -317,9 +317,70 @@ def glm_dcp_query_scratch_worker_fn(): ) assert torch.equal(actual, expected) + from vllm.v1.attention.ops.dcp import cp_lse_ag_out_rs + + def gather_lse(value, dim): + assert dim == 0 + result = value.new_empty((size * value.shape[0], *value.shape[1:])) + comm.all_gather(result, value.contiguous()) + return result + + def reduce_output(value, dim): + packed = value.movedim(0, dim).contiguous() + result = packed.new_empty((packed.shape[0] // size, *packed.shape[1:])) + comm.reduce_scatter(result, packed) + return result.movedim(0, dim).contiguous() + + group.all_gather = gather_lse + group.reduce_scatter = reduce_output + for rows in (32, 512, 64): + shape = (rows, impl._input_num_heads, 512) + partial = ( + scratch[: rows * impl._input_num_heads * 512 * 2] + .view(torch.bfloat16) + .view(shape) + ) + source = torch.arange(partial.numel(), device=world.device).remainder(127) + source = (source + rank).to(torch.bfloat16).view(shape) + lse = torch.arange(rows * impl._input_num_heads, device=world.device) + lse = lse.float().remainder(17).view(shape[:2]) + rank + lse[::7].fill_(float("-inf")) + for base_e in (True, False): + expected, expected_lse = cp_lse_ag_out_rs( + source.clone(), + lse.clone(), + group, + return_lse=True, + is_lse_base_on_e=base_e, + ) + partial.copy_(source) + actual, actual_lse = cp_lse_ag_out_rs( + partial, + lse.clone(), + group, + return_lse=True, + is_lse_base_on_e=base_e, + output_reduce_scatter=impl.reduce_scatter_dcp_output, + ) + assert torch.equal(actual, expected) + assert torch.equal(actual_lse, expected_lse) + assert actual.transpose(0, 1).is_contiguous() + + partial.fill_(rank) + impl.reduce_scatter_dcp_output(partial) + torch.accelerator.synchronize() + output_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(output_graph): + actual = impl.reduce_scatter_dcp_output(partial) + for iteration in range(3): + partial.fill_(rank + iteration * 8) + output_graph.replay() + expected_sum = sum(range(size)) + size * iteration * 8 + assert torch.equal(actual, torch.full_like(actual, expected_sum)) + @pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Need two GPUs") -def test_glm_dcp_query_scratch_preserves_values_and_graph_inputs(): +def test_glm_dcp_query_scratch_preserves_query_output_and_graph_inputs(): distributed_run(glm_dcp_query_scratch_worker_fn, 2) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index b10d5da158ad..c9f9c678ef62 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -885,6 +885,83 @@ def gather(value, dim): assert calls == [(query, 1)] +@pytest.mark.parametrize("world_size", [2, 4]) +def test_glm_dcp_output_reuses_query_storage_with_head_major_rank_slices( + monkeypatch, world_size: int +) -> None: + from vllm.v1.worker.workspace import WorkspaceManager + + manager = WorkspaceManager(torch.device("cpu")) + impl = object.__new__(B12xMLASparseImpl) + impl._is_glm_next = True + impl._decode_max_rows = 4 + impl._max_tokens = 32 + impl._input_num_heads = world_size * 4 + impl._q_head_dim = 8 + impl._scratch_nbytes = 32 * impl._input_num_heads * 8 * 2 + impl._decode_plan = object() + impl._ckv_local_capacity = 0 + impl._cache_record_bytes = 528 + impl.dcp_world_size = world_size + monkeypatch.setattr(b12x_mla_sparse, "current_workspace_manager", lambda: manager) + monkeypatch.setattr(b12x_mla_sparse, "should_nccl_symm_mem_ag_rs", lambda: False) + q_buffer, scratch = impl._borrow_workspaces() + manager.lock() + for rank in range(world_size): + for rows in (19, 7, 23): + partial = ( + scratch[: rows * impl._input_num_heads * 8 * 2] + .view(torch.bfloat16) + .view(rows, impl._input_num_heads, 8) + ) + partial.copy_(torch.arange(partial.numel()).view_as(partial) % 127) + original = partial.clone() + + def reduce_scatter(local, packed, rank=rank, original=original): + assert packed.data_ptr() == q_buffer.data_ptr() + assert local.data_ptr() == packed.data_ptr() + rank * local.nbytes + assert local.is_contiguous() and packed.is_contiguous() + assert torch.equal(packed, original.transpose(0, 1)) + local.mul_(world_size) + + comm = SimpleNamespace(disabled=False, reduce_scatter=reduce_scatter) + group = SimpleNamespace( + rank_in_group=rank, + device_communicator=SimpleNamespace(pynccl_comm=comm), + ) + monkeypatch.setattr( + b12x_mla_sparse, "get_dcp_group", lambda group=group: group + ) + actual = impl.reduce_scatter_dcp_output(partial) + expected = original[:, rank * 4 : (rank + 1) * 4] * world_size + assert torch.equal(actual, expected) + assert torch.equal(partial, original) + assert actual.transpose(0, 1).is_contiguous() + scratch.fill_(255) + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize( + "rows,disabled,symmetric", [(4, False, False), (8, True, False), (8, False, True)] +) +def test_glm_dcp_output_preserves_decode_and_special_transport_paths( + monkeypatch, rows, disabled, symmetric +) -> None: + impl = object.__new__(B12xMLASparseImpl) + impl._is_glm_next = True + impl._decode_max_rows = 4 + comm = SimpleNamespace(disabled=disabled) + monkeypatch.setattr( + b12x_mla_sparse, + "get_dcp_group", + lambda: SimpleNamespace(device_communicator=SimpleNamespace(pynccl_comm=comm)), + ) + monkeypatch.setattr( + b12x_mla_sparse, "should_nccl_symm_mem_ag_rs", lambda: symmetric + ) + assert impl.reduce_scatter_dcp_output(torch.empty(rows, 8, 8)) is None + + @pytest.mark.parametrize("record_bytes", [528, 304]) @pytest.mark.parametrize("rank", [0, 1]) @pytest.mark.parametrize("padded_tokens", [2, 4]) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index c9ec73b57f84..afe70dde979f 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -862,6 +862,9 @@ def __init__( is_lse_base_on_e=self.impl.lse_base_on_e, use_pcp=self.use_pcp, query_gather_fallback=getattr(self.impl, "gather_dcp_query", None), + output_reduce_scatter=getattr( + self.impl, "reduce_scatter_dcp_output", None + ), ) self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled() diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 6eb5dc44e7d0..cfc00970bc58 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -15,6 +15,9 @@ from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group +from vllm.distributed.device_communicators.all_reduce_utils import ( + should_nccl_symm_mem_ag_rs, +) from vllm.logger import init_logger from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, @@ -1651,6 +1654,33 @@ def supports_fused_mla_query_output( and self._q_head_dim == 576 ) + def reduce_scatter_dcp_output(self, output: torch.Tensor) -> torch.Tensor | None: + """Reuse consumed query storage for the prefill output collective. + + The corrected attention output remains in disjoint kernel scratch. + Pack it into the query buffer in head-major order, then let NCCL write + each rank's reduction at its in-place receive offset. The resulting + head-major matrices stay live until the following value projection. + Decode and symmetric-memory transports retain their existing path. + """ + if not self._is_glm_next or output.shape[0] <= self._decode_max_rows: + return None + group = get_dcp_group() + communicator = getattr(group, "device_communicator", None) + comm = getattr(communicator, "pynccl_comm", None) + if comm is None or comm.disabled or should_nccl_symm_mem_ag_rs(): + return None + + rows, heads, head_dim = output.shape + q_buffer = self._borrow_workspaces()[0] + assert output.dtype == q_buffer.dtype + packed = q_buffer.view(-1)[: output.numel()].view(heads, rows, head_dim) + packed.copy_(output.transpose(0, 1)) + local_heads = heads // self.dcp_world_size + local = packed.narrow(0, group.rank_in_group * local_heads, local_heads) + comm.reduce_scatter(local, packed) + return local.transpose(0, 1) + def gather_dcp_query(self, query: torch.Tensor) -> torch.Tensor: """Gather GLM prefill heads into caller-owned attention query storage. diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 6313f43e571b..4b37e59c210b 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -466,6 +466,7 @@ def cp_lse_ag_out_rs( is_lse_base_on_e=True, seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, + output_reduce_scatter: Callable[[torch.Tensor], torch.Tensor | None] | None = None, ): """ cp_attn_out: [ B, H, D ] @@ -480,7 +481,12 @@ def cp_lse_ag_out_rs( seq_lens=seq_lens, query_start_loc=query_start_loc, ) - if current_platform.is_cuda() and current_platform.is_device_capability_family(120): + reduced = output_reduce_scatter(out) if output_reduce_scatter is not None else None + if reduced is not None: + out = reduced + elif current_platform.is_cuda() and current_platform.is_device_capability_family( + 120 + ): # Preserve the head-major collective output for the following MLA # value GEMM. Its batch matrices are disjoint, avoiding cuBLAS's # SM120/121 overlapping-stride read defect and a redundant transpose copy. @@ -1468,6 +1474,8 @@ def __init__( is_lse_base_on_e: bool, use_pcp: bool, query_gather_fallback: Callable[[torch.Tensor], torch.Tensor] | None = None, + output_reduce_scatter: Callable[[torch.Tensor], torch.Tensor | None] + | None = None, ) -> None: parallel_config = vllm_config.parallel_config self.group = get_dcp_group() @@ -1477,6 +1485,7 @@ def __init__( self.use_a2a = parallel_config.dcp_comm_backend == "a2a" self.padded_num_heads = padded_num_heads self._query_gather_fallback = query_gather_fallback + self._output_reduce_scatter = output_reduce_scatter self.combine = self._init_combine( num_heads, @@ -1537,6 +1546,13 @@ def _init_combine( if use_pcp else cp_lse_ag_out_rs ) + if combine_fn is cp_lse_ag_out_rs and self._output_reduce_scatter is not None: + return functools.partial( + cp_lse_ag_out_rs, + cp_group=self.group, + is_lse_base_on_e=is_lse_base_on_e, + output_reduce_scatter=self._output_reduce_scatter, + ) return functools.partial( combine_fn, cp_group=self.group, From e7376107b7cef3193f6047db1bead55194d3d5fd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 12 Sep 2026 02:13:12 +0000 Subject: [PATCH 057/365] test(attention): exercise full-size GLM DCP collective storage Run caller-owned query gathering and in-place output reduction with 2048 rows, 64 global heads and 512 features, covering the 128 MiB production collective size as well as smaller variable shapes. Two-rank PyNCCL/Torch query parity, allocating-output LSE parity and live-input CUDA graph replay pass. Runtime files and image hashes are unchanged. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit c9eb3e5859b21209501043f46968c1b4095a5fa8) --- tests/distributed/test_pynccl.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 0cf5eefbec12..a3dabd7f3bf1 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -272,10 +272,11 @@ def glm_dcp_query_scratch_worker_fn(): impl = object.__new__(b12x_mla_sparse.B12xMLASparseImpl) impl._is_glm_next = True impl._decode_max_rows = 16 - impl._max_tokens = 512 - impl._input_num_heads = size * 4 + impl._max_tokens = 2048 + local_heads = 32 + impl._input_num_heads = size * local_heads impl._q_head_dim = 512 - impl._scratch_nbytes = 512 * impl._input_num_heads * 512 * 2 + impl._scratch_nbytes = impl._max_tokens * impl._input_num_heads * 512 * 2 impl._decode_plan = object() impl._ckv_local_capacity = 0 impl._cache_record_bytes = 528 @@ -286,10 +287,15 @@ def glm_dcp_query_scratch_worker_fn(): manager.lock() for transport in ("pynccl", "torch"): group.device_communicator.pynccl_comm = comm if transport == "pynccl" else None - for rows in (32, 512, 64): - base = torch.arange(rows * 4 * 512, device=world.device).remainder(127) + for rows in (32, 2048, 64): + base = torch.arange( + rows * local_heads * 512, device=world.device + ).remainder(127) parts = [ - (base + source).to(torch.bfloat16).view(4, rows, 512).transpose(0, 1) + (base + source) + .to(torch.bfloat16) + .view(local_heads, rows, 512) + .transpose(0, 1) for source in range(size) ] actual = impl.gather_dcp_query(parts[rank]) @@ -300,7 +306,9 @@ def glm_dcp_query_scratch_worker_fn(): assert torch.equal(actual, expected) group.device_communicator.pynccl_comm = comm - local = torch.empty((4, 64, 512), dtype=torch.bfloat16, device=world.device) + local = torch.empty( + (local_heads, 64, 512), dtype=torch.bfloat16, device=world.device + ) local.fill_(rank) query = local.transpose(0, 1) impl.gather_dcp_query(query) @@ -333,7 +341,7 @@ def reduce_output(value, dim): group.all_gather = gather_lse group.reduce_scatter = reduce_output - for rows in (32, 512, 64): + for rows in (32, 2048, 64): shape = (rows, impl._input_num_heads, 512) partial = ( scratch[: rows * impl._input_num_heads * 512 * 2] From fd047ddfc7f923df1f411a689a8c4e56a6d7e87d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 11 Sep 2026 21:51:36 +0000 Subject: [PATCH 058/365] Account for recurrent endpoint storage when sizing the KV pool Reserve private instruction, prompt, and response checkpoint blocks in boundary-policy layout selection, maximum-context admission, and capacity reporting. Keep aligned-policy accounting unchanged. This prevents a small hybrid pool from passing startup while every request waits for unavailable checkpoint blocks. Validation: 131 KV sizing and contiguous-packing CPU tests passed; the seven-group TP2 MTP3 reproducer requires 41 blocks without an instruction checkpoint and 49 with one, exceeding its 37 usable blocks. Ruff and git diff --check pass. GPU serving qualification is pending. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 94febe5c35530d630b161b5abf5617e484b70fbe) --- tests/v1/core/test_contiguous_kv_packing.py | 1 + tests/v1/core/test_kv_cache_utils.py | 104 ++++++++++++++++++++ vllm/v1/core/kv_cache_utils.py | 35 ++++++- 3 files changed, 139 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index dc56526e2329..15888fe56db2 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -72,6 +72,7 @@ def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5): def _mock_vllm_config(layout: str | None): config = MagicMock() + config.use_request_boundary_checkpoints = False config.cache_config = CacheConfig() config.cache_config.num_gpu_blocks_override = None config.cache_config.kv_cache_layout = layout diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1da2be251f9a..76fc5f2945c5 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2653,6 +2653,42 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): ) +@pytest.mark.parametrize("with_tail", [False, True]) +def test_balanced_glm_boundary_capacity_includes_private_endpoints( + monkeypatch, with_tail +): + """The GLM slot-sharing layout still reserves distinct endpoint block IDs.""" + config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + specs = ( + _glm5_like_kv_cache_spec_with_tail() + if with_tail + else _glm5_like_kv_cache_spec()[0] + ) + groups = get_kv_cache_groups(config, specs) + assert kv_cache_utils._glm5_next_tensor_layout(groups) is not None + block_bytes = kv_cache_utils._pool_bytes_per_block(groups) + live_bytes = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) + live_blocks = live_bytes // block_bytes + private_blocks = 3 * (len(groups) + 1) + monkeypatch.setattr( + VllmConfig, "use_request_boundary_checkpoints", property(lambda self: True) + ) + + required = (live_blocks + private_blocks) * block_bytes + assert ( + kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) == required + ) + assert ( + kv_cache_utils._get_kv_cache_group_allocation_cost(config, groups) == required + ) + pool = kv_cache_utils.get_kv_cache_config_from_groups(config, groups, required) + assert get_max_concurrency_for_kv_cache_config(config, pool) == 1 + assert ( + kv_cache_utils._estimate_max_model_len_from_groups(config, groups, live_bytes) + < config.model_config.max_model_len + ) + + def test_glm5_kpool_tail_does_not_drag_hash_block_size(): """The tail's kpool-sized scratch block (4 tokens) must not constrain the prefix-cache hash granularity: participating groups alone decide it.""" @@ -3243,6 +3279,7 @@ def test_glm5next_split_cache_preserves_physical_pages( def _glm5next_split_config() -> SimpleNamespace: return SimpleNamespace( + use_request_boundary_checkpoints=False, scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), speculative_config=None, model_config=SimpleNamespace(max_model_len=202_752), @@ -3380,6 +3417,7 @@ def test_glm5next_nvfp4_auto_geometry_capacity() -> None: prefix_cache_retention_interval=4096, ) vllm_config = SimpleNamespace( + use_request_boundary_checkpoints=False, model_config=SimpleNamespace(max_model_len=202_752), parallel_config=SimpleNamespace(decode_context_parallel_size=4), cache_config=SimpleNamespace(mamba_cache_mode="align"), @@ -4238,6 +4276,72 @@ def test_auto_fit_max_model_len_with_hybrid(): assert vllm_config.model_config.max_model_len == 1024 +@pytest.mark.parametrize("dcp", [1, 2, 4]) +@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (70, 15)]) +def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_pages): + """Live MTP state alone does not cover instruction/prompt/response copies.""" + config = SimpleNamespace( + use_request_boundary_checkpoints=True, + attention_config=SimpleNamespace(hisparse_config=None), + model_config=SimpleNamespace(max_model_len=1048576), + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), + cache_config=SimpleNamespace(mamba_cache_mode="align"), + ) + attention = MLAAttentionSpec( + block_size=2048, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + model_version="glm5_next", + ) + state = MambaSpec( + block_size=2048, + shapes=((1,),), + dtypes=(torch.uint8,), + mamba_cache_mode="align", + num_speculative_blocks=3, + ) + groups = [KVCacheGroupSpec([f"state-{i}"], state) for i in range(6)] + groups.append(KVCacheGroupSpec(["attention"], attention)) + block_bytes = kv_cache_utils._pool_bytes_per_block(groups) + # Six groups require five live/scratch states each; three endpoints + # additionally own eight blocks each. The null block is unavailable. + available = (num_blocks - 1) * block_bytes + estimate = kv_cache_utils._estimate_max_model_len_from_groups( + config, groups, available + ) + assert estimate == expected_pages * 2048 * dcp + assert config.model_config.max_model_len == 1048576 + + if expected_pages == 0: + with pytest.raises(ValueError, match="even a single token"): + kv_cache_utils._auto_fit_max_model_len(config, [groups], [available]) + else: + kv_cache_utils._auto_fit_max_model_len(config, [groups], [available]) + assert config.model_config.max_model_len == estimate + assert kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) == ( + available + ) + assert kv_cache_utils._get_kv_cache_group_allocation_cost(config, groups) == ( + available + ) + cache = KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ) + assert get_max_concurrency_for_kv_cache_config(config, cache) == ( + num_blocks / (30 + 24 + expected_pages) + ) + + # An aligned-policy control has no endpoint reservation. Its capacity + # calculation must keep the ordinary live-state contract. + config.use_request_boundary_checkpoints = False + config.model_config.max_model_len = 1048576 + control = kv_cache_utils._estimate_max_model_len_from_groups( + config, groups, available + ) + assert control == (num_blocks - 1 - 30) * 2048 * dcp + + def test_auto_fit_max_model_len_not_triggered(): """Test that auto-fit is not triggered when original_max_model_len is not -1.""" model_config = ModelConfig(max_model_len=16) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 490273b796a4..567cab0e4628 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1005,6 +1005,14 @@ def check_enough_kv_cache_memory( if groups else available_memory ) + if vllm_config.use_request_boundary_checkpoints: + _check_enough_kv_cache_memory( + check_memory, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + return _check_enough_kv_cache_memory( check_memory, lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), @@ -1068,6 +1076,18 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: return True +def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) -> int: + """Count private instruction, prompt, and response checkpoint storage.""" + if not num_groups or not vllm_config.use_request_boundary_checkpoints: + return 0 + + from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS + + # KVCacheManager.allocate_slots reserves one block per group and one + # auxiliary block per endpoint, separately from the live sequence state. + return NUM_BOUNDARY_CHECKPOINT_SLOTS * (num_groups + 1) + + def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1095,6 +1115,9 @@ def get_max_concurrency_for_kv_cache_config( host_blocks_per_request += required else: num_blocks_per_request += required + num_blocks_per_request += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_config.kv_cache_groups) + ) limits = [kv_cache_config.num_blocks / num_blocks_per_request] if host_blocks_per_request: assert kv_cache_config.hisparse_host_num_blocks is not None @@ -1748,6 +1771,9 @@ def _get_kv_cache_group_allocation_cost( ) for group in kv_cache_groups ) + num_blocks_per_request += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_groups) + ) return _get_kv_cache_bytes_per_block(kv_cache_groups) * num_blocks_per_request @@ -2776,7 +2802,10 @@ def _max_memory_usage_bytes_from_groups( return 0 if vllm_config.attention_config.hisparse_config is not None: - return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + ( + _request_boundary_reserve_blocks(vllm_config, len(kv_cache_groups)) + * _pool_bytes_per_block(kv_cache_groups) + ) if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( @@ -2800,6 +2829,9 @@ def _max_memory_usage_bytes_from_groups( ) if tail_names: total_blocks += 1 + total_blocks += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_groups) + ) return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page) bytes_per_block = _pool_bytes_per_block(kv_cache_groups) @@ -2814,6 +2846,7 @@ def _max_memory_usage_bytes_from_groups( spec.page_size_bytes, ) + total_blocks += _request_boundary_reserve_blocks(vllm_config, len(kv_cache_groups)) return bytes_per_block * total_blocks From 3afee67a6d4137715955de38456c0bf304f7948b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:48:50 +0000 Subject: [PATCH 059/365] Release unused sparse MLA projections after weight absorption MQA-only sparse attention retains its independently owned W_UK_T/W_UV matrices but releases the MXFP8 linear holder, packed weights and preparation provider. Reload repacks the linear and preserves absorbed matrix addresses. Profiling reserves expanded dense context only for backends that can execute dense MHA. Adapts the projection lifetime and profiling portions of 7f4aecc66e857d093e012a2c57d03711bd23be4a to the Karmic preparation API. Five focused CPU/native GPU tests pass, including exact MXFP8 absorption and reload with both sparse-only and dense-prefill configurations. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_mla_backends.py | 93 ++++++++++++++++++- .../layers/attention/mla_attention.py | 36 ++++--- 2 files changed, 116 insertions(+), 13 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 10bd0dc8d4ce..ed0d6d47c86e 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -346,7 +346,9 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): layer.dcp_q_replicate = False layer.quant_config = None layer.layer_name = "test" - layer.impl = SimpleNamespace(process_weights_after_loading=lambda act_dtype: None) + layer.impl = SimpleNamespace( + process_weights_after_loading=lambda act_dtype: None, is_sparse=False + ) monkeypatch.setattr( mla_attention_module, "set_default_quant_scales", lambda *_, **__: None @@ -370,6 +372,95 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100) +@pytest.mark.parametrize("has_mha_prefill", [False, True]) +def test_mla_profile_reserves_only_supported_dense_context( + monkeypatch, has_mha_prefill +): + layer = SimpleNamespace( + prefill_backend=object() if has_mha_prefill else None, + chunked_prefill_workspace_size=128, + num_heads=2, + qk_nope_head_dim=4, + v_head_dim=4, + ) + q = torch.ones(3, 2, 8, dtype=torch.bfloat16) + kv_c = torch.ones(3, 8, dtype=torch.bfloat16) + output = torch.full((3, 8), float("nan"), dtype=torch.bfloat16) + allocations = [] + empty = torch.empty + + def record_empty(shape, **kwargs): + allocations.append(shape) + return empty(shape, **kwargs) + + monkeypatch.setattr(torch, "empty", record_empty) + actual = MLAAttention.forward_impl(layer, q, kv_c, kv_c, kv_c, None, output) + assert actual is output + assert torch.count_nonzero(output) == 0 + assert allocations == ([(128, 2, 8)] if has_mha_prefill else []) + + +@pytest.mark.skipif(not torch.accelerator.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("has_mha_prefill", [False, True]) +@torch.inference_mode() +def test_mla_absorbed_mxfp8_weights_release_unused_linear_and_reload( + monkeypatch, has_mha_prefill +): + from vllm.model_executor.kernels.linear.mxfp8.b12x import B12xMxfp8LinearKernel + + layer = MLAAttention.__new__(MLAAttention) + torch.nn.Module.__init__(layer) + layer.kv_lora_rank = 32 + layer.num_heads = 2 + layer.qk_nope_head_dim = layer.v_head_dim = 8 + layer.kv_b_proj = torch.nn.Module() + layer.kv_b_proj.prefix = "test.absorbed.kv_b_proj" + layer.kv_b_proj.quant_method = None + layer.is_aiter_triton_fp4_bmm_enabled = False + layer.is_aiter_triton_fp8_bmm_enabled = False + layer.is_amx_bmm_enabled = False + layer.dcp_q_replicate = False + layer.quant_config = None + layer.layer_name = "test.absorbed" + layer.impl = SimpleNamespace( + process_weights_after_loading=lambda act_dtype: None, is_sparse=True + ) + layer.prefill_backend = object() if has_mha_prefill else None + monkeypatch.setattr( + mla_attention_module, "set_default_quant_scales", lambda *_, **__: None + ) + kernel = object.__new__(B12xMxfp8LinearKernel) + addresses = None + for offset in (0, 1): + source = torch.arange(1024, device="cuda").reshape(32, 32).remainder(15) + source = (source + offset).to(torch.float8_e4m3fn) + scales = torch.full((32, 1), 127, device="cuda", dtype=torch.uint8) + layer.kv_b_proj.weight = torch.nn.Parameter(source, requires_grad=False) + layer.kv_b_proj.weight_scale = torch.nn.Parameter(scales, requires_grad=False) + kernel.process_weights_after_loading(layer.kv_b_proj) + packed = layer.kv_b_proj.b12x_mxfp8_packed_weight + holder = layer.kv_b_proj.b12x_linear + layer.process_weights_after_loading(torch.float32) + expected = source.float().T.reshape(32, 2, 16) + torch.testing.assert_close( + layer.W_UK_T, expected[:, :, :8].permute(1, 2, 0), rtol=0, atol=0 + ) + torch.testing.assert_close( + layer.W_UV, expected[:, :, 8:].transpose(0, 1), rtol=0, atol=0 + ) + observed = (layer.W_UK_T.data_ptr(), layer.W_UV.data_ptr()) + if addresses is not None: + assert observed == addresses + addresses = observed + assert layer.kv_b_proj.b12x_mxfp8_packed_weight is ( + packed if has_mha_prefill else None + ) + assert layer.kv_b_proj.b12x_linear is (holder if has_mha_prefill else None) + assert layer.kv_b_proj.b12x_preparation_provider is ( + kernel if has_mha_prefill else None + ) + + # Validate parameter combinations during collection, before GPU fixtures run. PREFILL_BACKENDS_TO_TEST: list[MLAPrefillBackendEnum] = [] if current_platform.is_cuda(): diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index afe70dde979f..29b30379679b 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1142,18 +1142,18 @@ def forward_impl( ) if attn_metadata is None: - # During the profile run try to simulate to worse case output size - # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context` - # since this can be large - _ = torch.empty( - ( - self.chunked_prefill_workspace_size, - self.num_heads, - self.qk_nope_head_dim + self.v_head_dim, - ), - device=k_c_normed.device, - dtype=k_c_normed.dtype, - ) + if self.prefill_backend is not None: + # Dense MHA expands compressed keys through kv_b_proj. An + # MQA-only backend never materializes this context-sized tensor. + _ = torch.empty( + ( + self.chunked_prefill_workspace_size, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ), + device=k_c_normed.device, + dtype=k_c_normed.dtype, + ) # The zero fill is required when used with DP + EP # to ensure all ranks within a DP group compute the @@ -1625,6 +1625,18 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): self._b12x_query_layer_name = _encode_layer_name(self._b12x_query_prefix) set_b12x_preparation_provider(self, self) + if ( + self.impl.is_sparse + and self.prefill_backend is None + and getattr(self.kv_b_proj, "b12x_mxfp8_packed_weight", None) is not None + ): + # MQA-only execution owns W_UK_T/W_UV and never calls this linear. + # Release both packed-weight owners and its preparation provider. + # A weight reload repacks the linear before refreshing these matrices. + self.kv_b_proj.b12x_mxfp8_packed_weight = None + self.kv_b_proj.b12x_linear = None + set_b12x_preparation_provider(self.kv_b_proj, None) + def _b12x_query_name(self, tokens: int) -> str: return f"{self._b12x_query_prefix}.m{tokens}" From 7e8284e8fba80ffeed667d2ec93b15a9678f48a7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:18:04 +0000 Subject: [PATCH 060/365] perf(b12x): omit unused sparse MLA index-group storage B12X performs physical selection through its prepared indexer or caller-owned scratch. Do not allocate the generic group conversion buffers or events for that backend. Other sparse backends retain group allocation and notification behavior. Validation: four construction/notification contracts pass for implicit and explicit builders; stock TP2 allocation inventory identified twelve unused 33,611,788-byte buffers. Whole-model capacity qualification is separate. Existing sparse-MLA typing findings reproduce before this change; other hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 65 +++++++++++++++++++ .../layers/attention/sparse_mla_attention.py | 9 ++- .../attention/backends/mla/b12x_mla_sparse.py | 3 + 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index c9f9c678ef62..7a85dbb686f6 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -71,6 +71,71 @@ def get_simultaneous(self, *shapes_and_dtypes): return [torch.empty(shape, dtype=dtype) for shape, dtype in shapes_and_dtypes] +@pytest.mark.parametrize("grouped_backend", [False, True]) +@pytest.mark.parametrize("explicit_builder", [False, True]) +def test_sparse_backend_allocates_only_its_consumed_index_group( + monkeypatch, grouped_backend, explicit_builder +): + from vllm.model_executor.layers.attention import sparse_mla_attention as common + + impl = object.__new__(B12xMLASparseImpl) + if grouped_backend: + impl.uses_index_group = True + indices = torch.zeros((4, 8), dtype=torch.int32) + calls: list[object] = [] + group = SimpleNamespace( + set_logical_topk_ready=lambda index: calls.append(("ready", index)), + prepare_for_batch=lambda index, metadata: calls.append(("batch", index)), + ) + + def register(*args, **kwargs): + calls.append("allocate") + return group, 0 + + builder = SimpleNamespace(register_layer=register) + monkeypatch.setattr(common, "SparseMLAIndexGroupBuilder", lambda topk: builder) + + def initialize_base(self, *args): + self.num_heads, self.qk_nope_head_dim, self.qk_rope_head_dim = 1, 8, 0 + + monkeypatch.setattr(common.MLACommonBaseImpl, "__init__", initialize_base) + monkeypatch.setattr(common, "has_flashinfer", lambda: False) + monkeypatch.setattr(common, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(common, "_is_masked_mha_available", lambda **kwargs: False) + monkeypatch.setattr(common, "get_current_vllm_config", lambda: object()) + common.SparseMLACommonImpl.__init__( + impl, + num_heads=1, + head_size=8, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_ds_mla", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=8, + qk_nope_head_dim=8, + qk_rope_head_dim=0, + qk_head_dim=8, + v_head_dim=8, + kv_b_proj=None, + topk_indices_buffer=indices, + index_group_builder=builder if explicit_builder else None, + ) + assert impl.topk_indices_buffer is indices + impl.record_logical_topk_ready() + impl.prepare_for_batch(None) + if grouped_backend: + assert impl.index_group is group + assert calls == ["allocate", ("ready", 0), ("batch", 0)] + else: + assert impl.index_group is None + assert not calls + + def test_b12x_selector_routes_supported_attention_families() -> None: assert AttentionConfig(backend="b12x").backend == AttentionBackendEnum.B12X assert AttentionBackendEnum.B12X.get_class() is B12xPagedAttentionBackend diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index cf7bce52e908..9e6664913ee6 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -646,6 +646,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True + uses_index_group = True def __init__( self, @@ -688,9 +689,13 @@ def __init__( self.init_topk_indices_buffer(indexer, topk_indices_buffer) self.index_group: SparseMLAIndexGroup | None = None self.index_group_index = 0 - if index_group_builder is None and self.topk_indices_buffer is not None: + if ( + self.uses_index_group + and index_group_builder is None + and self.topk_indices_buffer is not None + ): index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) - if index_group_builder is not None: + if self.uses_index_group and index_group_builder is not None: vllm_config = get_current_vllm_config() self.index_group, self.index_group_index = ( index_group_builder.register_layer( diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index cfc00970bc58..793a17bee537 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1172,6 +1172,9 @@ class B12xMLASparseImpl(SparseMLACommonImpl[B12xMLASparseMetadata]): lse_base_on_e = True supports_dense_mha_prefill = False supports_pcp = False + # B12X maps selections into caller-owned scratch or uses the indexer's + # prepared physical selection. Generic group buffers are never consumed. + uses_index_group = False def __init__( self, From b7d22fb3772d6c93e67cf624f7bb9ed95b53af00 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:18:34 +0000 Subject: [PATCH 061/365] fix(glm): bound and prepare the DCP profiling cache Measure the single-request attention peak with one temporary cache block rather than a block for every configured request slot. Prepare pool-bound B12X plans before the dummy forward and release profile state on success or failure. Real serving capacity and CUDA-graph batch profiling are unchanged. Validation: 23 focused profile, state-lifecycle and sparse-group tests pass. Negative controls on the installed unmodified wheel fail. Existing GPUModelRunner typing and allocator-cleanup API hook findings reproduce on the unchanged parent; no new findings remain. Whole-model TP2/DCP2 qualification follows in a source-built wheel. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/worker/test_gpu_model_runner.py | 38 +++++++++++++++------ tests/v1/worker/test_gpu_model_runner_v2.py | 38 +++++++++++++++------ tests/v1/worker/test_gpu_worker.py | 14 ++++++-- vllm/v1/worker/gpu/model_runner.py | 12 ++++--- vllm/v1/worker/gpu_model_runner.py | 11 +++--- vllm/v1/worker/gpu_worker.py | 7 +++- 6 files changed, 89 insertions(+), 31 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 042b378010b4..1e3827bb02df 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1830,10 +1830,10 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): runner.initialize_kv_cache(kv_cache_config) -@pytest.mark.parametrize("dummy_run_fails", [False, True]) +@pytest.mark.parametrize("failure_stage", [None, "init", "prepare", "forward"]) def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( monkeypatch: pytest.MonkeyPatch, - dummy_run_fails: bool, + failure_stage: str | None, ): runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_config = SimpleNamespace( @@ -1844,11 +1844,22 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( runner.max_num_tokens = 4096 events: list[object] = [] - runner._init_minimal_kv_cache_for_profiling = lambda: events.append("init-kv") + def initialize(*, num_blocks): + assert num_blocks == 1 + events.append("init-kv") + if failure_stage == "init": + raise RuntimeError("expected DCP profile failure") + + def prepare(): + events.append("prepare") + if failure_stage == "prepare": + raise RuntimeError("expected DCP profile failure") + + runner._init_minimal_kv_cache_for_profiling = initialize def dummy_run(*args, **kwargs): events.append(("dummy-run", args, kwargs)) - if dummy_run_fails: + if failure_stage == "forward": raise RuntimeError("expected DCP profile failure") return torch.empty(1), torch.empty(1) @@ -1862,14 +1873,22 @@ def dummy_run(*args, **kwargs): lambda _: nullcontext(), ) - if dummy_run_fails: + if failure_stage is not None: with pytest.raises(RuntimeError, match="expected DCP profile failure"): - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) else: - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) assert events[0] == "init-kv" - assert events[1] == ( + assert events[-1] == "cleanup" + if failure_stage == "init": + assert events == ["init-kv", "cleanup"] + return + assert events[1] == "prepare" + if failure_stage == "prepare": + assert events == ["init-kv", "prepare", "cleanup"] + return + assert events[2] == ( "dummy-run", (4096,), { @@ -1879,8 +1898,7 @@ def dummy_run(*args, **kwargs): "single_request_prefill": True, }, ) - assert events[-1] == "cleanup" - if not dummy_run_fails: + if failure_stage is None: assert events[-2] == "sync" diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 35d9f2f5d491..806d8e068e13 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -396,10 +396,10 @@ def test_boundary_logits_only_dispatches_pending_cache_tasks( assert input_batch.positions.item() == 3 -@pytest.mark.parametrize("dummy_run_fails", [False, True]) +@pytest.mark.parametrize("failure_stage", [None, "init", "prepare", "forward"]) def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( monkeypatch: pytest.MonkeyPatch, - dummy_run_fails: bool, + failure_stage: str | None, ): runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_config = SimpleNamespace( @@ -410,10 +410,21 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( runner.max_num_tokens = 4096 events: list[object] = [] + def initialize(_, *, num_blocks): + assert num_blocks == 1 + events.append("init-kv") + if failure_stage == "init": + raise RuntimeError("expected DCP profile failure") + + def prepare(): + events.append("prepare") + if failure_stage == "prepare": + raise RuntimeError("expected DCP profile failure") + monkeypatch.setattr( model_runner_module, "_init_minimal_kv_cache_for_profiling", - lambda _: events.append("init-kv"), + initialize, ) monkeypatch.setattr( model_runner_module, @@ -423,21 +434,29 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( def dummy_run(*args, **kwargs): events.append(("dummy-run", args, kwargs)) - if dummy_run_fails: + if failure_stage == "forward": raise RuntimeError("expected DCP profile failure") return torch.empty(1), torch.empty(1) runner._dummy_run = dummy_run monkeypatch.setattr(torch.accelerator, "synchronize", lambda: events.append("sync")) - if dummy_run_fails: + if failure_stage is not None: with pytest.raises(RuntimeError, match="expected DCP profile failure"): - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) else: - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) assert events[0] == "init-kv" - assert events[1] == ( + assert events[-1] == "cleanup" + if failure_stage == "init": + assert events == ["init-kv", "cleanup"] + return + assert events[1] == "prepare" + if failure_stage == "prepare": + assert events == ["init-kv", "prepare", "cleanup"] + return + assert events[2] == ( "dummy-run", (4096,), { @@ -448,8 +467,7 @@ def dummy_run(*args, **kwargs): "profile_all_kv_cache_groups": True, }, ) - assert events[-1] == "cleanup" - if not dummy_run_fails: + if failure_stage is None: assert events[-2] == "sync" diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index cda0764b458a..2133d437d7ee 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -326,10 +326,18 @@ def profile_cudagraph_memory(prepare_profile_state): events.append("profile_cudagraph_memory") return graph_estimate + def profile_run(prepare_profile_state): + prepare_profile_state() + events.append("profile_run") + + def profile_glm_dcp_attention(prepare_profile_state): + prepare_profile_state() + events.append("profile_glm_dcp_attention") + model_runner = SimpleNamespace( model_memory_usage=0, - profile_run=lambda prepare: (prepare(), events.append("profile_run")), - profile_glm_dcp_attention=lambda: events.append("profile_glm_dcp_attention"), + profile_run=profile_run, + profile_glm_dcp_attention=profile_glm_dcp_attention, profile_cudagraph_memory=profile_cudagraph_memory, ) profile_result = SimpleNamespace( @@ -408,7 +416,9 @@ def final_snapshot(**_kwargs): "prepare_b12x_profile_state", "profile_run", "release_b12x_profile_state", + "prepare_b12x_profile_state", "profile_glm_dcp_attention", + "release_b12x_profile_state", "prepare_b12x_profile_state", "profile_cudagraph_memory", "release_b12x_profile_state", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5b57f9e55146..cfb3cd98933e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1110,15 +1110,17 @@ def _profile_deepseek_v4_attention( _teardown_profiling_state(self) @torch.inference_mode() - def profile_glm_dcp_attention(self) -> None: + def profile_glm_dcp_attention( + self, prepare_profile_state: Callable[[], None] | None = None + ) -> None: """Measure the GLM DCP query-gather peak before KV cache sizing. The general activation profile deliberately skips attention and splits its token budget across many requests. A GLM sparse-MLA prefill can put the complete scheduler budget in one request, causing the DCP query all-gather to require substantially more temporary memory. Bind a - minimal split cache, execute that shape, then release all temporary - cache and backend state before production cache allocation. + one-block split cache for that single request, prepare its bound kernel + plans, then release the temporary cache before production allocation. """ if ( self.model_config.architecture != "Glm5NextForConditionalGeneration" @@ -1126,8 +1128,10 @@ def profile_glm_dcp_attention(self) -> None: ): return - _init_minimal_kv_cache_for_profiling(self) try: + _init_minimal_kv_cache_for_profiling(self, num_blocks=1) + if prepare_profile_state is not None: + prepare_profile_state() self._dummy_run( self.max_num_tokens, context_len=self.dcp_size * self.cp_interleave, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1fc4022327a5..8e890bfc9e94 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6644,7 +6644,9 @@ def _profile_deepseek_v4_attention( self._cleanup_profiling_kv_cache() @torch.inference_mode() - def profile_glm_dcp_attention(self) -> None: + def profile_glm_dcp_attention( + self, prepare_profile_state: Callable[[], None] | None = None + ) -> None: """Profile GLM split-cache DCP attention before KV cache sizing. The generic activation profile omits attention metadata and spreads the @@ -6659,11 +6661,12 @@ def profile_glm_dcp_attention(self) -> None: ): return - with set_current_vllm_config(self.vllm_config): - self._init_minimal_kv_cache_for_profiling() - model_output: tuple[torch.Tensor, torch.Tensor] | None = None try: + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling(num_blocks=1) + if prepare_profile_state is not None: + prepare_profile_state() model_output = self._dummy_run( self.max_num_tokens, force_attention=True, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index e2af61a2e389..bc7d8949cf07 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -641,7 +641,12 @@ def determine_available_memory(self) -> int: self.model_runner.profile_run(self._prepare_b12x_profile_state) finally: self._release_b12x_profile_state() - self.model_runner.profile_glm_dcp_attention() + try: + self.model_runner.profile_glm_dcp_attention( + self._prepare_b12x_profile_state + ) + finally: + self._release_b12x_profile_state() # Profile CUDA graph memory if graphs will be captured. # ROCm is included: #44825 moved the profiler to From ed7b6184abf0ea63455ee0ef0ba0bc715fb40502 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 20:40:28 +0000 Subject: [PATCH 062/365] fix(ds4.1): release transient MHC and WO graph activations Retain caller-owned scratch without extending activation lifetimes through B12X bindings. Keep prepared plans, fresh bindings, precision and CUDA graph coverage unchanged. Native GPU tests reproduce retained operands in both MHC modes and both WO modes. All four corrected cases and the block-linear regression pass exact changed-input graph replay. TP4 DSpark K7 serving restores logical KV capacity from 1,920,034 to 4,466,215 tokens and passes bounded text, prefix and vision checks. C1 is -0.65%, C8 -2.32% and 32k prefill +0.44% in sequential samples; regression-free throughput is not claimed. Ruff findings and eleven source mypy findings reproduce on the untouched base. Tests pass mypy. Skip those pre-existing hook failures and whole-file formatting for this isolated local commit; other commit hooks remain enabled. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit a28be967e5779629dbc212a312a3ad8924c0623f) --- tests/v1/attention/test_b12x_v41_workspace.py | 241 ++++++++++++++++++ vllm/models/deepseek_v4_1/attention.py | 7 +- vllm/models/deepseek_v4_1/b12x_layers.py | 4 +- 3 files changed, 249 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index e6bd78e6719d..1a87261c3e7a 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -422,6 +422,144 @@ def run(rows): session.close() +@pytest.mark.parametrize("operation", ["pre", "post_pre"]) +def test_mhc_capture_retains_scratch_not_caller_activations( + native_workspace, monkeypatch, operation +): + """Intermediate MHC outputs must be reusable across shared-pool graphs.""" + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4_1 import b12x_layers + from vllm.utils.b12x import B12xWorkload, register_b12x_layer + + _, manager, workspace = native_workspace + monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (8, 16)) + device = torch.device("cuda") + hidden = 5120 + torch.manual_seed(146) + layer = torch.nn.Module() + for kind in ("attn", "ffn"): + setattr(layer, f"hc_{kind}_fn", torch.randn(24, hidden * 4, device=device) / 64) + setattr(layer, f"hc_{kind}_scale", torch.ones(3, device=device)) + setattr(layer, f"hc_{kind}_base", torch.zeros(24, device=device)) + setattr( + layer, + f"{kind}_norm", + SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device) + ), + ) + layer.hc_attn_fn_broadcast = None + module = layer._b12x_mhc = b12x_layers.B12xMHC( + SimpleNamespace( + hidden_size=hidden, + rms_norm_eps=1e-20, + hc_eps=1e-6, + hc_sinkhorn_iters=20, + hc_mult=4, + ) + ) + name = f"test.mhc.capture.{operation}" + module.bind_layer_name(name) + register_b12x_layer(name, layer) + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=16, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare( + tuple( + request + for unit in module.get_b12x_preparation_units(layer, workload) + for request in unit.requests + ) + ) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in module._plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + residual = torch.randn(16, 4, hidden, dtype=torch.bfloat16, device=device) + previous = torch.randn(16, hidden, dtype=torch.bfloat16, device=device) + pre = torch.full((16, 4), 0.25, device=device) + previous_post = pre.clone() + comb = torch.eye(4, device=device).expand(16, -1, -1).contiguous() + destinations = [ + torch.empty_like(residual), + torch.empty_like(pre), + torch.empty_like(comb), + torch.empty_like(previous), + torch.empty_like(pre), + ] + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + kwargs = ( + {} + if operation == "pre" + else dict( + previous_output=previous[:rows], + previous_post=previous_post[:rows], + previous_comb=comb[:rows], + ) + ) + values = module.pre( + residual[:rows], + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + layer.attn_norm.weight, + pre[:rows], + **kwargs, + ) + for destination, value in zip(destinations, values, strict=True): + references.append(weakref.ref(value)) + destination[:rows].copy_(value) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + residual.normal_() + previous.normal_() + run(rows) + expected = [output[:rows].clone() for output in destinations] + for output in destinations: + output.fill_(float("nan")) + graphs[rows].replay() + for output, reference in zip(destinations, expected, strict=True): + torch.testing.assert_close(output[:rows], reference, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + @pytest.mark.parametrize("draft_tokens", [5, 7]) def test_parallel_draft_reservation_keeps_decode_split_parallelism( native_workspace, draft_tokens @@ -824,6 +962,109 @@ def projection(n, k, value): del resources +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_output_projection_capture_releases_caller_activations( + native_workspace, is_prefill +): + """WO bindings may borrow inputs and results without pinning them per graph.""" + from b12x.preparation import PreparationSession + + from vllm.utils.b12x import B12xWorkload + + attention, manager, workspace = native_workspace + device = torch.device("cuda") + torch.manual_seed(147) + layer = _layer(attention) + layer.capacity = 16 + layer.n_local_groups = 2 + layer.n_local_heads = 16 + layer.head_dim = 512 + layer.rope_head_dim = 64 + layer.o_lora_rank = 1024 + layer.hidden_size = 5120 + layer._wo_plans = {} + angles = torch.randn(32, 32, device=device) + layer.rotary_emb = SimpleNamespace( + cos_sin_cache=torch.cat((angles.cos(), angles.sin()), dim=-1) + ) + for name, shape in (("wo_a", (2048, 4096)), ("wo_b", (5120, 2048))): + setattr( + layer, + name, + SimpleNamespace( + weight=(torch.randn(shape, device=device) / 32).to(torch.float8_e4m3fn), + weight_scale_inv=torch.ones( + shape[0] // 32, shape[1] // 32, device=device + ).to(torch.float8_e8m0fnu), + ), + ) + layer.setup_wo_projection() + workload = B12xWorkload( + stage="weights", + token_counts=(8, 16), + fixed_token_counts=(8,), + output_dtype=torch.bfloat16, + max_tokens=16, + max_seqs=1, + max_model_len=32, + ) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + session.prepare(layer._wo_preparation_unit(workload).requests) + session.freeze() + manager.reserve_all( + *( + (spec.shape, spec.dtype) + for plan in layer._wo_plans.values() + for spec in plan.scratch_specs() + ) + ) + manager.lock() + inputs = torch.randn(16, 16, 512, dtype=torch.bfloat16, device=device) + positions = torch.arange(16, device=device) + output = torch.empty(16, 5120, dtype=torch.bfloat16, device=device) + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def run(rows): + source = inputs[:rows] + 1 + result = layer._o_proj(source, positions[:rows], is_prefill=is_prefill) + references.extend((weakref.ref(source), weakref.ref(result))) + if result._base is not None: + references.append(weakref.ref(result._base)) + output[:rows].copy_(result) + + pool = torch.cuda.graph_pool_handle() + graphs, owners = {}, [] + try: + for rows in (16, 8): + run(rows) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + graphs[rows] = graph + with ( + session.capture(), + workspace.collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, pool=pool), + ): + run(rows) + owners.append(resources) + gc.collect() + assert all(reference() is None for reference in references) + assert resources + + for rows in (8, 16, 8, 16): + inputs.normal_() + positions.copy_(torch.randperm(16, device=device)) + run(rows) + expected = output[:rows].clone() + output.fill_(float("nan")) + graphs[rows].replay() + torch.testing.assert_close(output[:rows], expected, rtol=0, atol=0) + finally: + for graph in graphs.values(): + graph.reset() + session.close() + + def test_output_projection_uses_fused_block32_path(native_workspace, monkeypatch): from dataclasses import dataclass diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 661e1f862bee..06c590ce5bc8 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1706,9 +1706,10 @@ def _o_proj(self, o, positions, *, is_prefill=False): ) require_prepared(plan, "gemm.wo_projection", o.device) weights = self._wo_projection_weights + scratch = _scratch(plan) binding = wo_projection.bind_inv_rope( plan, - scratch=_scratch(plan), + scratch=scratch, o=o, positions=positions, cos_sin_cache=self.rotary_emb.cos_sin_cache, @@ -1725,7 +1726,9 @@ def _o_proj(self, o, positions, *, is_prefill=False): device=o.device, ) binding = replace(binding, output=output) - retain_cuda_graph_capture_resource(binding) + # The model owns the plan and weights; graph pools manage activation + # lifetimes. Keep only the borrowed workspace, not layer activations. + retain_cuda_graph_capture_resource(scratch) local = wo_projection.run_inv_rope( binding=binding, plan=plan, diff --git a/vllm/models/deepseek_v4_1/b12x_layers.py b/vllm/models/deepseek_v4_1/b12x_layers.py index de605dfd11c0..57292c9db2e1 100644 --- a/vllm/models/deepseek_v4_1/b12x_layers.py +++ b/vllm/models/deepseek_v4_1/b12x_layers.py @@ -586,7 +586,9 @@ def _mhc_pre( comb=comb, pre_out=pre_out, ) - retain_cuda_graph_capture_resource(binding) + # Caller-owned outputs follow PyTorch's graph-pool lifetimes. Retaining + # them through the binding prevents reuse across layers and graph shapes. + retain_cuda_graph_capture_resource(scratch) if previous_output is None: mhc.run_pre( residual, From de2b5e6ec742e5881d2678c1f0e375f4064b1cf9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 16 Sep 2026 21:16:53 +0000 Subject: [PATCH 063/365] fix(ds4.1): release captured query and context activations Retain caller-owned scratch and attention metadata without retaining transient query/index operands or DSpark context-projection results. Prepared plans, per-call bindings, numerical kernels, cache geometry and graph coverage remain unchanged. Real SM120 tests reproduce the immutable-image ownership defects, including detached query storage. Twenty-four native lifetime and exact changed-input replay cases pass across decode/prefill and three page geometries. TP4 DSpark K7 memory snapshots remove 858657792 additional live graph bytes per rank and pass text, prefix and vision checks. Independent uninstrumented throughput qualification is recorded separately. The test file passes mypy; existing source typing findings remain. Ruff diagnostics match the parent revision. Skip the reproducing Ruff/mypy findings and unrelated whole-file formatting while retaining the remaining commit hooks. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 42627f07f551a5a38575eb1baca9693d7856c35e) --- tests/v1/attention/test_b12x_v41_workspace.py | 86 +++++++++++++------ vllm/models/deepseek_v4_1/attention.py | 22 ++--- vllm/models/deepseek_v4_1/nvidia/dspark.py | 2 +- 3 files changed, 72 insertions(+), 38 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 1a87261c3e7a..549ae1e88fec 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -4,10 +4,12 @@ import gc import weakref +from functools import partial from types import SimpleNamespace import pytest import torch +from torch.multiprocessing.reductions import StorageWeakRef from vllm.config import CUDAGraphMode @@ -336,9 +338,11 @@ def test_index_preparation_reuses_reserved_workspace(native_workspace): torch.testing.assert_close(cache, torch.zeros_like(cache), rtol=0, atol=0) +@pytest.mark.parametrize("projection_kind", ["linear", "dspark_context"]) def test_block_linear_capture_retains_scratch_not_caller_activations( native_workspace, monkeypatch, + projection_kind, ): from vllm.models.deepseek_v4_1 import b12x_layers @@ -353,10 +357,22 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( layer.weight = torch.randn(256, 256, device=device).to(torch.float8_e4m3fn) layer.weight_scale_inv = torch.ones(8, 8, device=device).to(torch.float8_e8m0fnu) monkeypatch.setattr(b12x_layers, "_execution_capacities", lambda: (16,)) - method = b12x_layers.B12xFP8LinearMethod( - SimpleNamespace(weight_block_size=[32, 32]) - ) - method.process_weights_after_loading(layer) + if projection_kind == "linear": + method = b12x_layers.B12xFP8LinearMethod( + SimpleNamespace(weight_block_size=[32, 32]) + ) + method.process_weights_after_loading(layer) + project = partial(method.apply, layer) + plans = layer.b12x_plans + else: + from vllm.models.deepseek_v4_1.nvidia import dspark + + monkeypatch.setattr(dspark, "_execution_capacities", lambda: (16,)) + method = dspark._ContextKVProjection( + SimpleNamespace(fused_wqa_wkv=layer, q_lora_rank=0), 16 + ) + project = method + plans = method.plans workload = B12xWorkload( stage="weights", token_counts=(8, 16), @@ -376,7 +392,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( ) session.freeze() manager.reserve_all( - *((spec.shape, spec.dtype) for spec in layer.b12x_plans[0].scratch_specs()) + *((spec.shape, spec.dtype) for spec in plans[0].scratch_specs()) ) manager.lock() inputs = torch.randn(16, 256, dtype=torch.bfloat16, device=device) @@ -385,7 +401,7 @@ def test_block_linear_capture_retains_scratch_not_caller_activations( def run(rows): source = inputs[:rows] + 1 - output = method.apply(layer, source) + output = project(source) references.extend((weakref.ref(source), weakref.ref(output))) outputs[:rows].copy_(output) @@ -695,6 +711,7 @@ def test_attention_shared_scratch_graph_replay( device = torch.device("cuda", torch.accelerator.current_device_index()) torch.manual_seed(142) layer = _layer(attention, swa_page=swa_page) + layer.rotary_emb = SimpleNamespace(cos_sin_cache=torch.empty(0, device=device)) layer.config.cache_config.block_size = main_page layer.max_model_len = 4096 if rows == 36: @@ -755,13 +772,12 @@ def metadata(page): device=device, dtype=torch.uint8, ) - _write_cache( - kv, - cache, - torch.arange(length, device=device) + page, - page_size=page, - cache_kind=kind, - cache_format="deepseek_v41", + from b12x.attention._shared.mla.compressed_reference import ( + pack_deepseek_v41_cache_reference, + ) + + cache[1:].copy_( + pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -782,12 +798,6 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 - _write_index_keys( - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - page_size=main_page, - ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -811,17 +821,35 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) + attention.dsa_indexer.quantize_write_index_k_mxfp4( + layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + ) + activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): - attention.dsa_indexer.quantize_q_mxfp4( - layer._index_plan("decode", min(rows, layer.DECODE_CHUNK)), - iq, - q_mxfp4=packed, - q_scales=scales, - ) - layer.forward_mqa( - q, None, positions, out, index_query=(packed, scales, weights) + query = q.clone() + index_query = packed.clone(), scales.clone(), weights.clone() + activation_refs.extend( + (name, StorageWeakRef(tensor.untyped_storage())) + for name, tensor in zip( + ("query", "index_query", "index_scales", "index_weights"), + (query, *index_query), + ) ) + mode = "decode" if is_decode else "prefill" + chunk = layer.DECODE_CHUNK if is_decode else layer.INDEX_CHUNK + for offset in range(0, rows, chunk): + end = min(offset + chunk, rows) + attention.dsa_indexer.quantize_q_mxfp4( + layer._index_plan(mode, end - offset), + iq[offset:end], + q_mxfp4=index_query[0][offset:end], + q_scales=index_query[1][offset:end], + ) + layer.forward_mqa(query, None, positions, out, index_query=index_query) run() batched_output = out.clone() @@ -845,6 +873,10 @@ def run(): torch.cuda.graph(graph, stream=stream), ): run() + # Graph resources own scratch, not the per-layer query activations. + gc.collect() + retained = [name for name, reference in activation_refs if not reference.expired()] + assert not retained, f"Capture resources retained caller activations: {retained}" # Changed queries exercise both selection and attention on replay, after # all plan storage has been reused by another operation. for step, seed in enumerate((143, 144)): diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 06c590ce5bc8..4dd404ee4853 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1454,16 +1454,17 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): candidate_lengths=source._candidate_lens[offset:end], ) index_plan = self._index_plan(index_mode, count) + scratch = current_workspace_manager().get_simultaneous( + *( + (spec.shape, spec.dtype) + for spec in dsa_indexer.scratch_specs( + index_plan, device=iq_data.device + ) + ) + ) binding = dsa_indexer.bind( index_plan, - scratch=current_workspace_manager().get_simultaneous( - *( - (spec.shape, spec.dtype) - for spec in dsa_indexer.scratch_specs( - index_plan, device=iq_data.device - ) - ) - ), + scratch=scratch, q_mxfp4=iq_data[offset:end], q_scales=iq_scale[offset:end], query_weights=iw[offset:end], @@ -1475,7 +1476,7 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): output_indices=self.topk_indices_buffer[offset:end], **candidate_args, ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) dsa_indexer.score(binding) dsa_indexer.select(binding) @@ -1534,7 +1535,8 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): swa_lengths=swa_lengths[:rows], **kwargs, ) - retain_cuda_graph_capture_resource(binding) + # Keep scratch and metadata, not aliases of the transient query storage. + retain_cuda_graph_capture_resource(buffers) mla.run( binding=binding, swa_k_cache=self.swa_cache_layer.kv_cache, diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v4_1/nvidia/dspark.py index 5251cc82d9e4..61ba555c05d3 100644 --- a/vllm/models/deepseek_v4_1/nvidia/dspark.py +++ b/vllm/models/deepseek_v4_1/nvidia/dspark.py @@ -186,7 +186,7 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor: packed_weight=self.weight, output=out.view(rows, self.weight.out_features, 1), ) - retain_cuda_graph_capture_resource(binding) + retain_cuda_graph_capture_resource(scratch) block_fp8_linear.run(binding=binding) return out From 870009e720b73ebef11f42fc348e250cb689d1c6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 12:25:32 +0000 Subject: [PATCH 064/365] test(ds41): verify MHC ownership across breakable graph segments Exercise pre and lagged post-pre outputs across an eager segment boundary and changed-input replay on a shared graph pool. Require transient Python owners to release while caller scratch remains retained. Native graph semantics and model code are unchanged. Applicable hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 3ead61584ac374ca5825ffb8aa21ad8a1a05beca) --- tests/v1/attention/test_b12x_v41_workspace.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 549ae1e88fec..260f386aa122 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -439,12 +439,14 @@ def run(rows): @pytest.mark.parametrize("operation", ["pre", "post_pre"]) +@pytest.mark.parametrize("capture_kind", ["full", "breakable"]) def test_mhc_capture_retains_scratch_not_caller_activations( - native_workspace, monkeypatch, operation + native_workspace, monkeypatch, operation, capture_kind ): """Intermediate MHC outputs must be reusable across shared-pool graphs.""" from b12x.preparation import PreparationSession + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture from vllm.models.deepseek_v4_1 import b12x_layers from vllm.utils.b12x import B12xWorkload, register_b12x_layer @@ -537,6 +539,10 @@ def run(rows): pre[:rows], **kwargs, ) + capture = BreakableCUDAGraphCapture.current() + if capture is not None: + # The lagged outputs cross a segment boundary before consumption. + capture.add_eager(lambda: None) for destination, value in zip(destinations, values, strict=True): references.append(weakref.ref(value)) destination[:rows].copy_(value) @@ -547,14 +553,23 @@ def run(rows): for rows in (16, 8): run(rows) torch.accelerator.synchronize() - graph = torch.cuda.CUDAGraph() + graph = ( + torch.cuda.CUDAGraph() + if capture_kind == "full" + else BreakableCUDAGraphCapture(pool=pool) + ) graphs[rows] = graph + context = ( + torch.cuda.graph(graph, pool=pool) if capture_kind == "full" else graph + ) with ( session.capture(), workspace.collect_cuda_graph_capture_resources() as resources, - torch.cuda.graph(graph, pool=pool), + torch.cuda.stream(torch.cuda.Stream()), + context, ): run(rows) + torch.accelerator.synchronize() owners.append(resources) gc.collect() assert all(reference() is None for reference in references) From 36dd216bb1184507fa4938ad17cd42d5e90bb67e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:17:17 +0000 Subject: [PATCH 065/365] fix(ds4): reuse MHC autotuning output buffers Allocate one output set per preparation candidate so repeated CUDA graph timing samples do not accumulate functional-output copies beyond the primed memory budget. Serving kernels, operation queries, candidate search and caller-owned scratch are unchanged. Validation: 13 native SM120 tests pass, including exact functional parity and changed-input CUDA graph replay at 3 and 4096 tokens. The baseline fails output-pointer reuse. DS4 TP2 completes weights preparation; end-to-end serving remains gated by KV admission capacity. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 2b2fa507b89d9652c5d87db61c44ee89b5fe864f) --- .../kernels/test_b12x_linear.py | 112 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 22 +++- 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 72c4f6d130ff..6a1a540400e3 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -7,6 +7,7 @@ import sys import types from dataclasses import dataclass +from typing import Any import pytest import torch @@ -1874,6 +1875,117 @@ def test_b12x_mhc_declares_the_operations_each_layer_runs( mhc._plan_for("pre", 8) +@torch.inference_mode() +@pytest.mark.parametrize("operation", ["pre", "post_pre", "post_pre_bf16", "post"]) +@pytest.mark.parametrize("tokens", [3, 4096]) +def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( + monkeypatch, operation, tokens +): + """Tuning fixtures retain one output set per candidate, not per sample.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("native b12x mHC requires SM12x") + from b12x.norm.mhc import _impl + from b12x.preparation import PreparationSession + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x + + device = torch.device("cuda", torch.accelerator.current_device_index()) + hidden, mult = 4096, 4 + fn = torch.randn(24, mult * hidden, device=device) * 0.001 + norm = types.SimpleNamespace( + weight=torch.ones(hidden, dtype=torch.bfloat16, device=device), + variance_epsilon=1e-6, + ) + layer = types.SimpleNamespace( + hc_attn_fn=fn, + hc_ffn_fn=fn, + hc_attn_scale=torch.ones(3, device=device), + hc_ffn_scale=torch.ones(3, device=device), + hc_attn_base=torch.zeros(24, device=device), + hc_ffn_base=torch.zeros(24, device=device), + hc_attn_fn_broadcast=fn.view(24, mult, hidden).sum(1), + attn_norm=norm, + ffn_norm=norm, + hc_ffn_fn_bf16=fn.to(torch.bfloat16), + ) + mhc = dsv4_b12x.B12xMHCResidual( + hidden_size=hidden, + hc_mult=mult, + rms_eps=1e-6, + hc_eps=1e-6, + sinkhorn_iters=20, + ) + workload = B12xWorkload( + stage="weights", + token_counts=(tokens,), + fixed_token_counts=(), + output_dtype=torch.bfloat16, + max_tokens=tokens, + max_seqs=1, + max_model_len=tokens, + ) + (unit,) = mhc.get_b12x_preparation_units(layer, workload) + request = next( + item for item in unit.requests if item.name.endswith(f".{operation}.m{tokens}") + ) + native_operation = "post_pre" if operation == "post_pre_bf16" else operation + name = f"_b12x_mhc_{native_operation}_impl" + native = getattr(_impl, name) + captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + + def observe(*args, **kwargs): + captured_call[:] = [(args, kwargs)] + return native(*args, **kwargs) + + def outputs(value): + return value if isinstance(value, tuple) else (value,) + + with PreparationSession( + device=device, autotune=False, compile_workers=0 + ) as session: + session.prepare((request,)) + call = request.prepare_call(request.plan.prepared.state) + with monkeypatch.context() as hooks: + hooks.setattr(_impl, name, observe) + call.produce() + first = outputs(call.invoke()) + args, kwargs = captured_call.pop() + reference_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} + } + pointers = [tensor.data_ptr() for tensor in first] + second = outputs(call.invoke()) + assert [tensor.data_ptr() for tensor in second] == pointers + for actual, expected in zip( + second, outputs(native(*args, **reference_kwargs)), strict=True + ): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + try: + with session.capture(), torch.cuda.graph(graph): + call.invoke() + call.invoke() + for _ in range(2): + call.produce() + for tensor in first: + tensor.fill_(float("nan")) + allocated = torch.accelerator.memory_allocated(device) + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated(device) == allocated + assert [ + tensor.data_ptr() for tensor in outputs(call.output) + ] == pointers + expected = outputs(native(*args, **reference_kwargs)) + for actual, wanted in zip(first, expected, strict=True): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + finally: + graph.reset() + + def test_b12x_mhc_keeps_plans_across_workloads(monkeypatch) -> None: from vllm.models.deepseek_v4.nvidia import b12x as dsv4_b12x diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 98c210d79767..232d84b2dc94 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -435,6 +435,24 @@ def prepare(state): dtype=torch.float32, device=device, ) + # Each candidate owns one output set. Functional allocations inside + # repeated capture samples would reserve additional graph-pool copies + # that are absent from the tuner's primed-residency measurement. + residual_out = torch.empty( + (tokens, self.hc_mult, self.hidden_size), + dtype=torch.bfloat16, + device=device, + ) + outputs = ( + {} + if operation == "post" + else { + "residual_out": residual_out, + "y_out": torch.empty_like(x), + "post_out": torch.empty_like(post), + "comb_out": torch.empty_like(comb), + } + ) def produce(): residual.normal_() @@ -455,6 +473,7 @@ def produce(): norm_weight=attn_norm.weight, norm_eps=float(attn_norm.variance_epsilon), _state=state, + **outputs, ) elif operation in ("post_pre", "post_pre_bf16"): if operation == "post_pre": @@ -488,10 +507,11 @@ def produce(): norm_weight=norm.weight, norm_eps=float(norm.variance_epsilon), _state=state, + **outputs, ) else: run = lambda: _impl._b12x_mhc_post_impl( - x, residual, post, comb, _state=state + x, residual, post, comb, out=residual_out, _state=state ) def execute(): From 395967f6ca7a95c0ed72e7b916caf047424fcde0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 01:54:03 +0000 Subject: [PATCH 066/365] fix(workspace): reserve B12X scratch by model owner Keep target and DSpark draft scratch capacities independent during profile hooks and B12X preparation. Reserve shared modules in both lanes and propagate each lane capacity across microbatches before graph capture. Preserve the complete attention shape envelope and unscoped coordinator compatibility. Validation: 83 focused workspace, preparation and graph-profile tests pass. Stock DS4 TP2/DSpark K5 with a 4096-token budget admits 7.56 GiB KV instead of 6.12 GiB, retaining the 1M-token model limit. Serving throughput qualification is pending. The unchanged base reproduces eight coordinator mypy errors and the existing pickle-import lint finding; remaining hooks pass. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit ea358a87e72e6f44baf14c436d278b5f4d52b92c) --- tests/model_executor/test_b12x_warmup.py | 25 +++++ tests/v1/worker/test_workspace.py | 117 +++++++++++++++++++++ vllm/model_executor/warmup/b12x_prepare.py | 48 +++++++-- vllm/utils/b12x.py | 7 +- vllm/v1/worker/b12x_startup.py | 20 +++- vllm/v1/worker/gpu/model_runner.py | 29 +++-- vllm/v1/worker/workspace.py | 40 +++++-- 7 files changed, 259 insertions(+), 27 deletions(-) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index c54ce7c3856d..41690faa4fef 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -348,6 +348,31 @@ def test_scope_unit_calls_wraps_mapping_factories() -> None: request.prepare_call = {2: lambda state: _Call(run=lambda: _workspace_lane.get())} scoped = scope_b12x_unit_calls(_unit("u", request), 1) assert scoped.requests[0].prepare_call[2](object()).run() == 1 + assert scoped.workspace_lanes == (1,) + + +@pytest.mark.parametrize("draft_lane", [0, 1]) +def test_shared_preparation_scratch_covers_both_model_owners(draft_lane): + target, draft, shared = (torch.nn.Module() for _ in range(3)) + for module in (target, draft, shared): + module.b12x_preparation_provider = _Provider() + target.shared = shared + draft.shared = shared + units = b12x_prepare.collect_b12x_units( + _worker(target, draft=draft, draft_lane=draft_lane), _workload() + ) + ownership = b12x_prepare._request_workspace_lanes(units) + assert ownership[f"{id(target):x}"] == (0,) + assert ownership[f"{id(draft):x}"] == (draft_lane,) + assert ownership[f"{id(shared):x}"] == tuple(sorted({0, draft_lane})) + assert ( + sum( + f"{id(shared):x}" == request.name + for unit in units + for request in unit.requests + ) + == 1 + ) def test_blockscaled_holder_declares_provided_workspace_under_the_cap( diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 0eba7b71d592..1699fe90ea7f 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from typing import cast import pytest @@ -167,6 +168,122 @@ def test_workspace_lane_validation(monkeypatch) -> None: workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) +def test_profile_reserves_model_lanes_without_replicating_target_capacity( + monkeypatch, +) -> None: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + monkeypatch.setattr(workspace, "_manager", manager) + calls: list[tuple[int, int]] = [] + + class ScratchOwner(torch.nn.Module): + def __init__(self, size): + super().__init__() + self.size = size + + def reserve_profile_scratch(self): + calls.append((self.size, workspace._workspace_lane.get())) + manager.get_simultaneous(((self.size,), torch.uint8)) + + target = ScratchOwner(4096) + draft = ScratchOwner(512) + shared = ScratchOwner(256) + target.add_module("shared", shared) + draft.add_module("shared", shared) + runner = SimpleNamespace( + get_model=lambda: target, + get_draft_model=lambda: draft, + _draft_workspace_lane=1, + compilation_config=SimpleNamespace( + static_forward_context={ + "target": target, + "alias": target, + "draft": draft, + "shared": shared, + } + ), + ) + GPUModelRunner._reserve_profile_scratch(runner) + + assert calls == [(4096, 0), (512, 1), (256, 0), (256, 1)] + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + assert [buffer.numel() for buffer in buffers] == [4096, 512, 4096, 512] + pointers = [buffer.data_ptr() for buffer in buffers] + assert len(set(pointers)) == 4 + manager.lock() + manager.reserve_by_lane() + for ubatch in range(2): + active_ubatch[0] = ubatch + for lane, size in enumerate((4096, 512)): + with workspace.use_workspace_lane(lane): + (view,) = manager.get_simultaneous(((size,), torch.uint8)) + assert view.data_ptr() == pointers[ubatch * 2 + lane] + + +def test_lane_reservation_rejects_unreserved_microbatch_when_locked(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((256,), torch.uint8)) + manager.lock() + with pytest.raises(AssertionError, match="reserve_by_lane"): + manager.reserve_by_lane() + + +@pytest.mark.parametrize("scoped", [False, True]) +def test_preparation_reservation_preserves_per_lane_profile_capacity( + monkeypatch, scoped +): + from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator + + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((4096,), torch.uint8)) + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((256,), torch.uint8)) + requests = [] + for name, size in (("target", 1024), ("draft", 512), ("shared", 768)): + spec = SimpleNamespace(shape=(size,), dtype=torch.uint8) + requests.append( + SimpleNamespace( + name=name, + plan=SimpleNamespace(scratch_specs=lambda spec=spec: (spec,)), + ) + ) + job = SimpleNamespace( + advance=lambda **kwargs: SimpleNamespace(done=True, pending_compilation=False), + result=lambda: SimpleNamespace(close=lambda: None), + close=lambda: None, + ) + coordinator = B12xPreparationCoordinator( + SimpleNamespace(begin=lambda *args, **kwargs: job), + [(tuple(requests), False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=manager, + request_workspace_lanes=( + {"target": (0,), "draft": (1,), "shared": (0, 1)} if scoped else None + ), + ) + result = coordinator.advance() + assert result["done"] and result["error"] is None + buffers = [buffer for buffer in manager._current_workspaces if buffer is not None] + expected = [4096, 768, 4096, 768] if scoped else [4096] * 4 + assert [buffer.numel() for buffer in buffers] == expected + assert len({buffer.data_ptr() for buffer in buffers}) == 4 + manager.lock() + manager.reserve_by_lane() + + def test_cuda_graph_capture_resources_are_scoped_to_collector() -> None: outside = object() first = object() diff --git a/vllm/model_executor/warmup/b12x_prepare.py b/vllm/model_executor/warmup/b12x_prepare.py index dafff99c8e49..6ec0d9768778 100644 --- a/vllm/model_executor/warmup/b12x_prepare.py +++ b/vllm/model_executor/warmup/b12x_prepare.py @@ -228,6 +228,7 @@ def _units_from_modules( workload: B12xWorkload, *, seen: set[int] | None = None, + module_lanes: dict[int, tuple[int, ...]] | None = None, ) -> Iterable[B12xPreparationUnit]: if seen is None: seen = set() @@ -254,7 +255,12 @@ def _units_from_modules( f"{type(provider).__qualname__} returned a non-unit " "preparation value" ) - yield unit + lanes = ( + (workload.lane,) + if module_lanes is None + else module_lanes.get(id(module), (workload.lane,)) + ) + yield replace(unit, workspace_lanes=lanes) def mark_b12x_eager_shapes(worker: Worker) -> None: @@ -328,26 +334,43 @@ def collect_b12x_units( """Collect every unit of one stage from the model, the draft, and comms.""" mark_b12x_eager_shapes(worker) units: list[B12xPreparationUnit] = [] + model = worker.get_model() draft = worker.get_draft_model() + lane = _draft_lane(worker) if draft is not None else 0 + lanes_by_module: dict[int, set[int]] = {} + for owner, owner_lane in ((model, 0), (draft, lane)): + if owner is not None: + for module in owner.modules(): + lanes_by_module.setdefault(id(module), set()).add(owner_lane) + module_lanes = {key: tuple(sorted(lanes)) for key, lanes in lanes_by_module.items()} if draft is not None: - lane = _draft_lane(worker) workload = _draft_workload(worker, workload, lane=0) # Target and draft can share the same embedding and output-head modules. seen: set[int] = set() - units.extend(_units_from_modules(worker.get_model(), workload, seen=seen)) + units.extend( + _units_from_modules(model, workload, seen=seen, module_lanes=module_lanes) + ) if draft is not None: draft_workload = replace(workload, lane=lane) from vllm.v1.worker.workspace import use_workspace_lane with use_workspace_lane(lane): - draft_units = list(_units_from_modules(draft, draft_workload, seen=seen)) + draft_units = list( + _units_from_modules( + draft, draft_workload, seen=seen, module_lanes=module_lanes + ) + ) if lane: draft_units = [scope_b12x_unit_calls(unit, lane) for unit in draft_units] units.extend(draft_units) for provider in b12x_unit_providers(): hook = getattr(provider, "get_b12x_preparation_units", None) if callable(hook): - units.extend(hook(provider, workload)) + # Communicators serve target and draft executions independently. + units.extend( + replace(unit, workspace_lanes=tuple(sorted({0, lane}))) + for unit in hook(provider, workload) + ) # The state stage also carries the weights-stage units: their prepared # plans are skipped by the session, and any plan a layer declared for a # count that only the state stage exposes is prepared before capture. @@ -447,6 +470,14 @@ def b12x_batches(units: Iterable[B12xPreparationUnit], *, autotune: bool = True) return batches +def _request_workspace_lanes(units: Iterable[B12xPreparationUnit]): + return { + request.name: unit.workspace_lanes + for unit in units + for request in unit.requests + } + + def get_b12x_session(worker: Worker): """Return the worker's preparation session, creating it on first use.""" session = getattr(worker, "_b12x_session", None) @@ -492,10 +523,12 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): from vllm.v1.worker.workspace import current_workspace_manager batches = [] + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage="state" if stage == "bind" else stage) + units = collect_b12x_units(worker, workload) batches = b12x_batches( - collect_b12x_units(worker, workload), + units, autotune=( bool(worker.vllm_config.kernel_config.enable_b12x_autotune) and os.environ.get("B12X_AUTOTUNE", "1") != "0" @@ -531,6 +564,7 @@ def begin_b12x_preparation(worker: Worker, *, stage: str): global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if batches else None, + request_workspace_lanes=_request_workspace_lanes(units), ) @@ -608,6 +642,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: from vllm.v1.worker.workspace import current_workspace_manager requests: tuple[PreparationRequest, ...] = () + units = [] if b12x_native_supported(worker): workload = b12x_workload(worker, stage=stage) units = collect_b12x_units(worker, workload) @@ -621,6 +656,7 @@ def prepare_b12x_profile(worker: Worker, *, stage: str) -> B12xPreparedBatch: global_rank=int(worker.rank), world_group=get_world_group(), workspace=current_workspace_manager() if requests else None, + request_workspace_lanes=_request_workspace_lanes(units), ) outcome = coordinator.status() while not outcome["done"]: diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 11c16cfe4227..25e6ea3d2618 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -94,6 +94,7 @@ class B12xPreparationUnit: requests: tuple[PreparationRequest, ...] stage: Literal["weights", "state"] autotune: bool = True + workspace_lanes: tuple[int, ...] = (0,) def __post_init__(self): if not isinstance(self.name, str) or not self.name: @@ -103,6 +104,10 @@ def __post_init__(self): raise ValueError("invalid native preparation stage") if type(self.autotune) is not bool: raise TypeError("autotune must be boolean") + lanes = tuple(self.workspace_lanes) + if not lanes or any(type(lane) is not int or lane < 0 for lane in lanes): + raise ValueError("workspace lanes must be nonnegative integers") + object.__setattr__(self, "workspace_lanes", tuple(sorted(set(lanes)))) requests = tuple(self.requests) names = [request.name for request in requests] if len(names) != len(set(names)): @@ -210,7 +215,7 @@ def wrap_calls(calls): ) for request in unit.requests ) - return replace(unit, requests=requests) + return replace(unit, requests=requests, workspace_lanes=(lane,)) def b12x_preparation_token_counts( diff --git a/vllm/v1/worker/b12x_startup.py b/vllm/v1/worker/b12x_startup.py index 275bb0bf79fd..49fd5f51b5d7 100644 --- a/vllm/v1/worker/b12x_startup.py +++ b/vllm/v1/worker/b12x_startup.py @@ -38,11 +38,13 @@ def __init__( world_group, process_local_only: bool = False, workspace=None, + request_workspace_lanes: dict[str, tuple[int, ...]] | None = None, ) -> None: if type(global_rank) is not int or global_rank < 0: raise ValueError("global_rank must be a nonnegative integer") self.session = session self._workspace = workspace + self._request_workspace_lanes = request_workspace_lanes self._active_requests = () self.global_rank = global_rank self.world_group = world_group @@ -275,13 +277,25 @@ def _advance_local(self) -> None: job.result().close() if self._workspace is not None: + from vllm.v1.worker.workspace import use_workspace_lane + for request in self._active_requests: specs = tuple(request.plan.scratch_specs()) if specs: - self._workspace.get_simultaneous( - *((spec.shape, spec.dtype) for spec in specs) + lanes = ( + (0,) + if self._request_workspace_lanes is None + else self._request_workspace_lanes[request.name] ) - self._workspace.reserve_all() + for lane in lanes: + with use_workspace_lane(lane): + self._workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in specs) + ) + if self._request_workspace_lanes is None: + self._workspace.reserve_all() + else: + self._workspace.reserve_by_lane() self._active_requests = () self._job = None if self._batches: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43da8b0e4646..aba519befaff 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1004,15 +1004,26 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: self.pooling_runner.dummy_pooler_run(hidden_states) def _reserve_profile_scratch(self) -> None: - seen: set[int] = set() + model_lanes: dict[int, set[int]] = {} + for model, lane in ( + (self.get_model(), 0), + (self.get_draft_model(), self._draft_workspace_lane), + ): + if model is not None: + for module in model.modules(): + model_lanes.setdefault(id(module), set()).add(lane) + seen: set[tuple[int, int]] = set() for module in self.compilation_config.static_forward_context.values(): - if id(module) in seen: - continue - seen.add(id(module)) - reserve = getattr(module, "reserve_profile_scratch", None) - if reserve is not None: - reserve() - current_workspace_manager().reserve_all() + for lane in sorted(model_lanes.get(id(module), {0})): + key = (id(module), lane) + if key in seen: + continue + seen.add(key) + reserve = getattr(module, "reserve_profile_scratch", None) + if reserve is not None: + with use_workspace_lane(lane): + reserve() + current_workspace_manager().reserve_by_lane() @torch.inference_mode() def profile_run( @@ -1054,7 +1065,7 @@ def profile_run( torch.accelerator.synchronize() del hidden_states, sample_hidden_states self._profile_deepseek_v4_attention(prepare_profile_state) - current_workspace_manager().reserve_all() + current_workspace_manager().reserve_by_lane() self.reset_encoder_cache() gc.collect() diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 81adeee47065..ddde9f8b1df5 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -229,15 +229,40 @@ def reserve_all( required_bytes, max(map(self._workspace_size_bytes, self._current_workspaces), default=0), ) + self._reserve_slots( + [required_bytes] * len(self._current_workspaces), operation="reserve_all" + ) + + def reserve_by_lane(self) -> None: + """Cover every microbatch without copying capacity between model lanes. + + A target and its drafter have separate scratch contracts. A microbatch + can execute either model, but each model always uses its assigned lane. + """ + lane_bytes = [ + max( + map( + self._workspace_size_bytes, + self._current_workspaces[lane :: self._num_lanes], + ), + default=0, + ) + for lane in range(self._num_lanes) + ] + self._reserve_slots( + lane_bytes * self._num_ubatches, operation="reserve_by_lane" + ) + + def _reserve_slots(self, sizes: list[int], *, operation: str) -> None: undersized = [ workspace_id for workspace_id, workspace in enumerate(self._current_workspaces) - if self._workspace_size_bytes(workspace) < required_bytes + if self._workspace_size_bytes(workspace) < sizes[workspace_id] ] if self._locked and undersized: raise AssertionError( - "Workspace is locked but reserve_all requires " - f"{required_bytes / _MB:.2f} MB in slot(s) {undersized}." + f"Workspace is locked but {operation} requires " + f"larger capacity in slot(s) {undersized}." ) for workspace_id in undersized: @@ -246,14 +271,13 @@ def reserve_all( del current_workspace torch.accelerator.empty_cache() self._current_workspaces[workspace_id] = torch.empty( - (required_bytes,), dtype=torch.uint8, device=self._device + (sizes[workspace_id],), dtype=torch.uint8, device=self._device ) if envs.VLLM_DEBUG_WORKSPACE and undersized: logger.info( - "[WORKSPACE DEBUG] Reserved %.2f MB in execution slots %s", - required_bytes / _MB, - undersized, + "[WORKSPACE DEBUG] Reserved execution slots (slot, MB): %s", + [(slot, sizes[slot] / _MB) for slot in undersized], ) def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: @@ -327,7 +351,7 @@ def get_caller_info() -> str: # stream-ordered, so wait for the device first. Growth happens only # before the workspace is locked, never in steady-state serving. if self._device.type == "cuda" and current_workspace is not None: - torch.cuda.synchronize(self._device) + torch.accelerator.synchronize(self._device) self._current_workspaces[workspace_id] = None del current_workspace # Release the freed segment back to CUDA so the caching From 1fea6913a002224fe68861b77df112fb16eeed69 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 03:29:17 +0000 Subject: [PATCH 067/365] fix(gdn): bind caller-owned prefill activations directly Keep only sequence metadata on each GDN layer. Bind projected inputs and invocation-owned outputs through the native variable-row contract instead of retaining scheduler-capacity copies per layer. Priming uses temporary activation storage; recurrent-state ownership and kernel arithmetic are unchanged. Validation coverage includes token-capacity-independent metadata storage and CUDA graph replay with 49 or 64 rows, changing worklists, immutable inputs, and allocation counters. GPU qualification is recorded separately. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 983c38e2ef0b8f3acaaa029f952ba1c2f26f8033) --- tests/kernels/mamba/test_gdn_prefill_b12x.py | 33 ++++++++- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 72 ++++++++++--------- .../layers/mamba/ops/b12x_gdn_prefill.py | 58 +++++---------- 3 files changed, 87 insertions(+), 76 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_prefill_b12x.py b/tests/kernels/mamba/test_gdn_prefill_b12x.py index 50d66d32764f..153722421c08 100644 --- a/tests/kernels/mamba/test_gdn_prefill_b12x.py +++ b/tests/kernels/mamba/test_gdn_prefill_b12x.py @@ -10,6 +10,7 @@ from vllm.model_executor.layers.mamba.ops.b12x_gdn_prefill import ( B12xGdnPrefill, + GdnPrefillStaging, prefill_capacities, ) @@ -27,6 +28,25 @@ def test_prefill_capacity_family_covers_exact_scheduler_limit(capacity, expected assert prefill_capacities(capacity) == expected +@pytest.mark.parametrize("max_tokens", (16, 6019, 32768)) +def test_prefill_staging_memory_depends_on_sequences_not_token_capacity(max_tokens): + staging = GdnPrefillStaging.allocate( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + assert staging.nbytes == ((16 + 1) + 4 * 16 + 2) * 4 + assert staging.is_compatible( + max_tokens=max_tokens, + max_seqs=16, + key_heads=16, + value_heads=48, + device=torch.device("cpu"), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_prepared_gdn_decode_plan_is_scoped_to_the_bound_recurrent_pool( default_vllm_config, @@ -324,8 +344,10 @@ def test_prefill_pooled_state_graph_replays_changed_lengths_and_slots(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", (49, 64)) def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists( default_vllm_config, + rows, ): from b12x._lib.runtime_control import kernel_resolution_guard from b12x.preparation import PreparationSession @@ -353,7 +375,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pytest.skip("SM12x is required") torch.manual_seed(23) device = torch.device("cuda") - rows, seqs, key_heads, value_heads = 64, 3, 2, 6 + seqs, key_heads, value_heads = 3, 2, 6 width = (2 * key_heads + value_heads) * 128 layer = QwenGatedDeltaNetAttention.__new__(QwenGatedDeltaNetAttention) torch.nn.Module.__init__(layer) @@ -383,7 +405,7 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists pool, ) layer._b12x_preparation_prefix = "test.mixed-gdn" - layer._b12x_prefill_max_tokens = rows + layer._b12x_prefill_max_tokens = 64 layer._b12x_prefill_max_seqs = seqs layer._b12x_prefill_staging = None layer._b12x_decode_plan = None @@ -446,6 +468,11 @@ def test_mixed_gdn_graph_replays_prefill_decode_verification_and_empty_worklists ), attn_metadata=attention, ) + immutable_inputs = { + name: value.clone() + for name, value in inputs.items() + if isinstance(value, torch.Tensor) and name != "core_attn_out" + } saved_conv, saved_pool = conv.clone(), pool.clone() def stage(lengths, computed, drafts, swap): @@ -593,6 +620,8 @@ def oracle(): ) assert_close("mixed state", pool, expected_pool, ratio=5e-3) torch.testing.assert_close(conv, expected_conv, rtol=0, atol=0) + for name, value in immutable_inputs.items(): + torch.testing.assert_close(inputs[name], value, rtol=0, atol=0) finally: guard.__exit__(None, None, None) finally: diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1169d9304d43..f75bc330347b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -959,16 +959,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): if api is None: raise RuntimeError("b12x GDN prefill was not initialized") recurrent_state = self.kv_cache[1] - staging = self._b12x_prefill_staging - resident_nbytes = 0 - if staging is not None and staging.is_compatible( - max_tokens=self._b12x_prefill_max_tokens, - max_seqs=self._b12x_prefill_max_seqs, - key_heads=self._b12x_local_key_heads, - value_heads=self._b12x_local_value_heads, - device=recurrent_state.device, - ): - resident_nbytes = staging.nbytes caps = api.Caps( device=current_platform.current_device(), max_tokens=capacity, @@ -979,8 +969,6 @@ def _b12x_gdn_prefill_declaration(self, capacity: int): state_dtype=recurrent_state.dtype, checkpoint_export=True, null_state_index=0, - staging_key=(self._b12x_preparation_prefix, "gdn-prefill-staging"), - staging_resident_nbytes=resident_nbytes, ) return api.plan( caps, @@ -1148,7 +1136,6 @@ def _benchmark_b12x_gdn_prefill(self, state, capacity: int): def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): from b12x.preparation import PreparedCall - owners: tuple[torch.Tensor, ...] specs = tuple(state.layout.scratch_specs()) if len(specs) != 1: raise RuntimeError("b12x GDN prefill requires one scratch buffer") @@ -1159,11 +1146,30 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): # Trial and prepare factories own their scratch; the runtime path in # B12xGdnPrefill.run draws from the workspace manager instead. scratch = torch.empty(spec.shape, dtype=spec.dtype, device=device) + # Priming and candidate races use temporary inputs. Serving binds the + # caller's projection and output tensors directly, without retaining + # scheduler-capacity activation storage on every layer. + mixed_qkv = torch.empty( + ( + capacity, + (2 * self._b12x_local_key_heads + self._b12x_local_value_heads) * 128, + ), + dtype=torch.bfloat16, + device=device, + ) + a = torch.empty( + (capacity, self._b12x_local_value_heads), + dtype=torch.bfloat16, + device=device, + ) + b = torch.empty_like(a) + output = torch.empty( + (capacity, self._b12x_local_value_heads, 128), + dtype=torch.bfloat16, + device=device, + ) + owners: tuple[torch.Tensor, ...] if benchmark: - mixed_qkv = torch.empty_like(staging.mixed_qkv[:capacity]) - a = torch.empty_like(staging.a[:capacity]) - b = torch.empty_like(staging.b[:capacity]) - output = torch.empty_like(staging.output[:capacity]) cu_seqlens = torch.empty_like(staging.query_start_loc) indices = torch.full_like(staging.initial_indices, slot) final_indices = torch.full_like(staging.final_indices, slot) @@ -1187,12 +1193,6 @@ def _b12x_gdn_prefill_call(self, state, capacity: int, *, benchmark: bool): num_tokens, ) else: - mixed_qkv = staging.mixed_qkv[:capacity] - a, b, output = ( - staging.a[:capacity], - staging.b[:capacity], - staging.output[:capacity], - ) cu_seqlens = staging.query_start_loc indices, final_indices = staging.initial_indices, staging.final_indices checkpoint_indices, offsets = ( @@ -2346,6 +2346,15 @@ def _forward_core( raise RuntimeError( "b12x GDN prefill cache and runtime metadata must be bound" ) + prefill_output = torch.empty( + ( + conv_output_prefill.shape[0], + self._b12x_local_value_heads, + self.head_v_dim, + ), + device=conv_output_prefill.device, + dtype=conv_output_prefill.dtype, + ) runner.run( mixed_qkv=conv_output_prefill, a=a_prefill, @@ -2355,12 +2364,10 @@ def _forward_core( has_initial_state=prefill_has_initial_state, live_counts=attn_metadata.b12x_prefill_live_counts, checkpoint=attn_metadata.prefill_checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) - core_attn_out_non_spec = runner.output[ - : conv_output_prefill.shape[0] - ].unsqueeze(0) + core_attn_out_non_spec = prefill_output.unsqueeze(0) else: initial_state = ssm_state[prefill_state_indices] initial_state[~prefill_has_initial_state, ...] = 0 @@ -2989,6 +2996,7 @@ def _forward_core_b12x_mixed( query_start_loc=metadata.query_start_loc, metadata=metadata.convolution_metadata(rows), ).transpose(0, 1) + prefill_output = torch.empty_like(core_attn_out) runner.run( mixed_qkv=convolved, a=a.index_select(0, non_spec_indices), @@ -2998,22 +3006,22 @@ def _forward_core_b12x_mixed( has_initial_state=metadata.has_initial_state, live_counts=metadata.live_counts, checkpoint=metadata.checkpoint, - output=runner.output, + output=prefill_output, eps=self.layer_norm_epsilon, ) self._rms_norm_gated_cuda( - runner.output[:rows], + prefill_output, output_gate.index_select(0, non_spec_indices), - runner.output[:rows], + prefill_output, ) core_attn_out.zero_() width = core_attn_out.shape[-2] * core_attn_out.shape[-1] _scatter_b12x_gdn_output[(rows, triton.cdiv(width, 256))]( - runner.output, + prefill_output, non_spec_indices, metadata.live_counts, core_attn_out, - SOURCE_STRIDE=runner.output.stride(0), + SOURCE_STRIDE=prefill_output.stride(0), OUTPUT_STRIDE=core_attn_out.stride(0), WIDTH=width, BLOCK=256, diff --git a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py index b13ad1904238..67fd3fd419bd 100644 --- a/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py +++ b/vllm/model_executor/layers/mamba/ops/b12x_gdn_prefill.py @@ -69,16 +69,16 @@ def prefill_capacities(max_tokens: int) -> tuple[int, ...]: @dataclass(frozen=True) class GdnPrefillStaging: - """Reusable capacity buffers independent of a recurrent-state generation.""" + """Sequence metadata independent of a recurrent-state generation. + + Activations and outputs remain caller-owned. Native GDN bindings accept + fewer rows than a plan's capacity, so no padded activation copies are needed. + """ max_tokens: int max_seqs: int key_heads: int value_heads: int - mixed_qkv: torch.Tensor - a: torch.Tensor - b: torch.Tensor - output: torch.Tensor query_start_loc: torch.Tensor initial_indices: torch.Tensor final_indices: torch.Tensor @@ -97,26 +97,12 @@ def allocate( value_heads: int, device: torch.device, ) -> "GdnPrefillStaging": - mixed_qkv = torch.empty( - (max_tokens, (2 * key_heads + value_heads) * 128), - dtype=torch.bfloat16, - device=device, - ) - a = torch.empty((max_tokens, value_heads), dtype=torch.bfloat16, device=device) initial_indices = torch.zeros(max_seqs, dtype=torch.int32, device=device) return cls( max_tokens=max_tokens, max_seqs=max_seqs, key_heads=key_heads, value_heads=value_heads, - mixed_qkv=mixed_qkv, - a=a, - b=torch.empty_like(a), - output=torch.empty( - (max_tokens, value_heads, 128), - dtype=torch.bfloat16, - device=device, - ), query_start_loc=torch.zeros(max_seqs + 1, dtype=torch.int32, device=device), initial_indices=initial_indices, final_indices=torch.zeros_like(initial_indices), @@ -140,7 +126,7 @@ def is_compatible( and self.max_seqs >= max_seqs and self.key_heads == key_heads and self.value_heads == value_heads - and self.mixed_qkv.device == device + and self.query_start_loc.device == device ) @property @@ -148,10 +134,6 @@ def nbytes(self) -> int: return sum( tensor.numel() * tensor.element_size() for tensor in ( - self.mixed_qkv, - self.a, - self.b, - self.output, self.query_start_loc, self.initial_indices, self.final_indices, @@ -164,10 +146,10 @@ def nbytes(self) -> int: class B12xGdnPrefill: - """Layer-held GDN capacity family with reusable staging buffers. + """Layer-held GDN capacity family with reusable sequence metadata. The layer supplies one declared ``Plan`` for every admitted capacity. - This helper owns only capacity staging buffers; the recurrent-state pool + This helper owns only metadata staging buffers; the recurrent-state pool remains a binding-time caller resource and can be invalidated without discarding the plans or these buffers. """ @@ -217,10 +199,6 @@ def __init__( ): raise ValueError("GDN staging buffers do not cover the prepared capacity") self.staging = staging - self.mixed_qkv = staging.mixed_qkv - self.a = staging.a - self.b = staging.b - self.output = staging.output self.query_start_loc = staging.query_start_loc self.initial_indices = staging.initial_indices self.final_indices = staging.final_indices @@ -259,9 +237,6 @@ def run( capacity = next(capacity for capacity in self.capacities if rows <= capacity) plan = self.plans[capacity] (scratch,) = get_b12x_scratch_buffers(plan) - self.mixed_qkv[:rows].copy_(mixed_qkv) - self.a[:rows].copy_(a) - self.b[:rows].copy_(b) _stage_metadata[(1,)]( query_start_loc, state_indices, @@ -281,18 +256,18 @@ def run( MAX_SEQS=self.max_seqs, BLOCK=triton.next_power_of_2(self.max_seqs + 1), ) - q, k, v = self.mixed_qkv[:capacity].split( + q, k, v = mixed_qkv.split( (self.key_heads * 128, self.key_heads * 128, self.value_heads * 128), dim=-1, ) binding = self.api.bind( plan, scratch=scratch, - q=q.view(capacity, self.key_heads, 128), - k=k.view(capacity, self.key_heads, 128), - v=v.view(capacity, self.value_heads, 128), - a=self.a[:capacity], - b=self.b[:capacity], + q=q.view(rows, self.key_heads, 128), + k=k.view(rows, self.key_heads, 128), + v=v.view(rows, self.value_heads, 128), + a=a, + b=b, A_log=self.A_log, dt_bias=self.dt_bias, recurrent_state=self.recurrent_state, @@ -303,13 +278,12 @@ def run( checkpoint_offsets=self.checkpoint_offsets, num_seqs=self.num_seqs, num_tokens=self.num_tokens, - output=self.output[:capacity], + output=output[:rows], ) self.api.run( binding, scale=scale, eps=eps, - max_live_tokens=capacity, + max_live_tokens=rows, max_live_seqs=self.max_seqs, ) - output[:rows].copy_(self.output[:rows]) From de40bcd689a1e92f202fc7b6826d129ed368b97a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:15:40 +0000 Subject: [PATCH 068/365] fix(ds4): reserve attention scratch from prepared declarations State preparation reserves the actual decode and single-pass extend plans before KV admission. Remove the independent split-envelope reservation from metadata-free profiling while retaining model-lane ownership and locked workspace bounds. Validation: metadata plans require at most 282 MiB for TP2 K5 attention with 4096 prefill rows and 48 decode rows, versus a 1633 MiB independent split-envelope estimate. Added metadata-free profile regression coverage; native and complete-model qualification pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit bd00a4acf031745194b43401e30776e2ed37fa7e) --- tests/v1/worker/test_workspace.py | 45 ++++++++++ vllm/models/deepseek_v4/nvidia/b12x.py | 109 +------------------------ 2 files changed, 47 insertions(+), 107 deletions(-) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 1699fe90ea7f..d31428e706e7 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -313,3 +313,48 @@ def test_suspended_graph_resources_restore_collector_after_failure() -> None: raise ValueError("eager operation failed") assert workspace.retain_cuda_graph_capture_resource(owner) assert resources == [owner] + + +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( + monkeypatch, image_tokens +) -> None: + """Prepared state declarations own attention scratch before KV admission.""" + from b12x.attention import compressed_sparse_mla + + from vllm.models.deepseek_v4.nvidia import b12x as dsv4 + + requested = [] + monkeypatch.setattr( + workspace, + "_manager", + SimpleNamespace(get_simultaneous=lambda *specs: requested.append(specs)), + ) + monkeypatch.setattr( + dsv4, "_require_b12x_compressed_sparse_mla", lambda: compressed_sparse_mla + ) + monkeypatch.setattr( + dsv4, "get_forward_context", lambda: SimpleNamespace(attn_metadata=None) + ) + layer = dsv4.DeepseekV4B12xAttention.__new__(dsv4.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.compress_ratio = 128 + layer.max_model_len = 1048576 + layer.max_num_batched_tokens = 4096 + layer.window_size = 128 + layer.max_image_tokens = image_tokens + layer.swa_cache_layer = SimpleNamespace(block_size=256) + layer.vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + use_dspark=lambda: True, num_speculative_tokens=5 + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4096, max_num_seqs=8), + ) + q = torch.randn((3, 32, 512), dtype=torch.bfloat16) + original_q = q.clone() + output = torch.full_like(q, float("nan")) + layer.forward_mqa(q, torch.empty(0), torch.arange(3), output) + + torch.testing.assert_close(q, original_q, rtol=0, atol=0) + assert torch.count_nonzero(output) == 0 + assert requested == [] diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 232d84b2dc94..14512c22cde3 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -4,9 +4,7 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass, replace -from functools import cache from typing import TYPE_CHECKING, Any, ClassVar, cast import torch @@ -60,7 +58,6 @@ DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.worker.workspace import ( - current_workspace_manager, retain_cuda_graph_capture_resource, ) @@ -558,24 +555,6 @@ def _c128a_profile_widths(max_width: int) -> tuple[int, ...]: return tuple(sorted(widths)) -@cache -def _max_q_chunks( - max_rows: int, - width: int, - split_chunks_for_contract: Callable[..., int], - decode_row_capacity: int | None, -) -> int: - return max( - rows - * split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - for rows in range(1, max(int(max_rows), 1) + 1) - ) - - def _cache_page_view( cache: torch.Tensor, page_size: int, @@ -1289,16 +1268,6 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: out = tensor_model_parallel_all_reduce(out) return out - def reserve_profile_scratch(self) -> None: - super().reserve_profile_scratch() - device = self.q_norm.weight.device - if device.type == "cuda": - # Reserve every layer type before compiled profile execution can - # skip its attention body. This tensor describes heads/device only. - self._reserve_profile_workspace( - torch.empty((0, self.padded_heads, _DSV4_HEAD_DIM), device=device) - ) - def _get_cache_page_view( self, cache: torch.Tensor, @@ -1312,81 +1281,6 @@ def _get_cache_page_view( self._b12x_cache_page_views[key] = view return view - def _reserve_profile_workspace(self, q: torch.Tensor) -> None: - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) - - module = _require_b12x_compressed_sparse_mla() - indexed_width = 0 - if self.compress_ratio == 4: - if self.topk_indices_buffer is not None: - indexed_width = int(self.topk_indices_buffer.shape[-1]) - elif self.indexer is not None: - indexed_width = int(self.indexer.topk_tokens) - elif self.compress_ratio > 1: - indexed_width = _c128a_topk_width( - self.max_model_len, - self.compress_ratio, - ) - - indexed_widths = ( - _c128a_profile_widths(indexed_width) - if self.compress_ratio > 4 - else (indexed_width,) - ) - swa_widths = { - int(self.window_size), - int(self.window_size) + self.max_image_tokens, - } - speculative_config = self.vllm_config.speculative_config - if speculative_config is not None and speculative_config.use_dspark(): - swa_widths.add( - get_dspark_swa_index_width( - int(self.window_size), - speculative_config.num_speculative_tokens or 0, - ) - ) - swa_width = max(swa_widths) - width = max(swa_width + indexed_width, 1) - rows = max(int(self.max_num_batched_tokens), 1) - decode_row_capacity = _get_dspark_decode_row_capacity(self.vllm_config) - max_chunks_per_row = module.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=decode_row_capacity, - ) - # Split count is not monotonic in index width: a shorter prefix may - # use 12-token chunks while a longer prefix uses 64-token chunks. - # Cover every metadata width, not only the longest supported context. - max_q_chunks = max( - _max_q_chunks( - rows, - max(swa + indexed, 1), - module.split_chunks_for_contract, - decode_row_capacity, - ) - for swa in swa_widths - for indexed in indexed_widths - ) - # Allocation profiling needs geometry, not an executable declaration - # tied to live cache storage or an autotuning preparation session. - plan = plan_compressed_sparse_mla_scratch( - module.Caps( - device=q.device, - num_q_heads=int(q.shape[1]), - max_q_rows=rows, - max_width=width, - head_dim=_DSV4_HEAD_DIM, - v_head_dim=_DSV4_HEAD_DIM, - page_size=int(self.swa_cache_layer.block_size), - max_chunks_per_row=max_chunks_per_row, - max_q_chunks=max_q_chunks, - decode_row_capacity=decode_row_capacity, - ) - ) - current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) - def forward_mqa( self, q: torch.Tensor, @@ -1404,7 +1298,8 @@ def forward_mqa( attn_metadata = get_forward_context().attn_metadata if attn_metadata is None: output.zero_() - self._reserve_profile_workspace(q) + # State preparation reserves each declared operation's scratch + # before KV admission. This metadata-free pass skips attention. return assert isinstance(attn_metadata, dict) From af3ebecdcea914eb358f3698023a13db2d621246 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 04:20:04 +0000 Subject: [PATCH 069/365] test(ds4): retain query padding reservation with prepared scratch Verify that the B12X attention class inherits per-microbatch padded query reservation while attention scratch remains owned by prepared declarations. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 842bac6e0232a7796a01c9a3d0456e277a2ed9d9) --- tests/v1/worker/test_workspace.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index d31428e706e7..e5df32783a15 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -358,3 +358,31 @@ def test_dsv4_metadata_free_profile_does_not_reserve_split_attention( torch.testing.assert_close(q, original_q, rtol=0, atol=0) assert torch.count_nonzero(output) == 0 assert requested == [] + + +def test_dsv4_profile_retains_padded_query_reservation() -> None: + from vllm.models.deepseek_v4.attention import DeepseekV4Attention + from vllm.models.deepseek_v4.nvidia.b12x import DeepseekV4B12xAttention + + assert ( + DeepseekV4B12xAttention.reserve_profile_scratch + is DeepseekV4Attention.reserve_profile_scratch + ) + layer = DeepseekV4B12xAttention.__new__(DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.kv_cache_torch_dtype = torch.uint8 + device = torch.device("cuda:0") + layer.q_norm = SimpleNamespace(weight=SimpleNamespace(device=device)) + layer._q_padded_scratch_num_ubatches = 2 + layer.max_num_batched_tokens = 4096 + layer.padded_heads = 32 + layer.head_dim = 512 + layer._q_padded_scratch_dtype = torch.bfloat16 + reserved = [] + layer._reserve_q_padded_scratch_buffer = lambda *args: reserved.append(args) + + layer.reserve_profile_scratch() + + assert reserved == [ + (4096, 32, 512, torch.bfloat16, device, ubatch) for ubatch in range(2) + ] From b70649521e6c467f4a97e053b4fc7f8a4ee662ae Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 06:23:19 +0000 Subject: [PATCH 070/365] test(ds4): validate declared attention workspace envelopes Exercise compressed-prefix, speculative, and multimodal scratch declarations through the preparation coordinator. Verify that both microbatch workspaces cover serving shapes after locking, without invoking the removed legacy split-envelope allocator. Import SimpleNamespace in the standalone workspace suite. Production code and serving behavior are unchanged. Validation: standalone test failure reproducer and pre-commit hooks. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 4b7b86dd057e334bc5fb953eea732ee395bfdc15) --- .../attention/test_ds4_workspace_envelope.py | 149 +++++++++++------- 1 file changed, 94 insertions(+), 55 deletions(-) diff --git a/tests/v1/attention/test_ds4_workspace_envelope.py b/tests/v1/attention/test_ds4_workspace_envelope.py index 404915c1469f..71c8a568b776 100644 --- a/tests/v1/attention/test_ds4_workspace_envelope.py +++ b/tests/v1/attention/test_ds4_workspace_envelope.py @@ -1,76 +1,115 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from math import prod from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from vllm.models.deepseek_v4.nvidia import b12x as adapter +from vllm.utils.b12x import B12xWorkload +from vllm.v1.worker import workspace as workspace_module +from vllm.v1.worker.b12x_startup import B12xPreparationCoordinator +from vllm.v1.worker.workspace import WorkspaceManager @pytest.mark.skipif(not torch.cuda.is_available(), reason="B12X device plan required") @pytest.mark.parametrize("max_length", [786_688, 1_048_576]) @pytest.mark.parametrize("drafts", [0, 5]) -def test_profile_covers_shorter_compressed_prefixes(monkeypatch, max_length, drafts): - """Shorter prefixes can select more split chunks than the longest prefix.""" - from b12x.attention import compressed_sparse_mla as mla - from b12x.attention.compressed_sparse_mla._scratch import ( - plan_compressed_sparse_mla_scratch, - ) +@pytest.mark.parametrize("image_tokens", [0, 384]) +def test_declared_workspace_covers_shorter_compressed_prefixes( + monkeypatch, max_length, drafts, image_tokens +): + """Reserve decode-prefix and capacity-prefill plans before workspace lock. + Native plan scratch requirements are evaluated without compiling attention + kernels. The job stub completes only the coordinator's control protocol; + declaration, plan selection, and workspace reservation are real. + """ + device = torch.device("cuda", torch.accelerator.current_device_index()) spec = SimpleNamespace(use_dspark=lambda: True, num_speculative_tokens=drafts) - layer = SimpleNamespace( - compress_ratio=128, + layer = adapter.DeepseekV4B12xAttention.__new__(adapter.DeepseekV4B12xAttention) + torch.nn.Module.__init__(layer) + layer.prefix = "model.layers.0.self_attn" + layer.indexer = None + layer.compress_ratio = 128 + layer.max_model_len = max_length + layer.max_num_batched_tokens = 4096 + layer.max_image_tokens = image_tokens + layer.window_size = 128 + layer.padded_heads = 32 + layer._b12x_cache_page_views = {} + layer._b12x_mla_plans = {} + layer.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=256), + speculative_config=spec if drafts else None, + scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=4096), + ) + page_bytes = adapter._DSV4_CACHE_BYTES_PER_TOKEN + layer.swa_cache_layer = SimpleNamespace( + block_size=256, + kv_cache=torch.empty((2, 256 * page_bytes), dtype=torch.uint8, device=device), + ) + layer.kv_cache = torch.empty((2, 2 * page_bytes), dtype=torch.uint8, device=device) + workload = B12xWorkload( + stage="state", + token_counts=(1, 8, 48, 128, 256, 4096), + fixed_token_counts=(1, 8, 48, 128, 256), + output_dtype=torch.bfloat16, + max_tokens=4096, + max_seqs=8, max_model_len=max_length, - max_num_batched_tokens=4096, - max_image_tokens=0, - window_size=128, - vllm_config=SimpleNamespace( - speculative_config=spec if drafts else None, - scheduler_config=SimpleNamespace( - max_num_seqs=8, max_num_batched_tokens=4096 - ), - ), - swa_cache_layer=SimpleNamespace(block_size=256), + speculative_tokens=drafts, ) - reserved = [] - monkeypatch.setattr( - adapter, - "current_workspace_manager", - lambda: SimpleNamespace( - get_simultaneous=lambda *items: reserved.append( - sum(prod(shape) * dtype.itemsize for shape, dtype in items) - ) - ), + (unit,) = layer.get_b12x_preparation_units(layer, workload) + queries = [request.plan.query for request in unit.requests] + widths = {128, 256, 512, 1024, 2048, 4096} + widths.add(6272 if max_length == 786_688 else 8192) + assert {q.indexed_width for q in queries} == widths + assert {q.query_rows for q in queries if q.mode == "decode"} == set(range(1, 257)) + assert {q.query_rows for q in queries if q.mode == "extend"} == {4096} + + job = Mock() + job.advance.return_value = SimpleNamespace( + done=True, pending_compilation=False, ready_collectives=() ) - adapter.DeepseekV4B12xAttention._reserve_profile_workspace( - layer, torch.empty((0, 32, 512), device="cuda") + session = Mock() + session.begin.return_value = job + workspace = WorkspaceManager(device, num_ubatches=2) + coordinator = B12xPreparationCoordinator( + session, + [(unit.requests, False)], + global_rank=0, + world_group=None, + process_local_only=True, + workspace=workspace, ) - for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): - for index_width in (128, 256, 512, 1024, 2048, 4096): - width = 128 + index_width - capacity = 8 * (1 + drafts) if drafts else None - plan = plan_compressed_sparse_mla_scratch( - mla.Caps( - device=torch.device("cuda"), - num_q_heads=32, - max_q_rows=rows, - max_width=width, - head_dim=512, - v_head_dim=512, - page_size=256, - decode_row_capacity=capacity, - max_chunks_per_row=mla.split_chunks_for_contract( - rows=rows, - width=width, - decode_row_capacity=capacity, - ), + outcome = coordinator.advance() + assert outcome["done"] and outcome["error"] is None, outcome + workspace.lock() + pointers = [buffer.data_ptr() for buffer in workspace._current_workspaces] + sizes = [buffer.numel() for buffer in workspace._current_workspaces] + assert sizes[0] == sizes[1] > 0 + + for ubatch in range(2): + monkeypatch.setattr( + workspace_module, "dbo_current_ubatch_id", lambda ubatch=ubatch: ubatch + ) + for rows in (1, 48, 128, 240, 256, 257, 1024, 4096): + mode = "decode" if rows <= 256 else "extend" + swa_widths = {q.swa_width for q in queries if q.mode == mode} + for swa_width in swa_widths: + for index_width in widths: + plan = layer._b12x_mla_plan( + mode, + rows, + torch.empty((rows, swa_width), device="meta"), + torch.empty((rows, index_width), device="meta"), + ) + workspace.get_simultaneous( + *((spec.shape, spec.dtype) for spec in plan.scratch_specs()) + ) + assert ( + workspace._current_workspaces[ubatch].data_ptr() == pointers[ubatch] ) - ) - required = sum( - prod(shape) * dtype.itemsize - for shape, dtype in plan.shapes_and_dtypes() - ) - assert max(reserved) >= required, (rows, width, reserved, required) From b08616024fe52f9ad385dd55d6df005bdc893b21 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 09:53:36 +0000 Subject: [PATCH 071/365] fix(deepseek-v41): preserve tool namespaces and reminder messages Canonicalize explicit namespaces before Chat Completions validation can discard them. Keep definitions, named choices, assistant history and streamed responses on the same qualified function identity. Reject conflicting declarations without mutating input dictionaries, and accept the reference latest_reminder role in the V4.1 tokenizer. Retain ordinary tool names, reasoning budgets, image ordering and the existing parser. Add pinned reference prompt fixtures, API validation and streaming roundtrip coverage. Validation: 179 unit tests, pre-commit hooks, TP4 temperature-1 Chat Completions tool/history checks and image smoke checks pass. Reference contract: Qizhou Guo / DeepSeek, HF DeepSeek-V4.1-Flash dba1be0a40aa45a94ad051997016db3960a90277. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5b5dd1c055d4f9561d7bbfe221f004da73815448) --- docs/features/tool_calling.md | 42 ++ .../openai/test_tool_calls_serialization.py | 162 ++++++- tests/parser/engine/test_deepseek_v4.py | 93 ++++ .../fixtures/deepseek_v41/encoding.json | 410 ++++++++++++++++++ tests/tokenizers_/test_deepseek_v41.py | 66 +++ .../openai/chat_completion/protocol.py | 33 +- vllm/tokenizers/deepseek_v41.py | 9 +- vllm/tokenizers/deepseek_v41_encoding.py | 8 +- vllm/utils/tool_names.py | 67 +++ 9 files changed, 882 insertions(+), 8 deletions(-) create mode 100644 tests/tokenizers_/fixtures/deepseek_v41/encoding.json create mode 100644 vllm/utils/tool_names.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 9628bc431a1e..2e8178d392ad 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -338,6 +338,48 @@ Supported models: Flags: `--tool-call-parser deepseek_v31 --chat-template {see_above}` +### DeepSeek-V4.1 Models (`deepseek_v41`) + +Status: implemented. The `deepseek_v41` tokenizer and parser support spaced +DSML tool calls and namespace-qualified function names for +`deepseek-ai/DeepSeek-V4.1-Flash`. + +Use `--tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 +--reasoning-parser deepseek_v41 --enable-auto-tool-choice`. + +Chat Completions accepts a namespace on a function tool: + +```json +{ + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "lookup", + "description": "Look up an item.", + "parameters": {"type": "object", "properties": {"sku": {"type": "string"}}} + } +} +``` + +The namespace may also be a string, or appear inside `function`. Before +validation and rendering, it becomes `function.name="inventory::lookup"`. +The namespace description is prepended to the function description. +The same normalization applies to assistant tool-call history and named +`tool_choice` objects. A qualified name must agree with any explicit namespace; +conflicting declarations are rejected. Tools without namespaces are unchanged. + +Complete and streamed Chat Completions responses retain the qualified +`function.name`; they do not add a separate response `namespace` field. Dispatch +tools by that qualified name and preserve it when sending assistant history. +This extension is separate from the Responses API's `type="namespace"` schema. + +The tokenizer also accepts `{"role":"latest_reminder","content":"..."}` +between a system message and a user message, rendering the reference +`<｜latest_reminder｜>` marker for date, locale or similar context. This is a +DeepSeek-specific role, not a standard OpenAI message role. Other unknown roles +remain invalid. The encoding contract follows the +[DeepSeek V4.1 reference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/dba1be0a40aa45a94ad051997016db3960a90277/encoding/README.md). + ### OpenAI OSS Models (`openai`) Supported models: diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index cedc2575b80e..478379f5e13b 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -15,9 +15,16 @@ "ValueError: Unexpected tool call id ...". """ +import copy + import pytest +from pydantic import ValidationError -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.exceptions import VLLMValidationError def _make_tool_call(tc_id: str, name: str, args: str) -> dict: @@ -148,3 +155,156 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): # Verify after model_dump_json too _ = req.model_dump_json() assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +@pytest.mark.parametrize("iterator", [False, True]) +def test_namespaces_survive_definitions_choices_and_history( + location, namespace, iterator +): + def qualify(item, value): + target = item if location == "tool" else item["function"] + target["namespace"] = value + return item + + tools = [ + qualify({"type": "function", "function": {"name": "lookup"}}, namespace), + qualify({"type": "function", "function": {"name": "lookup"}}, "billing"), + ] + calls = [qualify(_make_tool_call("call_1", "lookup", "{}"), namespace)] + payload = { + "model": "DeepSeek-V4.1-Flash", + "tools": tools, + "tool_choice": qualify( + {"type": "function", "function": {"name": "lookup"}}, namespace + ), + "messages": [ + {"role": "user", "content": "Look up stock."}, + {"role": "assistant", "content": None, "tool_calls": calls}, + ], + } + original = copy.deepcopy(payload) + if iterator: + payload["messages"][1]["tool_calls"] = iter(calls) + request = ChatCompletionRequest.model_validate(payload) + assert [tool.function.name for tool in request.tools] == [ + "inventory::lookup", + "billing::lookup", + ] + assert request.tool_choice.function.name == "inventory::lookup" + assert request.messages[1]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + serialized = request.model_dump_json() + assert request.model_dump_json() == serialized + restored = ChatCompletionRequest.model_validate_json(serialized) + assert restored.model_dump_json() == serialized + if not iterator: + assert payload == original + assert calls == original["messages"][1]["tool_calls"] + + +def test_namespace_description_is_preserved_exactly_once(): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": "Stock operations."}, + "function": { + "name": "inventory::lookup", + "description": "Find a SKU.", + "strict": True, + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == "Stock operations.\nFind a SKU." + assert parsed.function.strict is True + assert tool == original + assert ChatCompletionToolsParam.model_validate( + parsed.model_dump() + ).model_dump() == (parsed.model_dump()) + + +@pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) +@pytest.mark.parametrize( + "namespace, name", + [ + ("inventory", "billing::lookup"), + ("a::b", "lookup"), + ("", "lookup"), + ({}, "lookup"), + ({"name": 7}, "lookup"), + (False, "lookup"), + ("inventory", "inventory::x::lookup"), + ("inventory", "inventory::"), + ], +) +def test_invalid_namespaces_are_client_errors(namespace, name, tool_choice): + with pytest.raises(ValidationError, match="[Nn]amespace"): + ChatCompletionRequest.model_validate( + { + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "namespace": namespace, + "function": {"name": name}, + } + ], + "tool_choice": tool_choice, + } + ) + + +def test_conflicting_namespace_locations_are_rejected(): + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup", "namespace": "billing"}, + } + ) + + +def test_consistent_namespace_locations_do_not_duplicate_the_prefix(): + parsed = ChatCompletionToolsParam.model_validate( + { + "type": "function", + "namespace": {"name": "inventory"}, + "function": {"name": "inventory::lookup", "namespace": "inventory"}, + } + ) + assert parsed.function.name == "inventory::lookup" + + +def test_conflicting_history_namespace_is_not_silently_dropped(): + call = _make_tool_call("call_a", "inventory::lookup", "{}") + call["namespace"] = "billing" + with pytest.raises(ValidationError, match="Conflicting tool namespaces"): + _make_request([{"role": "assistant", "content": None, "tool_calls": [call]}]) + + +def test_bare_choice_cannot_select_a_namespaced_tool(): + with pytest.raises(VLLMValidationError, match="does not match"): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "lookup"}], + "tools": [ + { + "type": "function", + "namespace": "inventory", + "function": {"name": "lookup"}, + } + ], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + } + ) + + +def test_plain_and_already_qualified_names_keep_their_spelling(): + for name in ("lookup", "inventory::lookup", "opaque::legacy::name"): + tool = {"type": "function", "function": {"name": name}} + assert ChatCompletionToolsParam.model_validate(tool).function.name == name diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 7cbf9791a7b0..9180a6f4b2d5 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -21,6 +21,7 @@ simulate_reasoning_streaming, simulate_tool_streaming, ) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.deepseek_v4 import ( DSML_INVOKE_END, @@ -36,10 +37,13 @@ _unwrap_wrapper_args, deepseek_v4_config, ) +from vllm.parser.deepseek_v41 import DeepSeekV41Parser +from vllm.parser.engine.adapters import make_adapters from vllm.parser.engine.registered_adapters import ( DeepSeekV4ParserReasoningAdapter, DeepSeekV4ParserToolAdapter, ) +from vllm.tokenizers.deepseek_v41_encoding import encode_messages _THINK_START_ID = 50 _THINK_END_ID = 51 @@ -48,6 +52,95 @@ _PARAM_CLOSE = "" +@pytest.mark.parametrize("chunk_size", [1, 2, 7, 17, 1000]) +def test_v41_namespaced_tools_keep_identity_in_full_stream_and_history(chunk_size): + names = ["inventory::lookup", "billing::lookup"] + request = ChatCompletionRequest( + messages=[{"role": "user", "content": "Look up A1."}], + tools=[ + { + "type": "function", + "namespace": namespace, + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": {"sku": {"type": "string"}}, + }, + }, + } + for namespace in ("inventory", "billing") + ], + ) + tokenizer = make_mock_tokenizer({"": 50, "": 51}) + _, adapter = make_adapters(DeepSeekV41Parser) + + def parser(): + return adapter( + tokenizer, request.tools, chat_template_kwargs={"thinking": False} + ) + + text = ( + "<｜DSML｜ calls>\n" + + "\n".join( + f'<｜DSML｜ invoke name="{name}">\n' + '<｜DSML｜ parameter name="sku" string="true">A1\n' + "" + for name in names + ) + + "\n" + ) + calls = parser().extract_tool_calls(text, request).tool_calls + assert [call.function.name for call in calls] == names + assert [json.loads(call.function.arguments) for call in calls] == [ + {"sku": "A1"} + ] * 2 + stream = parser() + collected: dict[int, dict[str, str]] = {} + + def collect(delta): + if delta is None: + return + for call in delta.tool_calls: + value = collected.setdefault(call.index, {"name": "", "arguments": ""}) + if call.function: + value["name"] += call.function.name or "" + value["arguments"] += call.function.arguments or "" + + previous = "" + for offset in range(0, len(text), chunk_size): + chunk = text[offset : offset + chunk_size] + collect( + stream.extract_tool_calls_streaming( + previous, + previous + chunk, + chunk, + (), + (), + (), + request, + ) + ) + previous += chunk + collect(stream.finish_streaming()) + assert [collected[i]["name"] for i in sorted(collected)] == names + assert [json.loads(collected[i]["arguments"]) for i in sorted(collected)] == ( + [{"sku": "A1"}] * 2 + ) + prompt = encode_messages( + [ + {"role": "user", "content": "Look up A1."}, + { + "role": "assistant", + "content": "", + "tool_calls": [call.model_dump(exclude_none=True) for call in calls], + }, + ], + thinking_mode="chat", + ) + assert all(f'name="{name}"' in prompt for name in names) + + def _param(name: str, is_str: str, value: str) -> str: return f"<{_PARAM_OPEN.format(name=name, is_str=is_str)}{value}{_PARAM_CLOSE}" diff --git a/tests/tokenizers_/fixtures/deepseek_v41/encoding.json b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json new file mode 100644 index 000000000000..532e842e96aa --- /dev/null +++ b/tests/tokenizers_/fixtures/deepseek_v41/encoding.json @@ -0,0 +1,410 @@ +{ + "reference": "https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/raw/dba1be0a40aa45a94ad051997016db3960a90277/encoding/encoding.py", + "reference_author": "DeepSeek; namespace contract by Qizhou Guo", + "cases": [ + { + "name": "reminder_chat", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>You are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "reminder_thinking", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "latest_reminder", + "content": "Date: 2026-09-17; locale: cs-CZ." + }, + { + "role": "user", + "content": "Hello." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\nYou are helpful.<｜latest_reminder｜>Date: 2026-09-17; locale: cs-CZ.<｜User｜>Hello.<｜Assistant｜>" + }, + { + "name": "mid_system_chat", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "mid_system_thinking", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi." + }, + { + "role": "system", + "content": "Answer briefly." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜System｜>Answer briefly.<｜Assistant｜>" + }, + { + "name": "reasoning_history_chat", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "reasoning_history_thinking", + "messages": [ + { + "role": "user", + "content": "What is 2+2?" + }, + { + "role": "assistant", + "reasoning_content": "Add the integers.", + "content": "4" + }, + { + "role": "user", + "content": "And three times that?" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And three times that?<｜Assistant｜>" + }, + { + "name": "interleaved_images_chat", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "interleaved_images_thinking", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "first" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/a.png" + } + }, + { + "type": "text", + "text": "second" + }, + { + "type": "image_url", + "image_url": { + "url": "/unused/b.png" + } + } + ] + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>first\n\n<｜deepseek_image｜>\n\nsecond\n\n<｜deepseek_image｜><｜Assistant｜>" + }, + { + "name": "namespaced_tool_history_false", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "namespaced_tool_history_true", + "messages": [ + { + "role": "user", + "content": "Check A1." + }, + { + "role": "assistant", + "reasoning_content": "Consult both tools.", + "content": "", + "tool_calls": [ + { + "id": "inventory", + "type": "function", + "namespace": "inventory", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + }, + { + "id": "billing", + "type": "function", + "namespace": "billing", + "function": { + "name": "lookup", + "arguments": "{\"sku\":\"A1\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "billing", + "content": "Paid." + }, + { + "role": "tool", + "tool_call_id": "inventory", + "content": "In stock." + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"inventory::lookup\", \"description\": \"inventory operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n{\"name\": \"billing::lookup\", \"description\": \"billing operations.\\nLook up A1.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"sku\": {\"type\": \"string\"}}}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>Check A1.<｜Assistant｜>Consult both tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"inventory::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜DSML｜ invoke name=\"billing::lookup\">\n<｜DSML｜ parameter name=\"sku\" string=\"true\">A1\n\n<｜end▁of▁sentence｜><｜User｜>In stock.\n\nPaid.<｜Assistant｜>", + "tools": [ + { + "type": "function", + "namespace": { + "name": "inventory", + "description": "inventory operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + }, + { + "type": "function", + "namespace": { + "name": "billing", + "description": "billing operations." + }, + "function": { + "name": "lookup", + "description": "Look up A1.", + "parameters": { + "type": "object", + "properties": { + "sku": { + "type": "string" + } + } + } + } + } + ] + }, + { + "name": "action_task", + "messages": [ + { + "role": "user", + "content": "Find the documentation.", + "task": "action" + } + ], + "kwargs": { + "thinking": true + }, + "expected": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>Find the documentation.<｜Assistant｜><｜action｜>" + }, + { + "name": "title_task", + "messages": [ + { + "role": "user", + "content": "Hello." + }, + { + "role": "assistant", + "content": "Hi.", + "task": "title" + } + ], + "kwargs": { + "thinking": false + }, + "expected": "<｜begin▁of▁sentence｜><｜User｜>Hello.<｜Assistant｜>Hi.<｜end▁of▁sentence｜><｜title｜>" + } + ] +} diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index d0916fb6077b..5035a7e71431 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/tests/tokenizers_/test_deepseek_v41.py @@ -4,12 +4,14 @@ import copy import json from pathlib import Path +from typing import Any import pytest from tests.tokenizers_.test_deepseek_v4 import FakeHfTokenizer from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tokenizers.deepseek_v41 import get_deepseek_v41_tokenizer +from vllm.tokenizers.deepseek_v41_encoding import encode_messages FIXTURES = Path(__file__).parent / "fixtures" @@ -185,3 +187,67 @@ def image_part(color): def test_unsupported_media_is_a_request_error(): with pytest.raises(ValueError, match="text and image content only"): render([{"role": "user", "content": [{"type": "input_audio"}]}]) + + +_FIXTURES = json.loads( + (Path(__file__).parent / "fixtures/deepseek_v41/encoding.json").read_text() +) + + +@pytest.mark.parametrize("case", _FIXTURES["cases"], ids=lambda case: case["name"]) +def test_prompt_matches_reference_fixture(case): + original = copy.deepcopy(case["messages"]) + request = ChatCompletionRequest( + messages=case["messages"], + tools=case.get("tools"), + ) + prompt = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tools=[tool.model_dump(exclude_none=True) for tool in request.tools] + if request.tools + else None, + tokenize=False, + **case["kwargs"], + ) + assert prompt == case["expected"] + assert case["messages"] == original + + +@pytest.mark.parametrize("location", ["tool", "function"]) +@pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) +def test_direct_encoder_and_api_preserve_the_same_tool_identity(location, namespace): + tool: dict[str, Any] = {"type": "function", "function": {"name": "lookup"}} + call: dict[str, Any] = { + "id": "call_a", + "type": "function", + "function": {"name": "lookup", "arguments": '{"sku":"A1"}'}, + } + for item in (tool, call): + (item if location == "tool" else item["function"])["namespace"] = namespace + messages = [ + {"role": "system", "content": "Use tools.", "tools": [tool]}, + {"role": "user", "content": "Find A1."}, + {"role": "assistant", "content": "", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call_a", "content": "In stock."}, + ] + original = copy.deepcopy(messages) + direct = encode_messages(messages, thinking_mode="chat") + request = ChatCompletionRequest(messages=messages) + rendered = get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + request.messages, + tokenize=False, + thinking=False, + ) + assert rendered == direct + assert '"name": "inventory::lookup"' in rendered + assert 'name="inventory::lookup"' in rendered + assert "In stock." in rendered + assert messages == original + + +def test_reminder_does_not_enable_arbitrary_message_roles(): + with pytest.raises(ValueError, match="Invalid role: SYSTEM"): + get_deepseek_v41_tokenizer(FakeHfTokenizer()).apply_chat_template( + [{"role": "SYSTEM", "content": "Hello"}], + tokenize=False, + ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index badd5a0344c6..f9d07659cdb8 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -51,6 +51,7 @@ ThinkingTokenBudget, ) from vllm.utils import random_uuid +from vllm.utils.tool_names import normalize_tool_namespace logger = init_logger(__name__) @@ -186,6 +187,11 @@ class ChatCompletionToolsParam(OpenAIBaseModel): function: FunctionDefinition defer_loading: bool | None = None + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + @model_validator(mode="after") def _propagate_defer_loading(self) -> "ChatCompletionToolsParam": if self.defer_loading is not None and self.function.defer_loading is None: @@ -209,6 +215,11 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): function: ChatCompletionNamedFunction type: Literal["function"] = "function" + @model_validator(mode="before") + @classmethod + def _normalize_namespace(cls, data: Any) -> Any: + return normalize_tool_namespace(data) + class ChatCompletionRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation @@ -534,23 +545,31 @@ def check_cache_salt_support(cls, data: Any) -> Any: def _normalize_messages_before(cls, data: Any) -> Any: """Pre-process message dicts before Pydantic field validation. - Performs two normalizations in a single pass: + Normalizes tool identities and message fields in a single pass: - Converts tool_calls generators/iterators to lists so one-shot generators are not consumed during union type matching. - Renames the deprecated ``reasoning_content`` field to ``reasoning`` so downstream code only needs to check one field. + - Qualifies tool-call namespaces before union validation can drop them. """ if not isinstance(data, dict): return data messages = data.get("messages") if not isinstance(messages, list): return data + data = dict(data) + messages = [dict(msg) if isinstance(msg, dict) else msg for msg in messages] + data["messages"] = messages for msg in messages: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None and not isinstance(tool_calls, list): - msg["tool_calls"] = list(tool_calls) + if tool_calls is not None: + msg["tool_calls"] = [ + normalize_tool_namespace(call) for call in tool_calls + ] + if isinstance(msg.get("tools"), list): + msg["tools"] = [normalize_tool_namespace(tool) for tool in msg["tools"]] reasoning_content = msg.pop("reasoning_content", None) if reasoning_content is not None and msg.get("reasoning") is None: msg["reasoning"] = reasoning_content @@ -916,6 +935,14 @@ def check_tool_usage(cls, data): if not isinstance(data, dict): return data + # Named choices must be compared with qualified identities before + # Pydantic serializes function definitions and discards extra fields. + data = dict(data) + if isinstance(data.get("tools"), list): + data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + if isinstance(data.get("tool_choice"), dict): + data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) + # Reject empty tools array, matching OpenAI API behavior if data.get("tools") == []: raise VLLMValidationError( diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index 3405cbbbe7dc..68ba93e5ec8a 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -23,7 +23,14 @@ def _normalize_messages( result = [dict(message) for message in copy.deepcopy(messages)] for message in result: role = message.get("role") - if role not in ("system", "developer", "user", "assistant", "tool"): + if role not in ( + "system", + "developer", + "user", + "assistant", + "tool", + "latest_reminder", + ): raise ValueError(f"Invalid role: {role}") if role == "developer": message["role"] = "system" diff --git a/vllm/tokenizers/deepseek_v41_encoding.py b/vllm/tokenizers/deepseek_v41_encoding.py index 995115f1f7d0..89907d152bef 100644 --- a/vllm/tokenizers/deepseek_v41_encoding.py +++ b/vllm/tokenizers/deepseek_v41_encoding.py @@ -11,6 +11,8 @@ import json from typing import Any, Dict, List, Optional, Union +from vllm.utils.tool_names import normalize_tool_namespace + bos_token: str = "<｜begin▁of▁sentence｜>" eos_token: str = "<｜end▁of▁sentence｜>" thinking_start_token: str = "" @@ -55,8 +57,8 @@ def to_json(value: Any) -> str: def tools_from_openai_format(tools): - """Extract function definitions from OpenAI-format tool list.""" - return [tool["function"] for tool in tools] + """Extract function definitions with namespace-qualified identities.""" + return [normalize_tool_namespace(tool)["function"] for tool in tools] def tool_calls_from_openai_format(tool_calls): @@ -66,7 +68,7 @@ def tool_calls_from_openai_format(tool_calls): "name": tool_call["function"]["name"], "arguments": tool_call["function"]["arguments"], } - for tool_call in tool_calls + for tool_call in map(normalize_tool_namespace, tool_calls) ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py new file mode 100644 index 000000000000..02fcbc9e7c2e --- /dev/null +++ b/vllm/utils/tool_names.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Canonical function identities for namespace-bearing chat tools. + +The namespace spelling follows Qizhou Guo's DeepSeek V4.1 reference encoder: +https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/commit/dba1be0a40aa45a94ad051997016db3960a90277 +""" + +from typing import Any + + +def _namespace_parts(namespace: Any) -> tuple[str, str | None]: + description = None + if isinstance(namespace, dict): + description = namespace.get("description") + namespace = namespace.get("name") + if not isinstance(namespace, str) or not namespace or "::" in namespace: + raise ValueError("Tool namespace must be a nonempty string without '::'") + if description is not None and not isinstance(description, str): + raise ValueError("Tool namespace description must be a string") + return namespace, description + + +def normalize_tool_namespace(tool: Any) -> Any: + """Qualify explicit namespaces without mutating a definition or tool call. + + Accept a string or {name, description} namespace on the tool or function. + Consistent duplicate declarations are allowed; conflicting identities are + rejected. Remove namespace metadata after folding it into the function name + and description, so normalization is idempotent across API and tokenizer + boundaries. Inputs without namespace metadata retain their existing names. + """ + if not isinstance(tool, dict) or not isinstance(tool.get("function"), dict): + return tool + function = tool["function"] + declarations = [ + _namespace_parts(value) + for value in (tool.get("namespace"), function.get("namespace")) + if value is not None + ] + if not declarations: + return tool + namespace, description = declarations[0] + if any(name != namespace for name, _ in declarations[1:]): + raise ValueError("Conflicting tool namespaces on the tool and function") + + name = function.get("name") + if not isinstance(name, str) or not name: + raise ValueError("Namespaced function name must be a nonempty string") + if "::" in name: + prefix, name = name.split("::", 1) + if prefix != namespace: + raise ValueError("Conflicting tool namespaces in the qualified name") + if not name or "::" in name: + raise ValueError("Namespaced function name must be nonempty without '::'") + + normalized = dict(tool) + normalized.pop("namespace", None) + normalized["function"] = function = dict(function) + function.pop("namespace", None) + function["name"] = f"{namespace}::{name}" + if description: + detail = function.get("description") + if detail is not None and not isinstance(detail, str): + raise ValueError("Function description must be a string") + function["description"] = description + "\n" + (detail or "") + return normalized From 76fa9c0e89c5087a3d64bfc02686fc4cdf7a3ff6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 10:55:43 +0000 Subject: [PATCH 072/365] fix(tools): preserve validation paths and namespace descriptions Normalize supported tool-call iterables without preempting scalar field validation. Retain the first non-null namespace description while rejecting conflicting names. Focused regressions reproduce both review findings; 50 tool-call serialization tests pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit fb50f44a5f025c9c2eed857fbac938e8513227a7) --- .../openai/test_tool_calls_serialization.py | 37 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 5 ++- vllm/utils/tool_names.py | 7 +++- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 478379f5e13b..156d6bb140c4 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -226,6 +226,43 @@ def test_namespace_description_is_preserved_exactly_once(): ).model_dump() == (parsed.model_dump()) +@pytest.mark.parametrize("tool_calls", [1, False, 1.5]) +def test_invalid_tool_calls_shape_reports_the_field(tool_calls): + with pytest.raises(ValidationError) as error: + _make_request( + [{"role": "assistant", "content": None, "tool_calls": tool_calls}] + ) + assert any("tool_calls" in item["loc"] for item in error.value.errors()) + + +@pytest.mark.parametrize( + "outer, inner, expected", + [ + (None, "Stock operations.", "Stock operations.\nFind a SKU."), + ("Outer description.", "Inner description.", "Outer description.\nFind a SKU."), + ("", "Inner description.", "Find a SKU."), + ], +) +def test_consistent_namespaces_preserve_first_non_null_description( + outer, inner, expected +): + tool = { + "type": "function", + "namespace": {"name": "inventory", "description": outer}, + "function": { + "name": "lookup", + "namespace": {"name": "inventory", "description": inner}, + "description": "Find a SKU.", + }, + } + original = copy.deepcopy(tool) + parsed = ChatCompletionToolsParam.model_validate(tool) + assert parsed.function.name == "inventory::lookup" + assert parsed.function.description == expected + assert tool == original + assert ChatCompletionToolsParam.model_validate(parsed.model_dump()) == parsed + + @pytest.mark.parametrize("tool_choice", ["none", "auto", "required"]) @pytest.mark.parametrize( "namespace, name", diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index f9d07659cdb8..e5913efef6c6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -4,6 +4,7 @@ # Adapted from # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py import time +from collections.abc import Iterable from typing import Annotated, Any, ClassVar, Literal from openai.types.chat.chat_completion_audio import ( @@ -564,7 +565,9 @@ def _normalize_messages_before(cls, data: Any) -> Any: if not isinstance(msg, dict): continue tool_calls = msg.get("tool_calls") - if tool_calls is not None: + if isinstance(tool_calls, Iterable) and not isinstance( + tool_calls, (str, bytes, dict) + ): msg["tool_calls"] = [ normalize_tool_namespace(call) for call in tool_calls ] diff --git a/vllm/utils/tool_names.py b/vllm/utils/tool_names.py index 02fcbc9e7c2e..1cf3f368d464 100644 --- a/vllm/utils/tool_names.py +++ b/vllm/utils/tool_names.py @@ -41,8 +41,11 @@ def normalize_tool_namespace(tool: Any) -> Any: if not declarations: return tool namespace, description = declarations[0] - if any(name != namespace for name, _ in declarations[1:]): - raise ValueError("Conflicting tool namespaces on the tool and function") + for declared_name, declared_description in declarations[1:]: + if declared_name != namespace: + raise ValueError("Conflicting tool namespaces on the tool and function") + if description is None: + description = declared_description name = function.get("name") if not isinstance(name, str) or not name: From 20cfba83b08ba757c9434f6cfba9e733db290359 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 11:13:52 +0000 Subject: [PATCH 073/365] perf(comm): index declared PCIe collective plans by their full call contract Bound lookup storage to declared shapes and avoid scans and exception construction when unsupported prefill shapes fall through to another transport. Preserve normalization weight identity, epsilon, strides and first-declaration precedence. Inspired by fernandaspets/vllm_sm120#2. Native suite: 36 passed including TP2 fused collective replay, one four-GPU case skipped. Serving throughput is not yet qualified. Co-authored-by: neolithic5452 Co-authored-by: Neo Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 5e693878601f1145e41582999439080687cc0bf1) --- benchmarks/benchmark_b12x_plan_lookup.py | 81 ++++++++++++++++ .../distributed/test_b12x_pcie_all_reduce.py | 86 +++++++++++++++++ .../b12x_pcie_all_reduce.py | 93 ++++++++++++------- 3 files changed, 229 insertions(+), 31 deletions(-) create mode 100644 benchmarks/benchmark_b12x_plan_lookup.py diff --git a/benchmarks/benchmark_b12x_plan_lookup.py b/benchmarks/benchmark_b12x_plan_lookup.py new file mode 100644 index 000000000000..fc5c8fe0dcba --- /dev/null +++ b/benchmarks/benchmark_b12x_plan_lookup.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure CPU-only lookup cost for declared B12X PCIe collective shapes. + +No GPU operation is executed. Use --module-file to compare pinned source files +in the same Python/runtime environment; these timings are not serving throughput. +""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +import sys +import timeit +from pathlib import Path + +import torch + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--module-file", type=Path, required=True) + parser.add_argument("--declarations", type=int, nargs="+", default=[64, 512, 4096]) + parser.add_argument("--iterations", type=int, default=1000) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + if min(*args.declarations, args.iterations, args.repeats) < 1: + parser.error("counts must be positive") + spec = importlib.util.spec_from_file_location("pcie_lookup_probe", args.module_file) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + report = { + "source_sha256": hashlib.sha256(args.module_file.read_bytes()).hexdigest(), + "iterations": args.iterations, + "repeats": args.repeats, + "scope": "CPU metadata lookup; no collective or model execution", + "results": [], + } + for count in args.declarations: + communicator = object.__new__(module.B12xPcieAllReduce) + invocations = [ + module.B12xPcieInvocation( + name=str(rows), + operation="all_reduce", + shape=(rows, 16), + dtype=torch.bfloat16, + ) + for rows in range(1, count + 1) + ] + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + if hasattr(communicator, "_index_declared_plans"): + communicator._index_declared_plans() + for case, rows, expected in (("hit", count, True), ("miss", count + 1, False)): + source = torch.empty((rows, 16), dtype=torch.bfloat16) + + def lookup(communicator=communicator, source=source): + return communicator._has_plan_for(source) + + assert lookup() is expected + for _ in range(100): + lookup() + samples = timeit.repeat(lookup, number=args.iterations, repeat=args.repeats) + report["results"].append( + { + "declarations": count, + "case": case, + "microseconds_per_call": [ + s * 1e6 / args.iterations for s in samples + ], + "median_us": statistics.median(samples) * 1e6 / args.iterations, + } + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_b12x_pcie_all_reduce.py b/tests/distributed/test_b12x_pcie_all_reduce.py index 4a4b6a21eeda..017883b4a3e5 100644 --- a/tests/distributed/test_b12x_pcie_all_reduce.py +++ b/tests/distributed/test_b12x_pcie_all_reduce.py @@ -58,9 +58,11 @@ def _make_communicator( communicator.twoshot_max_bytes = 0 plan = object() communicator._plans = {"prepared": plan} + communicator._plan_index = {} communicator._routes = {} communicator._invocations = {} communicator._plan_for = MagicMock(return_value=plan) + communicator._lookup_plan = MagicMock(return_value=plan) return communicator, runtime @@ -323,6 +325,7 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: communicator, _ = _make_communicator() communicator._plans = {} communicator._invocations = {} + communicator._lookup_plan.return_value = None with pytest.raises( PreparationResourceUnavailableError, @@ -331,6 +334,89 @@ def test_dispatch_rejects_missing_exact_preparation() -> None: B12xPcieAllReduce._plan_for(communicator, torch.randn(2, 4)) +def _indexed_communicator(invocations): + communicator = object.__new__(B12xPcieAllReduce) + communicator._invocations = {item.name: item for item in invocations} + communicator._plans = {item.name: object() for item in invocations} + communicator._index_declared_plans() + return communicator + + +def test_declared_plan_index_preserves_complete_fused_identity(): + operation = "all_reduce_fused_add_rms_norm" + weight_a = torch.ones(4) + weight_b = weight_a.clone() + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocations = [ + b12x_pcie_all_reduce.B12xPcieInvocation( + name=name, + operation=operation, + shape=(2, 4), + dtype=source.dtype, + norm_weight=weight, + epsilon=epsilon, + ) + for name, weight, epsilon in ( + ("a", weight_a, 1e-6), + ("b", weight_b, 1e-6), + ("epsilon", weight_a, 1e-5), + ) + ] + communicator = _indexed_communicator(invocations) + for invocation in invocations: + assert ( + communicator._plan_for( + source, + operation=operation, + weight=invocation.norm_weight, + epsilon=invocation.epsilon, + ) + is communicator._plans[invocation.name] + ) + assert not communicator._has_plan_for(source) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a.clone(), epsilon=1e-6 + ) + assert not communicator._has_plan_for( + source, operation=operation, weight=weight_a, epsilon=1e-4 + ) + + +def test_declared_plan_index_matches_shape_dtype_stride_and_first_declaration(): + source = torch.empty((2, 4), dtype=torch.bfloat16) + invocation = b12x_pcie_all_reduce.B12xPcieInvocation + communicator = _indexed_communicator( + [ + invocation("first", "all_reduce", (2, 4), source.dtype), + invocation( + "equivalent", "all_reduce", (2, 4), source.dtype, strides=(4, 1) + ), + invocation("strided", "all_reduce", (2, 4), source.dtype, strides=(8, 1)), + ] + ) + assert communicator._plan_for(source) is communicator._plans["first"] + strided = torch.empty((2, 8), dtype=source.dtype)[:, :4] + assert communicator._plan_for(strided) is communicator._plans["strided"] + assert not communicator._has_plan_for(source.float()) + assert not communicator._has_plan_for(source.flatten()) + communicator._invocations = {} + assert communicator._plan_for(source) is communicator._plans["first"] + communicator._index_declared_plans() + assert not communicator._has_plan_for(source) + + +def test_undeclared_plan_probe_neither_raises_nor_grows_the_index(): + communicator = _indexed_communicator([]) + for rows in range(1, 65): + source = torch.empty((rows, 4), dtype=torch.bfloat16) + assert not communicator._has_plan_for(source) + with pytest.raises( + PreparationResourceUnavailableError, match="no declared plan" + ): + communicator._plan_for(source) + assert communicator._plan_index == {} + + def test_fused_allreduce_has_an_independent_cutoff() -> None: communicator, runtime = _make_communicator( allreduce_max_bytes=16, fused_max_bytes=64 diff --git a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py index 16cf5b31fadf..e82d1510731f 100644 --- a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py @@ -248,6 +248,7 @@ def __init__( ] ] = [] self._plans: dict[str, object] = {} + self._plan_index: dict[tuple, object] = {} self._invocations: dict[str, B12xPcieInvocation] = {} self._routes: dict[str, str] = {} if len(self.global_ranks) != self.world_size: @@ -596,12 +597,14 @@ def get_b12x_preparation_units( plans: dict[str, object] = {} routes: dict[str, str] = {} collective_ranks = tuple(sorted(self.global_ranks)) - declarations = {} + declarations: dict[tuple, Any] = {} + prepare: Callable[..., Any] for invocation in invocations: route = self._route_invocation(invocation) if route is None: continue routes[invocation.name] = route + assert self._runtime is not None target = self._runtime._prepared_channel_for_stream( None, invocation.channel_id ) @@ -630,7 +633,7 @@ def get_b12x_preparation_units( plan = _oneshot_preparation.plan(query, runtime=target) declarations[key] = plan - def prepare(state, invocation=invocation, query=query): + def prepare_oneshot(state, invocation=invocation, query=query): registered = bool(query.setup["registered"]) inp = self._request_input(invocation, registered=registered) reset, produce, restore, owners = self._oneshot_prime_input( @@ -665,6 +668,8 @@ def produce_fused() -> None: restore=restore, owners=owners, ) + + prepare = prepare_oneshot elif route == "twoshot": assert self._twoshot is not None query = _twoshot_preparation.query_from_metadata( @@ -681,13 +686,15 @@ def produce_fused() -> None: plan = _twoshot_preparation.plan(query, runtime=self._twoshot) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_twoshot(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) return _twoshot_preparation.prepared_call( state, payload=inp, out=torch.empty_like(inp) ) + + prepare = prepare_twoshot else: assert self._dma is not None query = _dma_preparation.query_from_metadata( @@ -703,7 +710,7 @@ def prepare(state, invocation=invocation): plan = _dma_preparation.plan(query, runtime=self._dma) declarations[key] = plan - def prepare(state, invocation=invocation): + def prepare_dma(state, invocation=invocation): inp = self._request_input(invocation, registered=False) assert inp is not None inp.fill_(1) @@ -711,6 +718,7 @@ def prepare(state, invocation=invocation): state, inp=inp, out=torch.empty_like(inp) ) + prepare = prepare_dma plans[invocation.name] = plan requests.append( plan.request( @@ -732,6 +740,7 @@ def prepare(state, invocation=invocation): } self._plans = plans self._routes = routes + self._index_declared_plans() if not requests: return () return ( @@ -743,6 +752,47 @@ def prepare(state, invocation=invocation): ), ) + @staticmethod + def _plan_key(operation, shape, dtype, strides, weight=None, epsilon=None): + norm = None if operation == "all_reduce" else (id(weight), epsilon) + return operation, tuple(shape), dtype, tuple(strides), norm + + def _index_declared_plans(self) -> None: + index: dict[tuple, object] = {} + for name, invocation in self._invocations.items(): + key = self._plan_key( + invocation.operation, + invocation.shape, + invocation.dtype, + invocation.strides + if invocation.strides is not None + else self._contiguous_strides(invocation.shape), + invocation.norm_weight, + invocation.epsilon, + ) + # Equivalent declarations keep the first plan, matching model order. + index.setdefault(key, self._plans[name]) + self._plan_index = index + + def _lookup_plan( + self, + inp: torch.Tensor, + *, + operation: str = "all_reduce", + weight: torch.Tensor | None = None, + epsilon: float | None = None, + ): + return self._plan_index.get( + self._plan_key( + operation, + inp.shape, + inp.dtype, + inp.stride(), + weight, + epsilon, + ) + ) + def _plan_for( self, inp: torch.Tensor, @@ -751,32 +801,16 @@ def _plan_for( weight: torch.Tensor | None = None, epsilon: float | None = None, ): - strides = tuple(inp.stride()) - for name, invocation in self._invocations.items(): - if ( - invocation.operation == operation - and tuple(inp.shape) == invocation.shape - and inp.dtype == invocation.dtype - and strides - == ( - self._contiguous_strides(invocation.shape) - if invocation.strides is None - else invocation.strides - ) - and ( - operation == "all_reduce" - or ( - invocation.norm_weight is weight - and invocation.epsilon == epsilon - ) - ) - ): - return self._plans[name] + plan = self._lookup_plan( + inp, operation=operation, weight=weight, epsilon=epsilon + ) + if plan is not None: + return plan from vllm.utils.b12x import PreparationResourceUnavailableError raise PreparationResourceUnavailableError( f"PCIe collective has no declared plan for {operation} on shape " - f"{tuple(inp.shape)} {inp.dtype} strides {strides}; declared: " + f"{tuple(inp.shape)} {inp.dtype} strides {tuple(inp.stride())}; declared: " + ", ".join( f"{invocation.operation}{invocation.shape}" for invocation in self._invocations.values() @@ -791,11 +825,7 @@ def _has_plan_for(self, inp: torch.Tensor, **kwargs) -> bool: the captured and planned serving shapes, so an undeclared shape is a prefill chunk or warm-up batch outside the graph-replayed set. """ - from vllm.utils.b12x import PreparationResourceUnavailableError - - try: - self._plan_for(inp, **kwargs) - except PreparationResourceUnavailableError: + if self._lookup_plan(inp, **kwargs) is None: logger.debug( "b12x PCIe all-reduce declines undeclared shape %s %s", tuple(inp.shape), @@ -953,6 +983,7 @@ def close(self) -> None: self._runtime.close() self._runtime = None self._plans.clear() + self._plan_index.clear() self._invocations.clear() self._routes.clear() self.disabled = True From c362325c2071f495cdd42541dc476cf9a901ea9f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:36:27 +0000 Subject: [PATCH 074/365] fix(glm): invalidate physical-selection plans with cache bindings Preserve the KK prepared-state call interface and invalidate geometry-dependent declarations on cache bind/unbind. Extend the existing construction test to cover declaration lifetime. Reconciles the remaining ownership part of vLLM #789; native qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 8 ++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 5 +++++ 2 files changed, 13 insertions(+) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 843fb7e1356d..ffb61df5ec9e 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -751,9 +751,17 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( _, main = _packed_main_cache( device=torch.device("cpu"), blocks=2, layers=1, block_size=256, layer=0 ) + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() indexer.bind_main_kv_cache(main) assert indexer.indexer_op._num_q_heads == indexer._q_fp8.shape[1] == 32 + assert indexer._physical_selection_plan is None + indexer._physical_selection_plan = object() + indexer.unbind_main_kv_cache() + assert indexer._physical_selection_plan is None + assert indexer._index_cache is None + assert indexer._main_cache_num_blocks == 0 def test_glm53_packed_tail_accepts_nvfp4_main_record() -> None: diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index cc0311138364..3f15e9b2ce77 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -143,6 +143,7 @@ def __init__( self._expand_pooled_topk_to_physical_slots = ( b12x_sparse_mla.expand_pooled_topk_to_physical_slots ) + self._physical_selection_plan = None self.index_kpool_compress_ape = nn.Parameter( allocate_weights( torch.empty, @@ -360,6 +361,7 @@ def _index_cache_view( return index_cache, subpages_per_parent, parent_stride_pages def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: + self._physical_selection_plan = None index_cache, subpages, parent_stride_pages = self._index_cache_view(main_cache) block_size = int(main_cache.shape[1]) parent_table_width = self._max_parent_table_width( @@ -393,6 +395,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: def unbind_main_kv_cache(self) -> None: self._index_cache = None self._main_cache_num_blocks = 0 + self._physical_selection_plan = None def make_b12x_physical_selection_prepare_call( self, output: torch.Tensor, active_counts: torch.Tensor, *, state=None @@ -707,6 +710,8 @@ def forward( if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: if self._main_cache_num_blocks < 1: raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From 0f1d80f4f2e02503f303515f4d6c0a71c0a24108 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 20:49:49 +0000 Subject: [PATCH 075/365] test(b12x): prepare native cache writers in graph-lifetime coverage Keep Karmic cache preparation contracts while checking that graph resources release caller activations. Preserve parser stream chunks as explicit concatenated values. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/parser/engine/test_deepseek_v4.py | 12 ++++++--- tests/v1/attention/test_b12x_v41_workspace.py | 25 ++++++++++--------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 9180a6f4b2d5..5eb369cee969 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -305,8 +305,10 @@ def test_streaming_split_next_parameter_tag_is_buffered( chunks = [ DSML_TOOL_START, f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n", - f"<{_PARAM_OPEN.format(name='location', is_str='true')}" - "Paris a<｜DSML｜parameter", + ( + f"<{_PARAM_OPEN.format(name='location', is_str='true')}" + "Paris a<｜DSML｜parameter" + ), ' name="date" string="true">tomorrow', _PARAM_CLOSE, DSML_INVOKE_END, @@ -376,8 +378,10 @@ def test_streaming_with_trailing_content(self, mock_tokenizer, mock_request): parser = DeepSeekV4Parser(mock_tokenizer) chunks = [ DSML_TOOL_START, - f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" - f"{_param('location', 'true', 'NYC')}\n", + ( + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" + f"{_param('location', 'true', 'NYC')}\n" + ), DSML_TOOL_END, "Done.", ] diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 260f386aa122..847536b72ae4 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -787,12 +787,13 @@ def metadata(page): device=device, dtype=torch.uint8, ) - from b12x.attention._shared.mla.compressed_reference import ( - pack_deepseek_v41_cache_reference, - ) - - cache[1:].copy_( - pack_deepseek_v41_cache_reference(kv, page_size=page, cache_kind=kind) + _write_cache( + kv, + cache, + torch.arange(length, device=device) + page, + page_size=page, + cache_kind=kind, + cache_format="deepseek_v41", ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache @@ -813,6 +814,12 @@ def metadata(page): index_keys = torch.zeros((length, 128), dtype=torch.bfloat16, device=device) index_keys[:512, 0] = 1 index_keys[256:768, 1] = 1 + _write_index_keys( + index_keys, + index_k_cache=layer.indexer.k_cache.kv_cache, + slot_mapping=torch.arange(length, device=device) + main_page, + page_size=main_page, + ) layer.attn_sink = torch.zeros(layer.n_local_heads, device=device) q = torch.randn( (rows, layer.n_local_heads, 512), dtype=torch.bfloat16, device=device @@ -836,12 +843,6 @@ def metadata(page): units = layer.get_b12x_preparation_units(layer, workload) requests = tuple(request for unit in units for request in unit.requests) session.prepare(requests, autotune=False) - attention.dsa_indexer.quantize_write_index_k_mxfp4( - layer._index_plan("prefill", min(rows, layer.INDEX_CHUNK)), - index_keys, - index_k_cache=layer.indexer.k_cache.kv_cache, - slot_mapping=torch.arange(length, device=device) + main_page, - ) activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): From 52b8258869ac1222e9426059ff8559e8cc74d89c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:23:02 +0000 Subject: [PATCH 076/365] test(ds41): declare attention bounds in WO capture fixtures Prepare the output-projection fixture with the same declared attention envelope as the model. Both graph-lifetime cases pass on SM120 with the KK native wheel. Co-authored-by: OpenAI Codex --- tests/v1/attention/test_b12x_v41_workspace.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 847536b72ae4..c0c33fa5a341 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -1047,6 +1047,7 @@ def test_output_projection_capture_releases_caller_activations( ), ) layer.setup_wo_projection() + layer._declare_attention(device) workload = B12xWorkload( stage="weights", token_counts=(8, 16), From c5e26550571d222d24638879ee2e016cb5ab7163 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 21:37:43 +0000 Subject: [PATCH 077/365] test(b12x): preserve Karmic preparation and precision contracts Observe caller-owned MHC outputs while requiring preparation not to retain them; retain allocation, pointer and graph replay assertions. Exercise the generic ModelOpt W4A16 method and padded native scales. Initialize the complete NVFP4 fixture. Use an IEEE FP32 reference instead of TF32 for unquantized projections; existing numerical tolerances are unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../kernels/test_b12x_linear.py | 52 ++++++++++++++----- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 6a1a540400e3..5a2e6d0a7db7 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -493,7 +493,8 @@ def pack(weight: torch.Tensor, weight_scale: torch.Tensor): assert len(calls) == 1 weight, weight_scale = calls[0] assert weight.shape == (48, 128) - assert weight_scale.shape == (48, 4) + # Preserve the checkpoint's padded row layout for native packed scales. + assert weight_scale.shape == (64, 4) assert weight.dtype == torch.float8_e4m3fn assert weight_scale.dtype == torch.uint8 assert layer.weight.numel() == 0 @@ -933,7 +934,14 @@ def test_b12x_fp4_processes_scale_and_preserves_loader( if kernel_cls is B12xNvFp4LinearKernel: layer.weight = torch.empty((48, 64), dtype=torch.uint8) layer.weight_global_scale = torch.tensor(0.5) - packed = object() + layer.input_global_scale_inv = torch.tensor(1.0) + + @dataclass(frozen=True) + class PackedWeight: + in_features: int + values: torch.Tensor + + packed = PackedWeight(0, layer.weight) def pack_weight(weight, scale, *, recipe, global_scale): assert weight.data_ptr() == layer.weight.data_ptr() @@ -948,13 +956,16 @@ def pack_weight(weight, scale, *, recipe, global_scale): lambda: types.SimpleNamespace(pack_weight=pack_weight), ) kernel = object.__new__(kernel_cls) + if kernel_cls is B12xNvFp4LinearKernel: + kernel.config = NvFp4LinearLayerConfig(use_a16=False) kernel.process_weights_after_loading(layer) assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() assert layer.weight_scale.weight_loader is weight_loader if kernel_cls is B12xNvFp4LinearKernel: - assert layer.b12x_nvfp4_packed_weight is packed + assert layer.b12x_nvfp4_packed_weight.in_features == 128 + assert layer.b12x_nvfp4_packed_weight.values is packed.values def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: @@ -1055,7 +1066,7 @@ def test_b12x_w4a16_modelopt_vision_width_preserves_bf16_activations(monkeypatch from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, - ModelOptNvFp4W4A16LinearMethod, + build_linear_method, ) from vllm.v1.worker.workspace import ( current_workspace_manager, @@ -1077,11 +1088,13 @@ def reject_activation_quantization(*args, **kwargs): n, k = 1152, 4304 config = VllmConfig(kernel_config=KernelConfig(linear_backend="b12x")) with set_current_vllm_config(config), torch.no_grad(): - method = ModelOptNvFp4W4A16LinearMethod( + method = build_linear_method( ModelOptNvFp4Config( quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True, - ) + ), + "W4A16_NVFP4", + "vision_tower.mlp", ) layer = torch.nn.Module() with torch.device(device): @@ -1932,10 +1945,13 @@ def test_b12x_mhc_preparation_reuses_outputs_with_exact_graph_replay( name = f"_b12x_mhc_{native_operation}_impl" native = getattr(_impl, name) captured_call: list[tuple[tuple[torch.Tensor, ...], dict[str, Any]]] = [] + captured_outputs = [] def observe(*args, **kwargs): captured_call[:] = [(args, kwargs)] - return native(*args, **kwargs) + result = native(*args, **kwargs) + captured_outputs[:] = outputs(result) + return result def outputs(value): return value if isinstance(value, tuple) else (value,) @@ -1948,7 +1964,10 @@ def outputs(value): with monkeypatch.context() as hooks: hooks.setattr(_impl, name, observe) call.produce() - first = outputs(call.invoke()) + assert call.invoke() is None + first = tuple(captured_outputs) + assert call.invoke() is None + second = tuple(captured_outputs) args, kwargs = captured_call.pop() reference_kwargs = { key: value @@ -1956,7 +1975,6 @@ def outputs(value): if key not in {"residual_out", "y_out", "post_out", "comb_out", "out"} } pointers = [tensor.data_ptr() for tensor in first] - second = outputs(call.invoke()) assert [tensor.data_ptr() for tensor in second] == pointers for actual, expected in zip( second, outputs(native(*args, **reference_kwargs)), strict=True @@ -1975,9 +1993,9 @@ def outputs(value): graph.replay() torch.accelerator.synchronize() assert torch.accelerator.memory_allocated(device) == allocated - assert [ - tensor.data_ptr() for tensor in outputs(call.output) - ] == pointers + # Preparation must not pin full-prefill results in the session. + assert call.output is None + assert [tensor.data_ptr() for tensor in first] == pointers expected = outputs(native(*args, **reference_kwargs)) for actual, wanted in zip(first, expected, strict=True): assert torch.isfinite(actual).all() @@ -2722,7 +2740,7 @@ def expected(rows): assert set(plans) == {4, capacity} actual = run(11) assert plans[11].prepared is not None - assert plans[11].selection.source == "fixed" + assert plans[11].selection.source == "default" torch.testing.assert_close(actual, expected(11), rtol=0.02, atol=0.125) session.freeze() with kernel_resolution_guard("FP8 prepared exact-M execution"): @@ -2746,11 +2764,17 @@ def expected(rows): @pytest.mark.parametrize("output_dtype", [torch.bfloat16, torch.float32]) @torch.no_grad() -def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay(output_dtype): +def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay( + output_dtype, request +): if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: pytest.skip("native b12x projection requires SM12x") from vllm.models.deepseek_v4_1.b12x_layers import B12xLinearMethod + # A TF32 reference rounds FP32 inputs, unlike the native SIMT projection. + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = torch.nn.Module() layer.weight = torch.nn.Parameter( From 149d38334b7834a44078c90dd3e9ed5c139a7b9e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 22:04:32 +0000 Subject: [PATCH 078/365] test(ds41): model native FP8 prefill arithmetic in CED oracle Compare compact replay with an independent BF16 QK / FP8 PV reference and assert the complete causal candidate set. Retain numerical tolerances, poisoned discarded rows, padded outputs, caller scratch and CUDA graph checks. Production kernels and precision are unchanged. Validation: 333 native tests passed, three multi-GPU tests separately passed; all three CED geometries pass at unchanged tolerances. Reference adapted from B12X commit 4af78b86771d8491ec7b9d27168fdfbc1d67a773. Existing mypy and CUDA-API hook findings are unchanged. Co-authored-by: Luke Alonso Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_v41_workspace.py | 129 ++++++++++++++---- 1 file changed, 106 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index c0c33fa5a341..66ea9c7bd3e6 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -701,6 +701,73 @@ def test_mhc_fixed_capacity_buckets_preserve_decode_policy( assert bool(torch.isfinite(y).all()) +def _v41_fp8_operands(cache, kind): + """Decode native cache rows into E4M3 operands and per-64 scales. + + This independent operand reference follows B12X's + tests/_reference/v41_fp8.py. SWA pairs its E8M0 scales; indexed NVFP4 + uses a power-of-two scale that bounds four adjacent scale groups. + """ + if kind == "swa": + rows = cache.reshape(-1, 528) + codes = rows[:, :512].contiguous().view(torch.float8_e4m3fn).double() + original = rows[:, 512:528].double() + exponent = original.reshape(-1, 8, 2).amax(-1).clamp_min(1) - 127 + values = codes * torch.exp2(original - 127).repeat_interleave(32, dim=1) + else: + rows = cache.reshape(-1, 288) + packed = rows[:, :256].long() + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(1) + lut = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], + dtype=torch.float64, + device=cache.device, + ) + original = rows[:, 256:288].contiguous().view(torch.float8_e4m3fn).double() + values = lut[codes] * original.repeat_interleave(16, dim=1) + bound = original.reshape(-1, 8, 4).amax(-1) * (6 / 448) + exponent = torch.ceil(torch.log2(bound.clamp_min(2.0**-126))) + scales = torch.exp2(exponent) + codes = (values / scales.repeat_interleave(64, dim=1)).float() + return codes.to(torch.float8_e4m3fn).float(), scales.float() + + +def _v41_fp8_prefill_reference(q, values, scales, valid, sink): + """Model BF16 Q/K and per-64-key FP8 probability/value products. + + Prefill computes Q/K with BF16 operands. Each output group independently + quantizes scaled probabilities to E4M3 before its product with FP8 V; + partial numerators remain FP32 until the final LSE merge. + """ + keys = (values * scales.repeat_interleave(64, dim=-1)).bfloat16().float() + logits = torch.einsum("rhd,rkd->rhk", q.float(), keys) * 512**-0.5 + logits.masked_fill_(~valid[:, None], -torch.inf) + partials, lses = [], [] + for first in range(0, values.shape[1], 64): + local = logits[:, :, first : first + 64] + maximum = local.amax(-1) + p = torch.exp(local - maximum[:, :, None]) + p = torch.where(valid[:, None, first : first + 64], p, 0) + denominator = p.sum(-1) + groups = [] + for group in range(8): + weighted = p * scales[:, None, first : first + 64, group] + pscale = weighted.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 + pq = (weighted / pscale).to(torch.float8_e4m3fn).float() + vq = values[:, first : first + 64, group * 64 : (group + 1) * 64] + groups.append(torch.einsum("rhk,rkd->rhd", pq, vq) * pscale) + partials.append(torch.cat(groups, -1) / denominator.clamp_min(1e-30)[..., None]) + lses.append( + torch.where(denominator > 0, maximum + denominator.log(), -torch.inf) + ) + lse = torch.stack(lses, -1) + maximum = torch.maximum(lse.amax(-1), sink.float()[None]) + weights = torch.where(torch.isfinite(lse), torch.exp(lse - maximum[..., None]), 0) + denominator = weights.sum(-1) + torch.exp(sink.float()[None] - maximum) + output = (torch.stack(partials, -2) * weights[..., None]).sum(-2) + return (output / denominator.clamp_min(1e-30)[..., None]).bfloat16().float() + + @pytest.mark.parametrize("main_page,swa_page", [(64, 32), (128, 64), (256, 128)]) @pytest.mark.parametrize( "is_decode,rows,live_rows", @@ -1561,14 +1628,14 @@ def test_ced_compact_attention_bounded_oracle_and_frozen_replay( ): from contextlib import ExitStack - from b12x.attention._shared.mla.compressed_reference import ( - unpack_deepseek_v41_cache_reference, - ) from b12x.preparation import PreparationSession from vllm.utils.b12x import B12xWorkload attention, manager, workspace = native_workspace + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + request.addfinalizer(lambda: torch.set_float32_matmul_precision(precision)) torch.manual_seed(713) device = torch.device("cuda", torch.accelerator.current_device_index()) layer = _layer(attention, 20, swa_page=swa_page) @@ -1633,17 +1700,19 @@ def metadata(page): ) if kind == "swa": layer.swa_cache_layer.kv_cache = cache - swa_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + swa_values, swa_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) # NaN FP8 payloads in every old SWA row: reading below the # replay boundary cannot accidentally look like a valid zero page. cache[1 : 1 + boundary // page].fill_(0x7F) else: layer.kv_cache = cache - main_values = unpack_deepseek_v41_cache_reference( - cache, page_size=page, cache_kind=kind - )[page : page + length].clone() + main_values, main_scales = ( + part[page : page + length].clone() + for part in _v41_fp8_operands(cache, kind) + ) layer.indexer.k_cache.kv_cache = torch.zeros( ( length // main_page + 1, @@ -1683,21 +1752,35 @@ def run(): ) def oracle(live): - values = torch.cat((swa_values[boundary:], main_values)) - logical = torch.cat( + selected = layer.topk_indices_buffer[:live].long() + # Candidate order affects per-tile FP8 rounding. Verify its complete + # causal set independently before using that order in the numeric oracle. + columns = torch.arange(selected.shape[1], device=device)[None] + expected_ids = columns.expand_as(selected).clone() + expected_ids.masked_fill_(columns > positions[:live, None], length) + actual_ids = selected.masked_fill(selected < 0, length).sort(dim=1).values + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + swa_logical = torch.arange(boundary, length, device=device)[None] + swa_valid = swa_logical <= positions[:live, None] + main_valid = (selected >= 0) & (selected <= positions[:live, None]) + valid = torch.cat((swa_valid, main_valid), dim=1) + values = torch.cat( ( - torch.arange(boundary, length, device=device), - torch.arange(length, device=device), - ) - ) - logits = torch.einsum("rhd,kd->rhk", q[:live].float(), values) * 512**-0.5 - logits.masked_fill_( - logical[None, None] > positions[:live, None, None], -torch.inf - ) - logits = torch.cat( - (logits, layer.attn_sink[None, :, None].expand(live, -1, -1)), -1 + swa_values[None, boundary:].expand(live, -1, -1), + main_values[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + scales = torch.cat( + ( + swa_scales[None, boundary:].expand(live, -1, -1), + main_scales[selected.clamp_min(0)], + ), + dim=1, + ).masked_fill(~valid[..., None], 0) + return _v41_fp8_prefill_reference( + q[:live], values, scales, valid, layer.attn_sink ) - return torch.einsum("rhk,kd->rhd", logits.softmax(-1)[..., :-1], values) workload = B12xWorkload( stage="state", @@ -1744,9 +1827,9 @@ def oracle(live): visible.copy_((positions + 1).clamp_min(0).int()) starts[1] = live q.normal_() - expected = oracle(live) # Eager execution also catches live-count specialization leaks. run() + expected = oracle(live) torch.testing.assert_close( out[:live].float(), expected, rtol=0.04, atol=0.025 ) From b69517f63c8d49bcc8e8a8bf92c82eee2c67389c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:11:18 +0000 Subject: [PATCH 079/365] fix(dflash): initialize fused MXFP8 projection with generic ModelOpt Recognize MXFP8 through its QuantSpec weight key and give the fused context K/V projection an independent kernel lifecycle. Preserve input/output dtypes and bias metadata without changing checkpoint precision or draft sampling. Keep BF16 loading unchanged and adapt the draft-loader fixture to the shared pipeline-safe loader helper. Validation: 34 DFlash tests pass. Native B12X fused projection passes exact representable-input checks at 1/8/32 rows for quantized and A16 activations, including changed-input CUDA graph replay with zero additional allocated bytes. Applicable pre-commit hooks pass. Whole-model MXFP8 DFlash qualification remains pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/spec_decode/test_dflash2.py | 100 +++++++++++++++++++-- vllm/model_executor/models/qwen3_dflash.py | 45 ++++++---- 2 files changed, 120 insertions(+), 25 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 315f28235456..09dccc5054e9 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -14,6 +14,8 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): + from vllm.config import LoadConfig + draft_load_config = object() draft_model_config = SimpleNamespace(hf_config=SimpleNamespace()) speculative_config = SimpleNamespace( @@ -25,6 +27,7 @@ def test_dflash_loader_honors_draft_load_config(monkeypatch): vllm_config = SimpleNamespace( attention_config=SimpleNamespace(), cache_config=SimpleNamespace(), + load_config=LoadConfig(load_format="fastsafetensors"), speculative_config=speculative_config, ) loaded = SimpleNamespace(model=SimpleNamespace()) @@ -42,17 +45,16 @@ def fake_get_model(**kwargs): monkeypatch.setattr(dflash_utils, "replace", fake_replace) monkeypatch.setattr(dflash_utils, "get_model", fake_get_model) monkeypatch.setattr( - dflash_utils, - "get_pp_group", + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.compilation.backends.set_model_tag", - lambda _tag: nullcontext(), + "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group", + lambda: SimpleNamespace(world_size=2), ) monkeypatch.setattr( - "vllm.model_executor.models.qwen3_dflash.dflash_target_rope_is_neox_style", - lambda _model: None, + "vllm.compilation.backends.set_model_tag", + lambda _tag: nullcontext(), ) monkeypatch.setattr( "vllm.model_executor.models.qwen3_dflash.dflash_has_any_non_causal", @@ -61,6 +63,8 @@ def fake_get_model(**kwargs): assert dflash_utils.load_dflash_model(SimpleNamespace(), vllm_config) is loaded assert captured["load_config"] is draft_load_config + assert captured["vllm_config"].load_config.load_format == "auto" + assert vllm_config.load_config.load_format == "fastsafetensors" def test_dflash_reset_attn_releases_cache_layout_state(): @@ -158,7 +162,7 @@ def test_dflash_context_projection_rejects_mixed_quantization(mxfp8_layer: int): from torch import nn from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, + build_linear_method, ) from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model @@ -166,8 +170,8 @@ class UnreadableWeight: def __getitem__(self, _key): raise AssertionError("weights must not be read before validation") - mxfp8_method = object.__new__(ModelOptMxFp8LinearMethod) - methods = [None, None] + mxfp8_method = build_linear_method(None, "MXFP8", "") + methods = [build_linear_method(None, "FP8", ""), None] methods[mxfp8_layer] = mxfp8_method layers_attn = [ SimpleNamespace( @@ -186,6 +190,84 @@ def __getitem__(self, _key): model._build_context_kv_buffers(layers_attn, has_bias=False) +@pytest.mark.parametrize("mxfp8", [False, True]) +@pytest.mark.parametrize("has_bias", [False, True]) +def test_fused_context_projection_owns_its_linear_method(monkeypatch, mxfp8, has_bias): + """Fused K/V packing must not reconfigure the query projection's kernel.""" + from torch import nn + + import vllm.model_executor.layers.quantization.modelopt as modelopt + import vllm.model_executor.parameter as parameter + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + kernels = [] + + def select_kernel(spec, layer, runtime_dtypes, **kwargs): + kernel = SimpleNamespace( + process_weights_after_loading=lambda layer: None, + input_quant_key=lambda: None, + ) + kernels.append((kernel, layer.output_size_per_partition)) + assert runtime_dtypes.input_dtype == torch.bfloat16 + assert runtime_dtypes.out_dtype == torch.bfloat16 + assert runtime_dtypes.marlin_input_dtype == torch.float16 + assert layer.has_bias is has_bias + return kernel + + monkeypatch.setattr(modelopt, "select_linear_kernel", select_kernel) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + method = modelopt.build_linear_method(None, "MXFP8", "") if mxfp8 else None + original_kernel = object() + if method is not None: + method.kernel = original_kernel + method.input_dtype = method.out_dtype = torch.bfloat16 + method.marlin_input_dtype = torch.float16 + layers = [] + for index in range(2): + projection = nn.Module() + projection.quant_method = method + weight = torch.full((192, 64), index + 1, dtype=torch.bfloat16) + if mxfp8: + weight = weight.to(torch.float8_e4m3fn) + projection.weight_scale = nn.Parameter( + torch.full((192, 2), 127, dtype=torch.uint8), requires_grad=False + ) + projection.weight = nn.Parameter(weight, requires_grad=False) + if has_bias: + projection.bias = nn.Parameter(torch.full((192,), float(index + 1))) + attention = SimpleNamespace( + qkv_proj=projection, + q_size=128, + k_norm=SimpleNamespace(weight=torch.ones(32)), + ) + layers.append(SimpleNamespace(self_attn=attention)) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.layers = layers + model.hidden_norm = SimpleNamespace(weight=torch.ones(64, dtype=torch.bfloat16)) + model._fused_kv_linear = nn.Module() + model._fused_kv_quant_method = None + model._build_context_kv_buffers([layer.self_attn for layer in layers], has_bias) + expected = model._fused_kv_weight.clone() + model.process_weights_after_loading() + if mxfp8: + assert method is not None + assert model._fused_kv_quant_method is not method + assert method.kernel is original_kernel + assert kernels == [(model._fused_kv_quant_method.kernel, 128)] + assert model._fused_kv_linear.weight_block_size == [1, 32] + torch.testing.assert_close( + model._fused_kv_linear.weight.float(), expected.float() + ) + assert model._fused_kv_weight is None + assert model._fused_kv_weight_scale is None + else: + assert not kernels + assert model._fused_kv_quant_method is None + torch.testing.assert_close(model._fused_kv_weight, expected) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3c7c301814ff..0ea09746a4bb 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -27,6 +27,8 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -511,13 +513,11 @@ def _build_context_kv_buffers( layers_attn: list[nn.Module], has_bias: bool, ) -> None: - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_methods = [a.qkv_proj.quant_method for a in layers_attn] uses_mxfp8 = [ - isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods + isinstance(method, ModelOptLinearMethod) + and method.spec.weight == kMxfp8Static + for method in quant_methods ] if any(uses_mxfp8) and not all(uses_mxfp8): raise ValueError( @@ -593,12 +593,11 @@ def _build_fused_kv_buffers(self) -> None: def process_weights_after_loading(self) -> None: """Pack the serialized MXFP8 context projection for its GEMM backend.""" - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMxFp8LinearMethod, - ) - quant_method = self.layers[0].self_attn.qkv_proj.quant_method - if not isinstance(quant_method, ModelOptMxFp8LinearMethod): + if not ( + isinstance(quant_method, ModelOptLinearMethod) + and quant_method.spec.weight == kMxfp8Static + ): return if self._fused_kv_weight is None or self._fused_kv_weight_scale is None: raise RuntimeError( @@ -607,9 +606,23 @@ def process_weights_after_loading(self) -> None: ) output_size, input_size = self._fused_kv_weight.shape - self._fused_kv_linear.input_size_per_partition = input_size - self._fused_kv_linear.output_size_per_partition = output_size - self._fused_kv_linear.logical_widths = [output_size] + # Fused context K/V has an independent weight shape and kernel + # lifecycle; do not repack through the query projection's method. + fused_method = ModelOptLinearMethod( + quant_method.spec, quant_method.ctx, quant_method.fmt + ) + fused_method.input_dtype = quant_method.input_dtype + fused_method.out_dtype = quant_method.out_dtype + fused_method.marlin_input_dtype = quant_method.marlin_input_dtype + self._fused_kv_linear.has_bias = self._fused_kv_bias is not None + fused_method.create_weights( + self._fused_kv_linear, + input_size_per_partition=input_size, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=self.hidden_norm.weight.dtype, + ) self._fused_kv_linear.register_parameter( "weight", nn.Parameter(self._fused_kv_weight, requires_grad=False) ) @@ -617,13 +630,13 @@ def process_weights_after_loading(self) -> None: "weight_scale", nn.Parameter(self._fused_kv_weight_scale, requires_grad=False), ) - quant_method.process_weights_after_loading(self._fused_kv_linear) - self._fused_kv_quant_method = quant_method + fused_method.process_weights_after_loading(self._fused_kv_linear) + self._fused_kv_quant_method = fused_method self._fused_kv_weight = None self._fused_kv_weight_scale = None logger.info_once( "Using %s for the fused DFlash context K/V projection.", - type(quant_method.kernel).__name__, + type(fused_method.kernel).__name__, ) def _project_context_kv( From 4ebd9823cd887311568bc587f6e4a733527b28c1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:21:42 +0000 Subject: [PATCH 080/365] fix(v2): register native sampling warmup contracts Register the existing top-k/top-p compile-key declarations on sampling ranks. Seeded and processed-logprob requests can select native filtering even when unseeded warmup uses FlashInfer. This covers singleton split-row kernels omitted by the bounded 8/16-row warmup; sampling arithmetic and the production JIT monitor default are unchanged. Validation: six GPU warmup cases pass for vocabularies 257/129280 and row limits 1/7/128. Every reachable row count and top-k/top-p subset runs with a hook that rejects post-warmup compilation. Applicable pre-commit hooks pass. The DS4.1 strict-monitor reproducer requires a rebuilt-image rerun. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/sample/test_topk_topp_sampler.py | 24 +++++++++++++++++++---- vllm/v1/worker/gpu/model_runner.py | 5 +++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 4d4740524de0..cbf1d7655c9e 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from torch import Generator @@ -13,6 +15,7 @@ apply_top_k_top_p_probs, apply_top_k_top_p_pytorch, random_sample, + register_top_k_top_p_warmups, warmup_top_k_top_p, ) from vllm.v1.sample.sampler import Sampler @@ -74,21 +77,34 @@ def test_apply_top_k_top_p_probs_matches_processed_logits( reason="requires GPU Triton compilation", ) @pytest.mark.parametrize("vocab_size", [257, 129280]) +@pytest.mark.parametrize("max_num_rows", [1, 7, 128]) def test_filter_warmup_covers_speculative_rows_and_constraint_subsets( - monkeypatch: pytest.MonkeyPatch, vocab_size: int + monkeypatch: pytest.MonkeyPatch, vocab_size: int, max_num_rows: int ): - """Four K5 requests must not compile top-p-only filtering after warmup.""" + """Seeded filtering must cover singleton, split-row and large-batch keys.""" from triton import knobs + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry + device = torch.device(DEVICE_TYPE) - warmup_top_k_top_p(vocab_size, 24, device) + registry = JitWarmupRegistry( + SimpleNamespace( + model_config=SimpleNamespace(get_vocab_size=lambda: vocab_size), + scheduler_config=SimpleNamespace(max_num_seqs=max_num_rows), + num_speculative_tokens=0, + ) + ) + with registry.activate(): + register_top_k_top_p_warmups() + registry.warmup() + warmup_top_k_top_p(vocab_size, max_num_rows, device) torch.accelerator.synchronize() def unexpected_compile(**kwargs): pytest.fail("Filtering compiled a kernel after bounded sampler warmup") monkeypatch.setattr(knobs.runtime, "jit_post_compile_hook", unexpected_compile) - for num_rows in (8, 9, 15, 16, 17, 24): + for num_rows in range(1, max_num_rows + 1): logits = torch.zeros(num_rows, vocab_size, device=device) top_k = torch.full((num_rows,), 20, dtype=torch.int32, device=device) top_p = torch.full((num_rows,), 0.95, device=device) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index aba519befaff..5b57f9e55146 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,6 +78,7 @@ ModelRunnerOutput, RoutedExpertsTensors, ) +from vllm.v1.sample.ops.topk_topp_sampler import register_top_k_top_p_warmups from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.watermarking.spec_decode import ( @@ -472,6 +473,10 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: + # Seeded and processed-logprob requests use native filtering even + # when unseeded warmup requests select FlashInfer sampling. + with self.jit_warmup_registry.activate(): + register_top_k_top_p_warmups() sampler_kwargs: dict[str, Any] = { "max_num_reqs": self.max_num_reqs, "vocab_size": self.vocab_size, From 2dda0dab431f31e9e6a3cd4233968c3241cf8092 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:41:37 +0000 Subject: [PATCH 081/365] fix(chat): normalize namespaced tool iterables before matching Materialize Python tool collections before qualifying namespace metadata and validating named choices. Lists, tuples and generators retain the same serialized function identities without mutating definitions. JSON request behavior is unchanged. Validation: tuple inputs reproduce a false tool-choice mismatch before the fix. The namespace/history suite passes for lists, tuples and generators; combined CPU suites pass 90 tests with 43 GPU-only skips. Applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/entrypoints/openai/test_tool_calls_serialization.py | 7 +++++-- vllm/entrypoints/openai/chat_completion/protocol.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index 156d6bb140c4..c62bff2d1ab6 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -160,8 +160,9 @@ def test_multiple_tool_calls_materialised(num_tool_calls: int): @pytest.mark.parametrize("location", ["tool", "function"]) @pytest.mark.parametrize("namespace", ["inventory", {"name": "inventory"}]) @pytest.mark.parametrize("iterator", [False, True]) +@pytest.mark.parametrize("tools_collection", [list, tuple, iter]) def test_namespaces_survive_definitions_choices_and_history( - location, namespace, iterator + location, namespace, iterator, tools_collection ): def qualify(item, value): target = item if location == "tool" else item["function"] @@ -185,6 +186,7 @@ def qualify(item, value): ], } original = copy.deepcopy(payload) + payload["tools"] = tools_collection(tools) if iterator: payload["messages"][1]["tool_calls"] = iter(calls) request = ChatCompletionRequest.model_validate(payload) @@ -200,8 +202,9 @@ def qualify(item, value): assert request.model_dump_json() == serialized restored = ChatCompletionRequest.model_validate_json(serialized) assert restored.model_dump_json() == serialized - if not iterator: + if not iterator and tools_collection is list: assert payload == original + assert tools == original["tools"] assert calls == original["messages"][1]["tool_calls"] diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index e5913efef6c6..3607ec221b6a 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -941,8 +941,11 @@ def check_tool_usage(cls, data): # Named choices must be compared with qualified identities before # Pydantic serializes function definitions and discards extra fields. data = dict(data) - if isinstance(data.get("tools"), list): - data["tools"] = [normalize_tool_namespace(tool) for tool in data["tools"]] + tools = data.get("tools") + if isinstance(tools, Iterable) and not isinstance( + tools, (str, bytes, bytearray, dict) + ): + data["tools"] = [normalize_tool_namespace(tool) for tool in tools] if isinstance(data.get("tool_choice"), dict): data["tool_choice"] = normalize_tool_namespace(data["tool_choice"]) From 1dcd01c02651b53f71cc5457fff1eb8d6a712e82 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 17 Sep 2026 23:42:16 +0000 Subject: [PATCH 082/365] fix(glm): reject undeclared selection before recurrent cache updates Check physical-selection resources before dispatching pooled-cache writes on pure DCP1 decode. An unbound main cache or missing selection plan now fails without advancing request state. Supported selection arithmetic and the hot-path checks are unchanged. Validation: both missing-resource cases reproduce an early cache-write attempt before the fix and pass afterwards. All 67 pooled-indexer tests pass on stock SM120, including changed-input replay and prepared selection; applicable commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 60 +++++++++++++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 17 ++++-- 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index ffb61df5ec9e..5b6d3ca555bc 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -689,6 +689,66 @@ def test_glm53_physical_selection_provider_is_explicit() -> None: ) +@pytest.mark.parametrize( + "main_blocks, error", + [(0, "main cache is not bound"), (1, "physical selection has not been declared")], +) +def test_glm53_missing_physical_plan_does_not_mutate_recurrent_cache( + monkeypatch, main_blocks, error +) -> None: + from vllm.models.glm5next.nvidia import pooled_indexer as module + + hidden = torch.zeros((1, 128)) + metadata = SimpleNamespace( + num_reqs=1, + num_actual_tokens=1, + num_decode_tokens=1, + num_decodes=1, + max_query_len=1, + block_table=torch.zeros((1, 1), dtype=torch.int32), + selector_state_slot_ids=torch.zeros(1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + indexer = SimpleNamespace( + max_tokens=1, + wq_b=lambda value: (torch.zeros((1, 32 * 128)), None), + wk=lambda value: (hidden, None), + k_norm=lambda value: value, + index_kpool_compress_gate=torch.zeros((128, 128)), + index_kpool_compress_ape=torch.zeros((4, 128)), + _project_head_weights=lambda value: torch.zeros((1, 32)), + _q_fp8=torch.empty((1, 32, 128), dtype=torch.float8_e4m3fn), + _q_scale=torch.empty((1, 32)), + main_layer_name="attention", + _index_cache=torch.zeros((1, 64, 132), dtype=torch.uint8), + _tail=torch.zeros((1, 2, 4, 128)), + _state_slots=Glm5NextPooledIndexer._state_slots, + _parent_table_width=1, + _parent_stride_pages=1, + block_size=256, + dcp_world_size=1, + _emit_physical_selection=True, + _main_cache_num_blocks=main_blocks, + _physical_selection_plan=None, + ) + monkeypatch.setattr(module, "fwht128_quant_fp8", lambda *args, **kwargs: None) + monkeypatch.setattr( + module, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={"attention": metadata}), + ) + + def reject_cache_update(*args, **kwargs): + pytest.fail("Missing physical-selection resources must not update the cache") + + monkeypatch.setattr(module, "update_decode_pools", reject_cache_update) + with pytest.raises(RuntimeError, match=error): + Glm5NextPooledIndexer.forward( + indexer, hidden, hidden, torch.zeros(1, dtype=torch.int64), None + ) + + def _packed_main_cache( *, device: torch.device, diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 3f15e9b2ce77..63ee0d63546f 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -568,6 +568,16 @@ def forward( f"actual={actual_table_width}, required={self._parent_table_width}" ) + decode_only = decode_rows == live_rows + emit_physical_selection = ( + decode_only and self.dcp_world_size == 1 and self._emit_physical_selection + ) + if emit_physical_selection: + if self._main_cache_num_blocks < 1: + raise RuntimeError("GLM selector main cache is not bound") + if self._physical_selection_plan is None: + raise RuntimeError("GLM physical selection has not been declared") + update_decode_pools( index_cache, self._tail, @@ -586,7 +596,6 @@ def forward( ) parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] seq_lens = self._pool_seq_lens[:live_rows] - decode_only = decode_rows == live_rows decode_table = self._decode_block_table[:decode_rows] if decode_only: prepare_c4_decode_metadata( @@ -707,11 +716,7 @@ def forward( output = self.topk_indices_buffer[:rows] if live_rows < rows: output[live_rows:].fill_(-1) - if decode_only and self.dcp_world_size == 1 and self._emit_physical_selection: - if self._main_cache_num_blocks < 1: - raise RuntimeError("GLM selector main cache is not bound") - if self._physical_selection_plan is None: - raise RuntimeError("GLM physical selection has not been declared") + if emit_physical_selection: self._expand_pooled_topk_to_physical_slots( pool_ids[:live_rows], positions[:live_rows], From 3e525ab0489c64cd68bf6d9dbad7751732b4ac7e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:52:25 +0000 Subject: [PATCH 083/365] Preserve qualified tool identities in tuple message histories Normalize tuple histories accepted by the Python request API before Pydantic union validation. Copy message dictionaries so caller-owned histories remain unchanged. JSON/list requests retain their existing behavior. Invalid scalar and mapping collections remain rejected. Validation: the tuple reproducer loses inventory::lookup before this fix. All 72 tool-call serialization tests and applicable pre-commit hooks pass with the fix; no model math or sampling behavior changes. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../openai/test_tool_calls_serialization.py | 30 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 9 +++++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py index c62bff2d1ab6..ec223ed2f70d 100644 --- a/tests/entrypoints/openai/test_tool_calls_serialization.py +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -229,6 +229,36 @@ def test_namespace_description_is_preserved_exactly_once(): ).model_dump() == (parsed.model_dump()) +@pytest.mark.parametrize("collection", [list, tuple]) +def test_iterable_history_preserves_namespaced_tool_identity(collection): + call = _make_tool_call("call_history", "lookup", "{}") + call["namespace"] = "inventory" + messages = [{"role": "assistant", "content": None, "tool_calls": [call]}] + original = copy.deepcopy(messages) + request = ChatCompletionRequest.model_validate( + {"model": "test-model", "messages": collection(messages)} + ) + assert request.messages[0]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + serialized = request.model_dump_json() + restored = ChatCompletionRequest.model_validate_json(serialized) + assert restored.messages[0]["tool_calls"][0]["function"]["name"] == ( + "inventory::lookup" + ) + assert request.model_dump_json() == serialized + assert messages == original + + +@pytest.mark.parametrize("messages", ["text", b"text", bytearray(b"text"), {}]) +def test_invalid_history_collection_remains_a_client_error(messages): + with pytest.raises(ValidationError) as error: + ChatCompletionRequest.model_validate( + {"model": "test-model", "messages": messages} + ) + assert any(item["loc"] == ("messages",) for item in error.value.errors()) + + @pytest.mark.parametrize("tool_calls", [1, False, 1.5]) def test_invalid_tool_calls_shape_reports_the_field(tool_calls): with pytest.raises(ValidationError) as error: diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 3607ec221b6a..89ecef017748 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -552,11 +552,18 @@ def _normalize_messages_before(cls, data: Any) -> Any: - Renames the deprecated ``reasoning_content`` field to ``reasoning`` so downstream code only needs to check one field. - Qualifies tool-call namespaces before union validation can drop them. + + Args: + data: Unvalidated request payload, including Python message sequences. + + Returns: + A copied payload with normalized list/tuple message histories, or + the unchanged input when it is not a supported payload shape. """ if not isinstance(data, dict): return data messages = data.get("messages") - if not isinstance(messages, list): + if not isinstance(messages, (list, tuple)): return data data = dict(data) messages = [dict(msg) if isinstance(msg, dict) else msg for msg in messages] From 25443f3ad36da076ac3e0517365dbc9e0f66d24c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:40:29 +0000 Subject: [PATCH 084/365] fix(attention): release sparse MLA preparation probes after synchronization Retain synthetic query and selection tensors through the prepared call closures, not through the published plan. Runtime kernels bind fresh caller-owned inputs; the compiled launcher does not need warmup tensor storage. This prevents per-layer probe retention during profiling and serving preparation without changing attention arithmetic or scratch ownership. Validation: the installed-wheel negative control fails the lifetime assertion; the patched CPU contract and native decode/prefill live-input CUDA-graph tests pass (3 cases). The TP2 allocator trace records 2,785,017,856 query bytes and 348,637,184 selection bytes alive when preparation exhausts memory. Whole-model validation requires the source-built wheel. Ruff and applicable hooks pass; six pre-existing mypy findings match the recorded parent control. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 105 ++++++++++++++++++ .../attention/backends/mla/b12x_mla_sparse.py | 7 +- 2 files changed, 109 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 7a85dbb686f6..35a9ea87d188 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1521,6 +1521,111 @@ def bind(**kwargs): assert calls == ["prime", "run", "prime", "run"] assert reservations == [(((spec.nbytes,), torch.uint8),)] * 2 + probes = [ + weakref.ref(bound[name]) + for name in ( + "q", + "selected_indices", + "cache_seqlens_int32", + "nsa_cache_seqlens_int32", + ) + ] + # Publishing retains owners, but must not retain synchronized warmup inputs. + published_owners = (*call.owners, *second_call.owners) + del call, second_call + bound.clear() + assert all(probe() is None for probe in probes), published_owners + + +@pytest.mark.parametrize("mode,rows", [("decode", 4), ("extend", 64)]) +@torch.inference_mode() +def test_sparse_mla_prepared_launcher_rebinds_without_warmup_probes(mode, rows): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("native B12X sparse MLA requires SM12x") + from b12x.attention import sparse_mla + from b12x.preparation import PreparationSession + + from vllm.v1.worker.workspace import ( + current_workspace_manager, + init_workspace_manager, + reset_workspace_manager, + ) + + device = torch.device("cuda", torch.accelerator.current_device_index()) + cache = torch.empty((1, 64, 528), dtype=torch.uint8, device=device) + keys = torch.ones((64, 512), dtype=torch.bfloat16, device=device) + keys[1].fill_(2) + slots = torch.arange(64, dtype=torch.int64, device=device) + sparse_mla.concat_and_cache_glm_next_mla_fp8( + keys, cache, slots, plan=sparse_mla.plan_cache_writer(keys, cache, slots) + ) + caps = sparse_mla.Caps( + device=device, + num_q_heads=16, + max_q_rows=rows, + max_width=2051, + softmax_scale=256**-0.5, + dtype=torch.bfloat16, + kv_dtype=torch.uint8, + head_dim=512, + v_head_dim=512, + model_type=int(sparse_mla.ModelType.GLM_NEXT), + mode=mode, + max_batch=rows, + page_size=64, + ) + plan = sparse_mla.plan(caps) + (spec,) = plan.scratch_specs() + impl = object.__new__(B12xMLASparseImpl) + impl._bound_kv_cache = cache + impl._scratch_nbytes = spec.nbytes + init_workspace_manager(device) + try: + with PreparationSession(device=device, autotune=False) as session: + session.prepare( + ( + plan.request( + name=f"probe-lifetime.{mode}", + prepare_call=lambda state: impl._make_prepare_call(state, caps), + ), + ) + ) + assert plan.prepared.owners == () + session.freeze() + current_workspace_manager().lock() + query = torch.zeros((rows, 16, 512), dtype=torch.bfloat16, device=device) + selected = torch.zeros((rows, 2051), dtype=torch.int32, device=device) + lengths = torch.full((rows,), 64, dtype=torch.int32, device=device) + selected_lengths = torch.ones(rows, dtype=torch.int32, device=device) + scratch = torch.empty(spec.shape, dtype=spec.dtype, device=device) + state = plan.prepared.state + + def run(): + binding = state.bind( + scratch=scratch, + q=query, + selected_indices=selected, + cache_seqlens_int32=lengths, + nsa_cache_seqlens_int32=selected_lengths, + kv_cache=cache, + ) + output = state.run(binding, kv_cache=cache) + return output[0] if isinstance(output, tuple) else output + + actual = run() + torch.testing.assert_close(actual, torch.ones_like(actual), rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + for index in (1, 0, 1): + selected[:, 0].fill_(index) + graph.replay() + torch.testing.assert_close( + captured, torch.full_like(captured, index + 1), rtol=0, atol=0 + ) + finally: + reset_workspace_manager() + @pytest.mark.parametrize("query_projection", [False, True]) def test_mla_layer_discovers_backend_and_optional_query_plans( diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 793a17bee537..21155440d5f2 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1462,8 +1462,8 @@ def _make_prepare_call(self, state, caps): ) selected_lengths = torch.ones(rows, dtype=torch.int32, device=caps.device) (scratch_spec,) = state.scratch_specs() - # Preparation calls execute serially but retain their owners. Borrow - # caller-owned scratch so exact-row plans do not each pin a workspace. + # Preparation calls execute serially. Borrow caller-owned scratch so + # exact-row plans do not each allocate a workspace. workspace = current_workspace_manager() workspace.reserve_all(((self._scratch_nbytes,), torch.uint8)) (scratch,) = workspace.get_simultaneous( @@ -1487,7 +1487,8 @@ def produce(): return PreparedCall( run=lambda: state.run(binding, kv_cache=kv_cache), produce=produce, - owners=(q, selected, cache_lengths, selected_lengths, scratch, binding), + # Closures retain probes until preparation has synchronized. The + # published launcher binds live inputs and must not pin these probes. ) def get_b12x_preparation_units( From 0d26f8144421342a8d042598bd75746041955024 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:43:19 +0000 Subject: [PATCH 085/365] test(attention): cover GLM TP2 prefill preparation geometry Exercise 4096 query rows and 64 heads through the prepared sparse-MLA launcher, then replay changed selections in a CUDA graph. This covers the declaration shape that exhausted memory with retained warmup probes. The lifetime contract plus three native shape cases pass on stock SM120. Serving source is unchanged from 25443f3ad36; its published wheel remains the implementation under qualification. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_sparse_mla_api.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 35a9ea87d188..06bc53bf9da9 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1537,9 +1537,12 @@ def bind(**kwargs): assert all(probe() is None for probe in probes), published_owners -@pytest.mark.parametrize("mode,rows", [("decode", 4), ("extend", 64)]) +@pytest.mark.parametrize( + "mode,rows,heads", + [("decode", 4, 16), ("extend", 64, 16), ("extend", 4096, 64)], +) @torch.inference_mode() -def test_sparse_mla_prepared_launcher_rebinds_without_warmup_probes(mode, rows): +def test_sparse_mla_prepared_launcher_rebinds_without_warmup_probes(mode, rows, heads): if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: pytest.skip("native B12X sparse MLA requires SM12x") from b12x.attention import sparse_mla @@ -1561,7 +1564,7 @@ def test_sparse_mla_prepared_launcher_rebinds_without_warmup_probes(mode, rows): ) caps = sparse_mla.Caps( device=device, - num_q_heads=16, + num_q_heads=heads, max_q_rows=rows, max_width=2051, softmax_scale=256**-0.5, @@ -1593,7 +1596,7 @@ def test_sparse_mla_prepared_launcher_rebinds_without_warmup_probes(mode, rows): assert plan.prepared.owners == () session.freeze() current_workspace_manager().lock() - query = torch.zeros((rows, 16, 512), dtype=torch.bfloat16, device=device) + query = torch.zeros((rows, heads, 512), dtype=torch.bfloat16, device=device) selected = torch.zeros((rows, 2051), dtype=torch.int32, device=device) lengths = torch.full((rows,), 64, dtype=torch.int32, device=device) selected_lengths = torch.ones(rows, dtype=torch.int32, device=device) From 719e2a05466a0c132bd351a51a390cb6fa08cfc9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:45:30 +0000 Subject: [PATCH 086/365] test(attention): vary live queries and selection counts after preparation Check CUDA-graph replay with different query values, selected indices and active selection lengths after warmup probe storage is released. Uniform and saturated two-token attention cases have exact BF16 results. All four lifetime/native tests pass on SM120, including the 4096-row 64-head GLM shape; production source is unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_sparse_mla_api.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 06bc53bf9da9..23d5d3e38c00 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1620,11 +1620,20 @@ def run(): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured = run() - for index in (1, 0, 1): + selected[:, 1].fill_(1) + for index, count, query_first, expected in ( + (1, 1, 0, 2), + (0, 1, 0, 1), + (0, 2, 0, 1.5), + (0, 2, 256, 2), + (0, 2, -256, 1), + ): selected[:, 0].fill_(index) + selected_lengths.fill_(count) + query[..., 0].fill_(query_first) graph.replay() torch.testing.assert_close( - captured, torch.full_like(captured, index + 1), rtol=0, atol=0 + captured, torch.full_like(captured, expected), rtol=0, atol=0 ) finally: reset_workspace_manager() From fb126faa03896db335a08b7de8166226d9a791cc Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 14:27:59 +0000 Subject: [PATCH 087/365] fix(flashkda): use CTA-local workspace copies on SM120 SM120 lowers cluster-addressed bulk copies through a CUDA device syscall that grows the per-thread stack to 14608 bytes. Copy into CTA-owned shared memory directly while retaining Hopper and SM100 dispatch and all tensor layouts. The installed extension fails the fresh-context stack regression; a minimal CUDA copy reproducer confirms exact bytes with the CTA instruction and no stack growth. Full patched-extension and model qualification follow the source-wheel build. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- cmake/external_projects/flashkda.cmake | 5 ++ .../patches/flashkda-sm120-cta-copy.patch | 67 +++++++++++++++++++ tests/models/kimi_k3/test_kda.py | 55 +++++++++++++++ 3 files changed, 127 insertions(+) create mode 100644 cmake/external_projects/patches/flashkda-sm120-cta-copy.patch diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 10d3598ed390..5e75595d74a2 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -21,6 +21,11 @@ else() "-DSOURCE_DIR=" "-DPATCH_FILE=${CMAKE_CURRENT_LIST_DIR}/patches/flashkda-packed-checkpoints.patch" -P "${CMAKE_CURRENT_LIST_DIR}/apply_flashkda_checkpoint_patch.cmake" + COMMAND + "${CMAKE_COMMAND}" + "-DSOURCE_DIR=" + "-DPATCH_FILE=${CMAKE_CURRENT_LIST_DIR}/patches/flashkda-sm120-cta-copy.patch" + -P "${CMAKE_CURRENT_LIST_DIR}/apply_flashkda_checkpoint_patch.cmake" ) endif() diff --git a/cmake/external_projects/patches/flashkda-sm120-cta-copy.patch b/cmake/external_projects/patches/flashkda-sm120-cta-copy.patch new file mode 100644 index 000000000000..89e5157a2ff0 --- /dev/null +++ b/cmake/external_projects/patches/flashkda-sm120-cta-copy.patch @@ -0,0 +1,67 @@ +diff --git a/csrc/smxx/fwd_kernel2.cuh b/csrc/smxx/fwd_kernel2.cuh +--- a/csrc/smxx/fwd_kernel2.cuh ++++ b/csrc/smxx/fwd_kernel2.cuh +@@ -2,6 +2,25 @@ + + #include "utils.cuh" + ++__device__ __forceinline__ void copy_cta_workspace( ++ void const* source, uint64_t* barrier, void* destination, int32_t bytes ++) { ++#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 ++ // Both shared addresses belong to this CTA. The cluster-addressed bulk ++ // copy uses a device syscall on SM120 and reserves a large context stack. ++ uint32_t dst = cute::cast_smem_ptr_to_uint(destination); ++ uint32_t mbar = cute::cast_smem_ptr_to_uint(barrier); ++ asm volatile( ++ "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " ++ "[%0], [%1], %2, [%3];\n" ++ : ++ : "r"(dst), "l"(source), "r"(bytes), "r"(mbar) ++ : "memory"); ++#else ++ cute::SM90_BULK_COPY_G2S::copy(source, barrier, destination, bytes); ++#endif ++} ++ + template + struct K2Layouts { + using MMALayout = decltype(tile_to_shape( +@@ -378,7 +397,7 @@ + // K1 stores byte images of its swizzled shared-memory tensors. + // K2 uses the same layouts, so raw bulk copies restore them + // directly without tensor-map coordinate work or repacking. +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_kd + int64_t(ws_idx) * (CHUNK * D), + reinterpret_cast(tma_barrier), + shared_storage.input[stage].k_decayed.begin(), +@@ -386,7 +405,7 @@ +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_qd + int64_t(ws_idx) * (CHUNK * D), + reinterpret_cast(tma_barrier), + shared_storage.input[stage].q_decayed.begin(), + int32_t(CHUNK * D * sizeof(BF16))); +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_kr + int64_t(ws_idx) * (CHUNK * D), +@@ -396,7 +415,7 @@ +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_gt + int64_t(ws_idx) * D, + reinterpret_cast(tma_barrier), + shared_storage.input[stage].g_total.begin(), + int32_t(D * sizeof(float))); +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_inv + int64_t(ws_idx) * (CHUNK * CHUNK), +@@ -406,7 +425,7 @@ +- cute::SM90_BULK_COPY_G2S::copy( ++ copy_cta_workspace( + ws_mqk + int64_t(ws_idx) * (CHUNK * CHUNK), + reinterpret_cast(tma_barrier), + shared_storage.input[stage].Mqk.begin(), + int32_t(CHUNK * CHUNK * sizeof(BF16))); + + ++load_write; diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index f753c0acb78f..9ddf49896c6e 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -6,6 +6,9 @@ Uses torch.rand for q/k/v to match FLA's test pattern. """ +import subprocess +import sys +import textwrap from types import SimpleNamespace import pytest @@ -1144,6 +1147,58 @@ def test_fused_kda_decode_rejects_speculative_conv_state(): ) +def test_flashkda_sm120_local_copy_preserves_cuda_stack_limit(): + """CTA-local workspace loads must not reserve a cluster-copy syscall stack.""" + capability = current_platform.get_device_capability() + if not current_platform.is_cuda() or capability is None or capability.major != 12: + pytest.skip("This regression concerns the SM120 bulk-copy lowering") + if not is_flashkda_supported(128, torch.bfloat16, torch.float32, -5.0): + pytest.skip("FlashKDA is not supported on this platform") + + # A fresh context prevents another kernel's retained stack from hiding the + # first-launch reservation. No model weights or memory-pressure allocation. + subprocess.run( + [ + sys.executable, + "-c", + textwrap.dedent(""" + import torch + import vllm._flashkda_C + from cuda.bindings import runtime + + def stack_limit(): + status, limit = runtime.cudaDeviceGetLimit( + runtime.cudaLimit.cudaLimitStackSize) + assert int(status) == 0 + return limit + + q = torch.zeros((1, 32, 2, 128), dtype=torch.bfloat16, device="cuda") + beta = torch.zeros((1, 32, 2), dtype=q.dtype, device=q.device) + state = torch.zeros((1, 2, 128, 128), device=q.device) + final = torch.full_like(state, float("nan")) + out = torch.full_like(q, float("nan")) + a_log = torch.zeros(2, device=q.device) + bias = torch.zeros((2, 128), device=q.device) + lengths = torch.tensor([0, 32], dtype=torch.int32, device=q.device) + scratch = torch.empty( + torch.ops._flashkda_C.get_workspace_size(32, 2, 1), + dtype=torch.uint8, device=q.device) + torch.accelerator.synchronize() + before = stack_limit() + torch.ops._flashkda_C.fwd( + q, q, q, q, beta, 128**-0.5, out, scratch, + a_log, bias, -5.0, state, final, lengths) + torch.accelerator.synchronize() + after = stack_limit() + assert (out == 0).all() and (final == 0).all() + assert after <= before, f"CUDA stack grew: {before} -> {after}" + """), + ], + check=True, + timeout=120, + ) + + @torch.inference_mode() def test_flashkda_near_collinear_keys_remain_finite(): """Guard against unstable inversion of near-collinear key blocks.""" From 28eb7e1e4d4cdf2ee9cd82c5feae58ceaa1d9b8c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 14:46:15 +0000 Subject: [PATCH 088/365] fix(kda): release B12X prefill preparation activations Allocate synthetic prefill activations within each preparation call instead of retaining a full token-capacity bank on every KDA layer. Do not publish transient scratch through PreparedCall.owners; synchronized call closures retain it until preparation finishes. Live metadata, request-owned inputs, kernel arithmetic, and graph workspace contracts are unchanged. Validation: 14 adapter/lifetime/native tests pass on RTX PRO 6000 Workstation; two lifetime assertions fail against the installed unmodified wheel. An independent 12-case native B12X oracle/checkpoint/graph suite also passes. Whole-model Spark TP2 startup remains pending. The failing whole-model allocation census attributes 5,713,188,504 bytes per rank to prefill probe buffers and metadata. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_model.py | 162 ++++++++++++++++-- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 66 ++----- 2 files changed, 164 insertions(+), 64 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 8dd7d4c8dca8..bc320db1c0a5 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -1856,6 +1856,7 @@ def __init__(self, **kwargs): torch.nn.Module.__init__(layer) layer.head_dim = 4 layer.local_num_heads = 3 + layer.model_config = SimpleNamespace(dtype=torch.float32) layer.gate_lower_bound = -5.0 layer.A_log = torch.empty(3) layer.dt_bias = torch.empty(12) @@ -1874,13 +1875,6 @@ def __init__(self, **kwargs): layer._b12x_kda_state_indices = torch.empty((1, 1), dtype=torch.int32) layer._b12x_kda_num_seqs = torch.empty(1, dtype=torch.int32) layer._b12x_kda_num_tokens = torch.empty(1, dtype=torch.int32) - layer._b12x_prefill_q = torch.empty((2, 3, 4)) - layer._b12x_prefill_k = torch.empty((2, 3, 4)) - layer._b12x_prefill_v = torch.empty((2, 3, 4)) - layer._b12x_prefill_raw_g = torch.empty((2, 3, 4)) - prefill_beta_storage = torch.empty((2, beta_width)) - layer._b12x_prefill_raw_beta = prefill_beta_storage.narrow(1, beta_offset, 3) - layer._b12x_prefill_output = torch.empty((2, 3, 4)) layer._b12x_prefill_cu_seqlens = torch.empty(2, dtype=torch.int32) layer._b12x_prefill_initial_indices = torch.empty(1, dtype=torch.int32) layer._b12x_prefill_null_indices = torch.empty(1, dtype=torch.int32) @@ -1922,12 +1916,6 @@ def run(self, binding, **kwargs): "num_tokens": layer._b12x_kda_num_tokens, } prefill_sources = { - "q": layer._b12x_prefill_q, - "k": layer._b12x_prefill_k, - "v": layer._b12x_prefill_v, - "raw_g": layer._b12x_prefill_raw_g, - "raw_beta": layer._b12x_prefill_raw_beta, - "output": layer._b12x_prefill_output, "cu_seqlens": layer._b12x_prefill_cu_seqlens, "initial_state_indices": layer._b12x_prefill_initial_indices, "checkpoint_state_indices": layer._b12x_prefill_null_indices, @@ -1949,6 +1937,154 @@ def run(self, binding, **kwargs): source_alignment = min(16, source_pointer & -source_pointer) trial_alignment = min(16, trial_pointer & -trial_pointer) assert trial_alignment == source_alignment + for name in ("q", "k", "v", "raw_g", "output", "raw_beta"): + tensor = state.prefill[name] + assert tensor.shape == ((2, 3) if name == "raw_beta" else (2, 3, 4)) + assert tensor.dtype == layer.model_config.dtype + assert tensor.is_contiguous() + + +@pytest.mark.parametrize("benchmark", [False, True]) +def test_b12x_kda_prefill_preparation_releases_synthetic_activations( + monkeypatch, benchmark +): + """Only live metadata persists after the synchronized preparation call.""" + monkeypatch.setattr(current_platform, "current_device", lambda: "cpu") + monkeypatch.setattr(kimi_gdn_linear_attn, "get_b12x_kda_prefill", lambda: object()) + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + torch.nn.Module.__init__(layer) + layer.kda_prefill_backend = "b12x" + layer.head_dim, layer.local_num_heads = 128, 2 + layer.model_config = SimpleNamespace(dtype=torch.bfloat16) + layer.A_log = torch.zeros(2) + layer.dt_bias = torch.zeros(2, 128) + layer.gate_lower_bound = -5.0 + layer.kv_cache = (torch.empty(1), torch.zeros(2, 2, 128, 128)) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=64, max_num_seqs=4) + ) + layer._initialize_b12x_kda_prefill(config) + assert sum(t.numel() * t.element_size() for t in layer.buffers()) < 1024 + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + class State: + layout = SimpleNamespace( + scratch_specs=lambda: [SimpleNamespace(shape=(128,), dtype=torch.uint8)] + ) + + def bind(self, **tensors): + references.extend( + weakref.ref(tensors[name]) + for name in ("scratch", "q", "k", "v", "raw_g", "raw_beta", "output") + ) + return SimpleNamespace(**tensors) + + def run(self, binding, **kwargs): + assert binding.q.shape == (64, 2, 128) + assert binding.q.is_contiguous() + assert binding.num_tokens.item() == 64 + binding.output.copy_(binding.q) + + call = layer._b12x_kda_prefill_call(State(), benchmark=benchmark) + assert call.owners == () + call.produce() + call.run() + assert all(reference() is not None for reference in references) + del call + assert all(reference() is None for reference in references) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +def test_b12x_kda_prefill_live_inputs_after_preparation(): + """Released synthetic inputs cannot replace live request/state tensors.""" + from b12x.preparation import PreparationSession + from b12x.sequence.kda_prefill.reference import prefill_kda + + api = kimi_gdn_linear_attn.get_b12x_kda_prefill() + device = torch.device(current_platform.current_device()) + if api is None or not api.is_supported(device): + pytest.skip("B12X KDA prefill is unavailable") + torch.manual_seed(73) + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + torch.nn.Module.__init__(layer) + layer.kda_prefill_backend = "b12x" + layer.head_dim, layer.local_num_heads = 128, 2 + layer.model_config = SimpleNamespace(dtype=torch.bfloat16) + layer.cache_config = SimpleNamespace(mamba_cache_dtype="auto") + layer.A_log = torch.zeros(2, device=device) + layer.dt_bias = torch.zeros(2, 128, device=device) + layer.gate_lower_bound = -5.0 + pool = torch.zeros(8, 2, 128, 128, device=device) + layer.kv_cache = (torch.empty(1, device=device), pool) + layer._initialize_b12x_kda_prefill( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=128, max_num_seqs=4) + ) + ) + plan = layer._b12x_kda_prefill_declaration(8) + layer._b12x_prefill_plan = plan + request = plan.request( + name="glm-kda-prefill-live-inputs", + prepare_call=layer._prepare_b12x_kda_prefill, + benchmark_call=layer._benchmark_b12x_kda_prefill, + ) + session = PreparationSession(device=device, autotune=False) + prepared = session.prepare((request,)) + try: + (spec,) = plan.scratch_specs() + scratch = torch.empty(spec.shape, dtype=spec.dtype, device=device) + q, k, v, gate = ( + torch.randn(64, 2, 128, dtype=torch.bfloat16, device=device) + for _ in range(4) + ) + beta = torch.randn(64, 2, dtype=q.dtype, device=device) + lengths = torch.tensor([0, 64], dtype=torch.int32, device=device) + indices = torch.tensor([3], dtype=torch.int32, device=device) + offsets = torch.zeros(1, dtype=torch.int32, device=device) + initial = torch.ones(1, dtype=torch.bool, device=device) + pool.normal_() + expected_pool = pool.clone() + expected = prefill_kda( + q, + k, + v, + gate, + beta, + layer.A_log, + layer.dt_bias, + expected_pool, + lengths, + indices, + indices, + indices, + offsets, + 1, + 64, + lower_bound=-5.0, + ) + output = torch.full_like(q, float("nan")) + layer._run_b12x_kda_prefill( + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=gate, + raw_beta=beta, + cu_seqlens=lengths, + state_indices=indices, + has_initial_state=initial, + checkpoint=None, + recurrent_state=pool, + output=output, + ) + torch.accelerator.synchronize() + torch.testing.assert_close(output, expected, rtol=0.02, atol=0.02) + torch.testing.assert_close(pool[3], expected_pool[3], rtol=0.01, atol=0.005) + torch.testing.assert_close(pool[:3], expected_pool[:3], rtol=0, atol=0) + torch.testing.assert_close(pool[4:], expected_pool[4:], rtol=0, atol=0) + finally: + prepared.close() + session.close() @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 47ad93288542..33a6fd2b55f1 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -756,24 +756,7 @@ def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: self._b12x_prefill_api = api self._b12x_prefill_max_tokens = int(scheduler_config.max_num_batched_tokens) self._b12x_prefill_max_seqs = int(scheduler_config.max_num_seqs) - max_tokens, max_seqs = ( - self._b12x_prefill_max_tokens, - self._b12x_prefill_max_seqs, - ) - shape = (max_tokens, self.local_num_heads, self.head_dim) - for name, tensor_shape in ( - ("_b12x_prefill_q", shape), - ("_b12x_prefill_k", shape), - ("_b12x_prefill_v", shape), - ("_b12x_prefill_raw_g", shape), - ("_b12x_prefill_raw_beta", (max_tokens, self.local_num_heads)), - ("_b12x_prefill_output", shape), - ): - self.register_buffer( - name, - torch.zeros(tensor_shape, dtype=self.model_config.dtype, device=device), - persistent=False, - ) + max_seqs = self._b12x_prefill_max_seqs self.register_buffer( "_b12x_prefill_cu_seqlens", torch.zeros(max_seqs + 1, dtype=torch.int32, device=device), @@ -1055,13 +1038,21 @@ def _b12x_kda_prefill_call(self, state, *, benchmark: bool): # Trial and prepare factories own their scratch; the runtime binding # in _run_b12x_kda_prefill draws from the workspace manager instead. scratch = torch.empty(spec[0], dtype=spec[1], device=slots.device) + # Synthetic activations belong to one preparation call, not each model + # layer or the published plan. Serving binds the live request tensors. + shape = ( + self._b12x_prefill_max_tokens, + self.local_num_heads, + self.head_dim, + ) + q, k, v, raw_g, output = ( + torch.empty(shape, dtype=self.model_config.dtype, device=slots.device) + for _ in range(5) + ) + raw_beta = torch.empty( + shape[:2], dtype=self.model_config.dtype, device=slots.device + ) if benchmark: - q = self._b12x_trial_tensor(self._b12x_prefill_q) - k = self._b12x_trial_tensor(self._b12x_prefill_k) - v = self._b12x_trial_tensor(self._b12x_prefill_v) - raw_g = self._b12x_trial_tensor(self._b12x_prefill_raw_g) - raw_beta = self._b12x_trial_tensor(self._b12x_prefill_raw_beta) - output = self._b12x_trial_tensor(self._b12x_prefill_output) cu_seqlens = self._b12x_trial_tensor(self._b12x_prefill_cu_seqlens) indices = self._b12x_trial_tensor(self._b12x_prefill_initial_indices) checkpoint_indices = self._b12x_trial_tensor( @@ -1071,32 +1062,7 @@ def _b12x_kda_prefill_call(self, state, *, benchmark: bool): num_seqs = self._b12x_trial_tensor(self._b12x_prefill_num_seqs) num_tokens = self._b12x_trial_tensor(self._b12x_prefill_num_tokens) saved_state = slots[slot : slot + 1].clone() - owners: tuple[torch.Tensor, ...] = ( - scratch, - q, - k, - v, - raw_g, - raw_beta, - output, - cu_seqlens, - indices, - checkpoint_indices, - offsets, - num_seqs, - num_tokens, - ) else: - q, k, v = ( - self._b12x_prefill_q, - self._b12x_prefill_k, - self._b12x_prefill_v, - ) - raw_g, raw_beta = ( - self._b12x_prefill_raw_g, - self._b12x_prefill_raw_beta, - ) - output = self._b12x_prefill_output cu_seqlens = self._b12x_prefill_cu_seqlens indices = self._b12x_prefill_initial_indices checkpoint_indices = self._b12x_prefill_null_indices @@ -1106,7 +1072,6 @@ def _b12x_kda_prefill_call(self, state, *, benchmark: bool): self._b12x_prefill_num_tokens, ) saved_state = None - owners = (scratch,) def produce(): for tensor in (q, k, v, raw_g, raw_beta): @@ -1152,7 +1117,6 @@ def reset(): produce=produce, reset=reset, restore=reset if saved_state is not None else None, - owners=owners, ) def _run_b12x_kda_prefill( From 7a6d0cf0ea40816df436d7e8f0a26f77a2888764 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 15:17:06 +0000 Subject: [PATCH 089/365] fix(attention): release GLM profiling cache views on unbind Detach the sparse MLA cache-writer binding and the pooled C4 indexer cache/plan references when their parent layer releases its KV pool. Preserve geometry and permit binding a replacement pool before preparation. The TP2 allocation census identifies 441188352 bytes retained by these references after graph profiling. Validation: 13 focused CPU/native SM120 tests pass, including live graph replay and rebind. Both new ownership assertions fail against the installed unmodified wheel. The six mypy findings are unchanged from the qualified parent baseline; other hooks pass. Whole-model admission is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 5 +++ .../v1/attention/test_b12x_sparse_mla_api.py | 37 +++++++++++++++++++ .../layers/attention/mla_attention.py | 3 ++ .../models/deepseek_v4/nvidia/b12x_indexer.py | 5 +++ vllm/models/glm5next/nvidia/pooled_indexer.py | 1 + .../attention/backends/mla/b12x_mla_sparse.py | 4 ++ 6 files changed, 55 insertions(+) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 21248d24b3b2..3be850cc9911 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -821,10 +821,15 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( assert indexer.indexer_op._num_q_heads == indexer.scratch.q_fp8.shape[1] == 32 assert indexer._physical_selection_plan is None indexer._physical_selection_plan = object() + indexer.indexer_op._plans[("decode", 8)] = object() indexer.unbind_main_kv_cache() assert indexer._physical_selection_plan is None assert indexer._index_cache is None assert indexer._main_cache_num_blocks == 0 + assert indexer.indexer_op._index_cache is None + assert indexer.indexer_op._plans == {} + indexer.bind_main_kv_cache(main) + assert indexer.indexer_op._index_cache is indexer._index_cache def test_glm53_packed_tail_accepts_nvfp4_main_record() -> None: diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 23d5d3e38c00..7850c95b4321 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1695,6 +1695,43 @@ def get_b12x_preparation_units(self, owner, workload): assert [request.name for request in units[1].requests] == ["test.query.m8"] +@pytest.mark.parametrize("b12x_backend", [False, True]) +def test_mla_unbind_releases_backend_cache_and_allows_rebinding(b12x_backend): + """Profiling-cache views must not survive the owning layer's teardown.""" + from vllm.model_executor.layers.attention.mla_attention import MLAAttention + from vllm.v1.worker.utils import clear_layer_kv_caches + + layer = MLAAttention.__new__(MLAAttention) + torch.nn.Module.__init__(layer) + impl = object.__new__(B12xMLASparseImpl) if b12x_backend else SimpleNamespace() + if b12x_backend: + impl._uses_glm_dsa_nvfp4_cache = False + impl._is_glm_next = False + layer.impl = impl + layer.indexer = None + layer._vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(enable_jit_warmup=False) + ) + cache = torch.empty((2, 1, 64, 528), dtype=torch.uint8) + cache_ref = weakref.ref(cache) + layer.bind_kv_cache(cache) + if b12x_backend: + impl._cache_writer_plan = SimpleNamespace(cache=impl._bound_kv_cache) + del cache + + clear_layer_kv_caches([layer, layer]) + assert layer.kv_cache.numel() == 0 + assert cache_ref() is None + if b12x_backend: + assert impl._bound_kv_cache is None + assert impl._cache_writer_plan is None + replacement = torch.empty((3, 1, 64, 528), dtype=torch.uint8) + layer.bind_kv_cache(replacement) + assert layer.kv_cache.data_ptr() == replacement.data_ptr() + if b12x_backend: + assert impl._bound_kv_cache is layer.kv_cache + + def test_b12x_sparse_mla_plan_lookup_declares_unplanned_decode_rows_once( monkeypatch, ) -> None: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 29b30379679b..53e913f0ee03 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -921,6 +921,9 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: bind_indexer(self.kv_cache) def unbind_kv_cache(self) -> None: + unbind_impl = getattr(self.impl, "unbind_kv_cache", None) + if unbind_impl is not None: + unbind_impl() unbind_indexer = getattr(self.indexer, "unbind_main_kv_cache", None) if unbind_indexer is not None: unbind_indexer() diff --git a/vllm/models/deepseek_v4/nvidia/b12x_indexer.py b/vllm/models/deepseek_v4/nvidia/b12x_indexer.py index b85fc94343e0..10994fd6b8b7 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x_indexer.py +++ b/vllm/models/deepseek_v4/nvidia/b12x_indexer.py @@ -347,6 +347,11 @@ def describe(workload): self, describe, group=get_dcp_group() ) + def clear_b12x_index_cache(self) -> None: + """Detach pool-derived views and plans before releasing the KV pool.""" + self._plans.clear() + self._index_cache = None + @property def _num_q_heads(self) -> int: return self._index_num_q_heads or int(getattr(self.k_cache, "num_q_heads", 1)) diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 9e94d04e6455..27b71348f772 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -395,6 +395,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: self.block_size = block_size def unbind_main_kv_cache(self) -> None: + self.indexer_op.clear_b12x_index_cache() self._index_cache = None self._main_cache_num_blocks = 0 self._physical_selection_plan = None diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 21155440d5f2..6776c452d1ee 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1813,6 +1813,10 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: self._cache_writer_plan = None self._bound_kv_cache = kv_cache + def unbind_kv_cache(self) -> None: + self._cache_writer_plan = None + self._bound_kv_cache = None + def do_kv_cache_update( self, kv_c_normed: torch.Tensor, From 9ccde629d4ad7fa5f8e28e8345d39c89aaae6be1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 15:23:34 +0000 Subject: [PATCH 090/365] fix(kda): bound decode preparation probes to their call lifetime Construct synthetic decode activations only when declaring or preparing a kernel, preserving the fused beta view stride and alignment. Do not publish temporary scratch as plan owners. Live request bindings, recurrent state, dtypes, and CUDA graph execution are unchanged. The GLM TP2 census contains 322256896 bytes of unnecessary persistent activation probes across 34 KDA layers. Validation: 37 KDA component tests and two prepared native decode graph replay cases pass on SM120. Both initialization-lifetime cases fail against the installed parent image. Lint and type hooks pass. Whole-model cache admission is validated by the source-wheel build separately. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/test_glm5next_model.py | 135 ++++++++++++++---- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 94 ++++-------- 2 files changed, 135 insertions(+), 94 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index bc320db1c0a5..790607dbdc20 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -1690,6 +1690,7 @@ def test_glm_adaptive_kda_backend_is_scoped_to_aligned_state_cache( def test_glm5next_b12x_kda_plan_reserves_null_state_zero(monkeypatch) -> None: captured_caps = {} + probes: list[weakref.ReferenceType[torch.Tensor]] = [] class FakeApi: @staticmethod @@ -1699,6 +1700,10 @@ def Caps(**kwargs): @staticmethod def invocation_from_tensors(caps, **tensors): + probes.extend( + weakref.ref(tensors[name]) + for name in ("mixed_qkv", "raw_g", "raw_beta", "z", "output") + ) return None @staticmethod @@ -1711,22 +1716,19 @@ def plan(caps, *, invocation=None): layer._b12x_kda_max_tokens = 16 layer._b12x_kda_max_seqs = 4 layer._b12x_kda_state_index_columns = 4 + layer._b12x_kda_beta_row_width = 8 + layer._b12x_kda_beta_offset = 0 layer.local_num_heads = 8 layer.head_dim = 128 layer.model_config = SimpleNamespace(dtype=torch.bfloat16) layer.b12x_kda_null_state_index = 0 empty = torch.empty(0) for name in ( - "_b12x_kda_mixed_qkv", - "_b12x_kda_raw_g", - "_b12x_kda_raw_beta", - "_b12x_kda_z", "_b12x_kda_query_start_loc", "_b12x_kda_num_accepted_tokens", "_b12x_kda_state_indices", "_b12x_kda_num_seqs", "_b12x_kda_num_tokens", - "_b12x_kda_output", ): setattr(layer, name, empty) layer.A_log = empty @@ -1749,6 +1751,7 @@ def plan(caps, *, invocation=None): assert plan == captured_caps assert captured_caps["null_state_index"] == 0 + assert all(probe() is None for probe in probes) @pytest.mark.parametrize("extra_padding", [0, 16]) @@ -1798,7 +1801,9 @@ def is_supported(device): layer._initialize_b12x_kda_decode(vllm_config) - assert layer._b12x_kda_raw_g.shape == (16, 16, 128) + _, raw_g, raw_beta, _, _ = layer._b12x_kda_decode_probes() + assert raw_g.shape == (16, 16, 128) + assert sum(t.numel() * t.element_size() for t in layer.buffers()) < 1024 beta_offset = ( 4 * layer.local_projection_size + layer.head_dim if use_full_rank_gate @@ -1807,8 +1812,8 @@ def is_supported(device): live_beta = torch.empty(16, live_width).narrow( 1, beta_offset, layer.local_num_heads ) - assert layer._b12x_kda_raw_beta.shape == live_beta.shape - assert layer._b12x_kda_raw_beta.stride() == live_beta.stride() + assert raw_beta.shape == live_beta.shape + assert raw_beta.stride() == live_beta.stride() @pytest.mark.parametrize("quantization", [None, "fp8", "fp8_online", "modelopt_mixed"]) @@ -1862,10 +1867,13 @@ def __init__(self, **kwargs): layer.dt_bias = torch.empty(12) layer.o_norm = SimpleNamespace(weight=torch.empty(3)) layer.kv_cache = (torch.empty(1), torch.empty((2, 3, 4, 4))) - layer._b12x_kda_mixed_qkv = torch.empty((2, 9)) + layer._b12x_kda_max_tokens = 2 + layer._b12x_kda_mixed_qkv = torch.empty((2, 36)) layer._b12x_kda_raw_g = torch.empty((2, 3, 4)) beta_width = 15 if use_full_rank_gate else 12 beta_offset = 8 if use_full_rank_gate else 4 + layer._b12x_kda_beta_row_width = beta_width + layer._b12x_kda_beta_offset = beta_offset beta_storage = torch.empty((2, beta_width)) layer._b12x_kda_raw_beta = beta_storage.narrow(1, beta_offset, 3) layer._b12x_kda_z = torch.empty((2, 3, 4)) @@ -1944,6 +1952,74 @@ def run(self, binding, **kwargs): assert tensor.is_contiguous() +@pytest.mark.parametrize("benchmark", [False, True]) +def test_b12x_kda_decode_preparation_releases_synthetic_activations( + monkeypatch, benchmark +): + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + torch.nn.Module.__init__(layer) + layer.enable_b12x_kda_decode = True + layer.head_dim, layer.local_num_heads = 128, 2 + layer.local_projection_size = 256 + layer.use_full_rank_gate = True + layer.in_proj_qkvgfab = SimpleNamespace(output_size_per_partition=1154) + layer.model_config = SimpleNamespace(dtype=torch.bfloat16) + layer.num_spec = 3 + layer.gate_lower_bound = -5.0 + layer.A_log = torch.zeros(2) + layer.dt_bias = torch.zeros(2, 128) + layer.o_norm = SimpleNamespace(weight=torch.ones(128), eps=1e-6) + layer.kv_cache = (torch.empty(1), torch.zeros(8, 2, 128, 128)) + monkeypatch.setattr( + kimi_gdn_linear_attn, + "current_platform", + SimpleNamespace(current_device=lambda: "cpu", is_cuda=lambda: True), + ) + monkeypatch.setattr( + kimi_gdn_linear_attn, + "get_b12x_gdn_decode", + lambda: SimpleNamespace( + bind_kda=None, run_kda=None, is_supported=lambda _: True + ), + ) + monkeypatch.setattr( + KimiGatedDeltaNetAttention, + "get_state_dtype", + lambda self: (torch.bfloat16, torch.float32), + ) + layer._initialize_b12x_kda_decode( + SimpleNamespace(scheduler_config=SimpleNamespace(max_num_seqs=32)) + ) + assert sum(t.numel() * t.element_size() for t in layer.buffers()) < 1024 + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + class State: + layout = SimpleNamespace( + scratch_specs=lambda: [SimpleNamespace(shape=(128,), dtype=torch.uint8)] + ) + + def bind_kda(self, **tensors): + references.extend( + weakref.ref(tensors[name]) + for name in ("scratch", "mixed_qkv", "raw_g", "raw_beta", "z", "output") + ) + return SimpleNamespace(**tensors) + + def run(self, binding, **kwargs): + assert binding.mixed_qkv.shape == (128, 768) + assert binding.raw_beta.stride() == (1154, 1) + assert binding.num_tokens.item() == 1 + binding.output.copy_(binding.z) + + call = layer._b12x_kda_decode_call(State(), benchmark=benchmark) + assert call.owners == () + call.produce() + call.run() + assert all(reference() is not None for reference in references) + del call + assert all(reference() is None for reference in references) + + @pytest.mark.parametrize("benchmark", [False, True]) def test_b12x_kda_prefill_preparation_releases_synthetic_activations( monkeypatch, benchmark @@ -2089,9 +2165,13 @@ def test_b12x_kda_prefill_live_inputs_after_preparation(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("tokens", [3, 12]) -def test_glm_adaptive_kda_graph_matches_independent_request_states(monkeypatch, tokens): +def test_glm_adaptive_kda_graph_matches_independent_request_states( + monkeypatch, tokens, request +): from dataclasses import replace + from b12x.preparation import PreparationSession + from tests.v1.attention.test_gdn_metadata_builder import _create_gdn_builder from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata from vllm.models.glm5next.nvidia.kda import Glm5NextKDAMetadataBuilder @@ -2152,8 +2232,8 @@ def test_glm_adaptive_kda_graph_matches_independent_request_states(monkeypatch, layer._b12x_kda_num_accepted_tokens = torch.ones( requests, dtype=torch.int32, device=device ) - # The declaration takes its operand layouts from the layer's staging - # tensors, as the served layer declares it in bind_kv_cache. + # This fixture's live beta rows are contiguous; served projection views + # use the fused projection's row width instead. layer.model_config = SimpleNamespace(dtype=torch.bfloat16) layer.b12x_kda_null_state_index = 0 monkeypatch.setattr( @@ -2161,23 +2241,9 @@ def test_glm_adaptive_kda_graph_matches_independent_request_states(monkeypatch, "get_state_dtype", lambda self: (torch.bfloat16, torch.float32), ) - width = heads * dim - layer._b12x_kda_mixed_qkv = torch.zeros( - (tokens, 3 * width), dtype=torch.bfloat16, device=device - ) layer.use_full_rank_gate = True - layer._b12x_kda_raw_g = torch.zeros( - (tokens, heads, dim), dtype=torch.bfloat16, device=device - ) - layer._b12x_kda_raw_beta = torch.zeros( - (tokens, heads), dtype=torch.bfloat16, device=device - ) - layer._b12x_kda_z = torch.zeros( - (tokens, heads, dim), dtype=torch.bfloat16, device=device - ) - layer._b12x_kda_output = torch.zeros( - (tokens, heads, dim), dtype=torch.bfloat16, device=device - ) + layer._b12x_kda_beta_row_width = heads + layer._b12x_kda_beta_offset = 0 layer._b12x_kda_query_start_loc = torch.zeros( requests + 1, dtype=torch.int32, device=device ) @@ -2186,6 +2252,19 @@ def test_glm_adaptive_kda_graph_matches_independent_request_states(monkeypatch, ) plan = layer._b12x_kda_decode_declaration(33) layer._b12x_kda_plan = plan + session = PreparationSession(device=device, autotune=False) + request.addfinalizer(session.close) + prepared = session.prepare( + ( + plan.request( + name="glm-kda-decode-live-inputs", + prepare_call=layer._prepare_b12x_kda_decode, + benchmark_call=layer._benchmark_b12x_kda_decode, + ), + ) + ) + request.addfinalizer(prepared.close) + assert plan.prepared.owners == () (layer._b12x_kda_scratch,) = get_b12x_scratch_buffers(plan) context = SimpleNamespace( attn_metadata={layer.prefix: metadata}, additional_kwargs={} diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 33a6fd2b55f1..e4700dd312aa 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -611,48 +611,14 @@ def _initialize_b12x_kda_decode(self, vllm_config: VllmConfig) -> None: self._b12x_kda_max_tokens = max_tokens self._b12x_kda_max_seqs = max_seqs self._b12x_kda_state_index_columns = state_index_columns - width = self.local_num_heads * self.head_dim - self.register_buffer( - "_b12x_kda_mixed_qkv", - torch.zeros( - (max_tokens, 3 * width), dtype=self.model_config.dtype, device=device - ), - persistent=False, - ) # KDA consumes one forget-gate value per head dimension in either gate # mode. Keep beta's row stride equal to its fused-projection view. - raw_beta_row_width = self.in_proj_qkvgfab.output_size_per_partition - raw_beta_offset = ( + self._b12x_kda_beta_row_width = self.in_proj_qkvgfab.output_size_per_partition + self._b12x_kda_beta_offset = ( 4 * self.local_projection_size + self.head_dim if self.use_full_rank_gate else 3 * self.local_projection_size ) - self.register_buffer( - "_b12x_kda_raw_beta_storage", - torch.zeros( - (max_tokens, raw_beta_row_width), - dtype=self.model_config.dtype, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_b12x_kda_raw_beta", - self._b12x_kda_raw_beta_storage.narrow( - 1, raw_beta_offset, self.local_num_heads - ), - persistent=False, - ) - for name, shape in ( - ("_b12x_kda_raw_g", (max_tokens, self.local_num_heads, self.head_dim)), - ("_b12x_kda_z", (max_tokens, self.local_num_heads, self.head_dim)), - ("_b12x_kda_output", (max_tokens, self.local_num_heads, self.head_dim)), - ): - self.register_buffer( - name, - torch.zeros(shape, dtype=self.model_config.dtype, device=device), - persistent=False, - ) self.register_buffer( "_b12x_kda_query_start_loc", torch.zeros(max_seqs + 1, dtype=torch.int32, device=device), @@ -681,6 +647,25 @@ def _initialize_b12x_kda_decode(self, vllm_config: VllmConfig) -> None: persistent=False, ) + def _b12x_kda_decode_probes(self) -> tuple[torch.Tensor, ...]: + """Allocate call-local probes with the live decode operand layouts.""" + rows, heads, dim = ( + self._b12x_kda_max_tokens, + self.local_num_heads, + self.head_dim, + ) + dtype, device = self.model_config.dtype, self._b12x_kda_num_tokens.device + mixed_qkv = torch.empty((rows, 3 * heads * dim), dtype=dtype, device=device) + beta_storage = torch.empty( + (rows, self._b12x_kda_beta_row_width), dtype=dtype, device=device + ) + raw_beta = beta_storage.narrow(1, self._b12x_kda_beta_offset, heads) + raw_g, z, output = ( + torch.empty((rows, heads, dim), dtype=dtype, device=device) + for _ in range(3) + ) + return mixed_qkv, raw_g, raw_beta, z, output + def _b12x_kda_decode_declaration(self, max_state_slots: int): api = self._b12x_kda_api if api is None: @@ -701,14 +686,15 @@ def _b12x_kda_decode_declaration(self, max_state_slots: int): qk_l2norm=True, null_state_index=self.b12x_kda_null_state_index, ) + mixed_qkv, raw_g, raw_beta, z, output = self._b12x_kda_decode_probes() return api.plan( caps, invocation=api.invocation_from_tensors( caps, - mixed_qkv=self._b12x_kda_mixed_qkv, - raw_g=self._b12x_kda_raw_g, - raw_beta=self._b12x_kda_raw_beta, - z=self._b12x_kda_z, + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + z=z, A_log=self.A_log, dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), norm_weight=self.o_norm.weight, @@ -718,7 +704,7 @@ def _b12x_kda_decode_declaration(self, max_state_slots: int): state_indices=self._b12x_kda_state_indices, num_seqs=self._b12x_kda_num_seqs, num_tokens=self._b12x_kda_num_tokens, - output=self._b12x_kda_output, + output=output, ), ) @@ -939,12 +925,8 @@ def _b12x_kda_decode_call(self, state, *, benchmark: bool): # Trial and prepare factories own their scratch; the runtime binding # in _run_b12x_kda_decode_post_conv draws from the workspace manager. scratch = torch.empty(spec[0], dtype=spec[1], device=slots.device) + mixed_qkv, raw_g, raw_beta, z, output = self._b12x_kda_decode_probes() if benchmark: - mixed_qkv = self._b12x_trial_tensor(self._b12x_kda_mixed_qkv) - raw_g = self._b12x_trial_tensor(self._b12x_kda_raw_g) - raw_beta = self._b12x_trial_tensor(self._b12x_kda_raw_beta) - z = self._b12x_trial_tensor(self._b12x_kda_z) - output = self._b12x_trial_tensor(self._b12x_kda_output) query_start_loc = self._b12x_trial_tensor(self._b12x_kda_query_start_loc) accepted = self._b12x_trial_tensor(self._b12x_kda_num_accepted_tokens) state_indices = self._b12x_trial_tensor(self._b12x_kda_state_indices) @@ -952,13 +934,6 @@ def _b12x_kda_decode_call(self, state, *, benchmark: bool): num_tokens = self._b12x_trial_tensor(self._b12x_kda_num_tokens) saved_state = slots[slot : slot + 1].clone() else: - mixed_qkv, raw_g, raw_beta, z, output = ( - self._b12x_kda_mixed_qkv, - self._b12x_kda_raw_g, - self._b12x_kda_raw_beta, - self._b12x_kda_z, - self._b12x_kda_output, - ) query_start_loc = self._b12x_kda_query_start_loc accepted, state_indices = ( self._b12x_kda_num_accepted_tokens, @@ -1008,19 +983,6 @@ def reset(): produce=produce, reset=reset, restore=reset if saved_state is not None else None, - owners=( - scratch, - mixed_qkv, - raw_g, - raw_beta, - z, - output, - query_start_loc, - accepted, - state_indices, - num_seqs, - num_tokens, - ), ) def _prepare_b12x_kda_prefill(self, state): From 951279a0999d32c927b313b98d7e9bba34f4e8be Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 16:59:50 +0000 Subject: [PATCH 091/365] fix(cudagraph): reserve retained native initialization once Measure non-Torch memory across the V2 capture span and reconcile it after backend-plan cleanup. Exclude only overlap already charged as late persistent memory, bounded by the measured graph delta. Bootstrap-only and cleanup-only allocations, transient native graph storage, activation headroom, and graph coverage remain reserved. Validation: 98 worker/profile tests pass, including native CUDA graph and cudaMalloc coverage. Eight worker overlap cases fail against the parent implementation. Ruff, formatting and type hooks pass. The CUDA API lint finding is unchanged at gpu_worker.py:592 in the parent revision; skip only that hook. Whole-model TP2 admission remains a separate serving gate. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- ...gpu_model_runner_v2_cudagraph_profiling.py | 80 +++++++++++++++++-- tests/v1/worker/test_gpu_worker.py | 25 +++++- vllm/v1/worker/gpu/cudagraph_utils.py | 15 ++++ vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu_worker.py | 31 +++++-- 5 files changed, 134 insertions(+), 18 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index 99f57ce4b20c..33a27e9272a8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py @@ -141,6 +141,7 @@ def _fake_set_current_vllm_config(_cfg): # retained; default to a constant (nothing retained). monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None) monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr(cgu.torch.accelerator, "memory_reserved", lambda: 0) monkeypatch.setattr( cgu.torch.accelerator, "get_memory_info", lambda: (1 << 30, 1 << 30) ) @@ -214,6 +215,49 @@ def test_profile_cudagraph_memory_piecewise_only_returns_measured(monkeypatch): assert result == captured_bytes +@pytest.mark.parametrize( + "init_native,capture_native,retained_native", + [ + (0, 256, 256), + (512, 256, 64), + (0, 256, 0), + (512, 0, 0), + (0, 0, 256), + (0, 256, 512), + ], +) +def test_profile_cudagraph_memory_records_native_growth_inside_capture_only( + monkeypatch, init_native, capture_native, retained_native +): + """Worker reconciles capture growth after all prepared plans are released.""" + _patch_module(monkeypatch) + mib = 1 << 20 + runner = _make_profiling_runner( + CUDAGraphMode.FULL_AND_PIECEWISE, + piecewise_only=True, + captured_bytes=1000 * mib, + ) + + def memory_info(): + native = init_native + if "teardown" in runner.events: + native += retained_native + elif "capture" in runner.events: + native += capture_native + # Torch-owned memory is not native initialization and is fully reserved. + return ((4096 - native - 512) * mib, 4096 * mib) + + monkeypatch.setattr(cgu.torch.accelerator, "get_memory_info", memory_info) + monkeypatch.setattr(cgu.torch.accelerator, "memory_reserved", lambda: 512 * mib) + + assert cgu.profile_cudagraph_memory(runner) == 1000 * mib + assert runner.cudagraph_native_memory_profile == ( + init_native * mib, + (init_native + capture_native) * mib, + 1000 * mib, + ) + + def test_profile_cudagraph_memory_tears_down_on_capture_error(monkeypatch): _patch_module(monkeypatch) runner = _make_profiling_runner(CUDAGraphMode.FULL) @@ -367,6 +411,8 @@ def test_profile_cudagraph_memory_clears_captured_graphs(monkeypatch): # CUDA graph executables must be destroyed and synchronized before their # B12X channel checkpoints and tensor workspaces are released. assert lifecycle == [ + "synchronize", # Native-memory measurement before capture. + "synchronize", # Native-memory measurement after capture. "synchronize", "reset-piecewise", "reset-breakable", @@ -495,8 +541,10 @@ def _capture_model(*, profile_only: bool = False) -> int: @create_new_process_for_each_test("spawn") @pytest.mark.skipif(not cgu.current_platform.is_cuda(), reason="requires CUDA") -def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch): - """Profiling graph memory must be freed before teardown completes.""" +@pytest.mark.parametrize("native_bytes", [0, 256 << 20]) +def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch, native_bytes): + """Release profiling graphs and avoid reserving persistent cudaMalloc twice.""" + from cuda.bindings import runtime @contextlib.contextmanager def _fake_set_current_vllm_config(_cfg): @@ -513,6 +561,7 @@ def _fake_set_current_vllm_config(_cfg): runner.maybe_remove_all_loras = lambda _: None allocation_bytes = 64 << 20 memory: dict[str, int] = {} + native_pointer = None torch.accelerator.synchronize() gc.collect() @@ -533,6 +582,11 @@ def _init(r): ) def _capture_model(*, profile_only: bool = False) -> int: + nonlocal native_pointer + before_free = torch.accelerator.get_memory_info()[0] + if native_bytes: + error, native_pointer = runtime.cudaMalloc(native_bytes) + assert error == runtime.cudaError_t.cudaSuccess for owner in ( runner.cudagraph_manager, runner.speculator.cudagraph_manager, @@ -545,17 +599,27 @@ def _capture_model(*, profile_only: bool = False) -> int: owner.graph_capture_resources["profile"] = [output] torch.accelerator.synchronize() memory["captured"] = torch.accelerator.memory_reserved() - return memory["captured"] - memory["before"] + memory["measured"] = before_free - torch.accelerator.get_memory_info()[0] + return memory["measured"] monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config) monkeypatch.setattr(cgu, "_init_minimal_kv_cache_for_profiling", _init) runner.capture_model = _capture_model - cgu.profile_cudagraph_memory(runner) - memory["after"] = torch.accelerator.memory_reserved() - - assert memory["captured"] - memory["before"] >= 3 * allocation_bytes - assert memory["after"] == memory["before"] + try: + graph_estimate = cgu.profile_cudagraph_memory(runner) + memory["after"] = torch.accelerator.memory_reserved() + + assert memory["captured"] - memory["before"] >= 3 * allocation_bytes + assert memory["after"] == memory["before"] + native_before, native_after, measured = runner.cudagraph_native_memory_profile + assert measured == graph_estimate == memory["measured"] + assert native_after - native_before >= native_bytes + finally: + if native_pointer is not None: + assert ( + runtime.cudaFree(native_pointer)[0] == runtime.cudaError_t.cudaSuccess + ) def test_teardown_profiling_state_clears_mamba_align_metadata(monkeypatch): diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 2133d437d7ee..b8ef5dc843fa 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -311,12 +311,27 @@ def run_model(scheduler_output, intermediate_tensors): ) @pytest.mark.parametrize("graph_estimate", [0, 4]) @pytest.mark.parametrize("estimate_graphs", [False, True]) +@pytest.mark.parametrize( + "native_profile,final_native,overlap", + [ + (None, 0, 0), # Runners without native capture measurements retain the budget. + ((100, 103, 4), 103, 3), # Persistent capture initialization. + ((100, 103, 4), 101, 1), # Prepared-plan release frees part of the growth. + ((100, 103, 4), 100, 0), # All native capture storage was temporary. + ((103, 103, 4), 107, 0), # Bootstrap and cleanup growth are not capture cost. + ((100, 103, 1), 107, 1), # The measured capture delta bounds the overlap. + ((100, 103, 4), 107, 3), # Cleanup-only growth cannot increase the discount. + ], +) def test_cudagraph_memory_profile_prepares_and_releases_b12x_state( monkeypatch, final_free_memory, expected_available_memory, graph_estimate, estimate_graphs, + native_profile, + final_native, + overlap, ): """KV admission counts retained allocations, not released profiling blocks.""" events: list[object] = [] @@ -335,6 +350,7 @@ def profile_glm_dcp_attention(prepare_profile_state): events.append("profile_glm_dcp_attention") model_runner = SimpleNamespace( + cudagraph_native_memory_profile=native_profile, model_memory_usage=0, profile_run=profile_run, profile_glm_dcp_attention=profile_glm_dcp_attention, @@ -381,7 +397,9 @@ def fake_memory_profiling(*args, **kwargs): def final_snapshot(**_kwargs): assert events[-3:] == ["collect", "synchronize", "empty_cache"] events.append("final_snapshot") - return SimpleNamespace(free_memory=final_free_memory) + return SimpleNamespace( + free_memory=final_free_memory, non_torch_memory=final_native + ) monkeypatch.setattr(gpu_worker.gc, "collect", lambda: events.append("collect")) monkeypatch.setattr( @@ -427,6 +445,7 @@ def final_snapshot(**_kwargs): "empty_cache", "final_snapshot", ] + graph_estimate -= min(overlap, graph_estimate, 90 - final_free_memory) applied_graph_estimate = graph_estimate if estimate_graphs else 0 assert available == expected_available_memory - applied_graph_estimate assert worker.peak_activation_memory == 5 @@ -488,8 +507,8 @@ def test_post_capture_recommendation_counts_measured_graph_memory_once( monkeypatch.setattr(gpu_worker, name, lambda *args: None) monkeypatch.setattr("vllm.utils.jit_monitor.activate", lambda **kwargs: None) - worker._compile_or_warm_up_model_after_preparation = ( - lambda: gpu_worker.Worker._compile_or_warm_up_model_after_preparation(worker) + worker._compile_or_warm_up_model_after_preparation = lambda: ( + gpu_worker.Worker._compile_or_warm_up_model_after_preparation(worker) ) worker._b12x_session = None worker._get_cudagraph_capture_context = nullcontext diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 7f5069b47f08..da2549011d80 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1181,6 +1181,7 @@ def profile_cudagraph_memory( partition reclaims the storages of earlier cudagraph recordings once the real capture records new ones, leading to use-after-free crashes). """ + runner.cudagraph_native_memory_profile = None if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: return 0 @@ -1253,7 +1254,14 @@ def profile_cudagraph_memory( mem_samples: list[int] = [] manager._capture_mem_samples = mem_samples + native_before_capture = _non_torch_memory_bytes() measured = int(runner.capture_model(profile_only=True)) + native_after_capture = _non_torch_memory_bytes() + runner.cudagraph_native_memory_profile = ( + native_before_capture, + native_after_capture, + measured, + ) # The measured delta covers PIECEWISE, encoder and speculator graphs # plus the sampled FULL graphs; swap the sampled FULL cost for the @@ -1303,6 +1311,13 @@ def profile_cudagraph_memory( platform_cls._global_graph_pool = saved_global_pool +def _non_torch_memory_bytes() -> int: + """Measure device memory not owned by the Torch caching allocator.""" + torch.accelerator.synchronize() + free, total = torch.accelerator.get_memory_info() + return total - free - torch.accelerator.memory_reserved() + + def _extrapolate_full_graph_memory( mem_samples: list[int], graph_descs: list[BatchExecutionDescriptor], diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index cfb3cd98933e..242a18efed53 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -202,6 +202,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): + self.cudagraph_native_memory_profile: tuple[int, int, int] | None = None self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index bc7d8949cf07..fa308b65fe87 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -669,13 +669,6 @@ def determine_available_memory(self) -> int: finally: self._release_b12x_profile_state() - # Respect the opt-in flag as originally designed. - cudagraph_memory_estimate_applied = ( - cudagraph_memory_estimate - if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS - else 0 - ) - # Prepared plans can retain profiling tensors beyond graph teardown. # Reclaim their released allocator blocks before measuring persistence. gc.collect() @@ -692,6 +685,30 @@ def determine_available_memory(self) -> int: - final_profile_snapshot.free_memory, 0, ) + native_profile = getattr( + self.model_runner, "cudagraph_native_memory_profile", None + ) + if native_profile is not None: + native_before, native_after, measured = native_profile + # Count capture-initialized modules/communication storage once. + # Sample retention only after prepared plans have been released; + # bootstrap-only and cleanup-only growth do not discount graphs. + retained_native_capture = max( + min(native_after, final_profile_snapshot.non_torch_memory) + - native_before, + 0, + ) + cudagraph_memory_estimate -= min( + retained_native_capture, + late_persistent_memory, + max(measured, 0), + cudagraph_memory_estimate, + ) + cudagraph_memory_estimate_applied = ( + cudagraph_memory_estimate + if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS + else 0 + ) init_free_memory = self.init_snapshot.free_memory free_gpu_memory = final_profile_snapshot.free_memory rocm_fallback = maybe_rocm_profiling_fallback(profile_result) From 09e7a0b4aeaf0829ed06bb75b909efa8b5ae9652 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 17:35:57 +0000 Subject: [PATCH 092/365] fix(cache): reserve immutable boundary restore sources in KV capacity Count checkpoint source ownership and partial-tail copy-on-write alongside private endpoints when sizing request-boundary cache pools. Reject pools that admit a cold prompt but cannot restore its checkpoint; retain the existing aligned-policy budget and allocator invariants. Validation: 396 KV-cache and boundary tests pass, including the 62-block deadlock reproducer and 65-block successful restore. Five environment-dependent tests are rerun with their model configs and GPU visibility. Nine capacity cases reject the parent implementation. Full-model TP2 qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../core/test_boundary_admission_horizon.py | 80 ++++++++++++++++++- tests/v1/core/test_kv_cache_utils.py | 12 +-- vllm/v1/core/kv_cache_utils.py | 7 +- 3 files changed, 89 insertions(+), 10 deletions(-) diff --git a/tests/v1/core/test_boundary_admission_horizon.py b/tests/v1/core/test_boundary_admission_horizon.py index 18faf1c33b79..3a789f0679a8 100644 --- a/tests/v1/core/test_boundary_admission_horizon.py +++ b/tests/v1/core/test_boundary_admission_horizon.py @@ -2,18 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Per-reader attention horizons, low-pressure concurrency, and release fences.""" +from types import SimpleNamespace + import pytest import torch from tests.v1.core import test_boundary_admission as base from tests.v1.core.test_boundary_admission import initialize_hash as initialize_hash +from vllm.v1.core.kv_cache_utils import ( + _estimate_max_model_len_from_groups, + _pool_bytes_per_block, +) from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, MambaSpec, MLAAttentionSpec, ) -from vllm.v1.request import RequestStatus +from vllm.v1.request import Request, RequestStatus pytestmark = pytest.mark.cpu_test @@ -89,12 +95,82 @@ def restore(cache, req, lookahead=4, defer=True, pending=()): return allocated +@pytest.mark.parametrize("num_blocks", [62, 65, 80]) +def test_small_pool_capacity_covers_instruction_checkpoint_restore(num_blocks): + """Accepted pool geometry must admit both cold prompts and their restores.""" + attention = MLAAttentionSpec( + block_size=2048, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + model_version="glm5_next", + ) + recurrent = MambaSpec( + block_size=256, + shapes=((1,),), + dtypes=(torch.uint8,), + mamba_cache_mode="align", + num_speculative_blocks=3, + ) + groups = [KVCacheGroupSpec([str(i)], recurrent) for i in range(6)] + groups.append(KVCacheGroupSpec(["attention"], attention)) + config = SimpleNamespace( + use_request_boundary_checkpoints=True, + attention_config=SimpleNamespace(hisparse_config=None), + model_config=SimpleNamespace(max_model_len=1048576), + parallel_config=SimpleNamespace(decode_context_parallel_size=2), + cache_config=SimpleNamespace(mamba_cache_mode="align"), + ) + max_model_len = _estimate_max_model_len_from_groups( + config, groups, (num_blocks - 1) * _pool_bytes_per_block(groups) + ) + if num_blocks == 62: + # Cold admission fits, but restore needs 64 non-null blocks. Startup + # must reject this pool instead of leaving an idle scheduler waiting. + assert max_model_len == 0 + return + assert max_model_len >= 6357 + cache = base.make_kv_cache_manager( + KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ), + max_model_len=max_model_len, + max_in_flight_tokens=8192, + enable_caching=True, + use_eagle=True, + num_prefill_lookahead=1, + dcp_world_size=2, + scheduler_block_size=4096, + hash_block_size=256, + enable_boundary_checkpoints=True, + ) + producer = request("producer", length=6357, cap=1024) + producer.recurrent_instruction_boundary = 6345 + cache.get_computed_blocks(producer) + for count in (4096, 2249): + cache.new_step_starts() + assert cache.allocate_slots(producer, count, num_lookahead_tokens=3) + producer.num_computed_tokens += count + base.drain(cache) + checkpoint = cache.publish_boundary_checkpoint(producer, 6345, kind="instruction") + assert checkpoint is not None + cache.free(producer) + assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 + + reader = request("reader", length=6357, cap=1024) + reader.recurrent_instruction_boundary = 6345 + assert restore(cache, reader, lookahead=3) is not None + base.drain(cache) + cache.free(reader) + assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 + + def test_all_16_short_warm_readers_admit_without_checkpoint_loss(): cache = manager() for index in range(16): seed(cache, str(index)) published = set(cache.boundary_checkpoints._entries) - active = [] + active: list[Request] = [] for index in range(16): req = request("warm-" + str(index), str(index)) allocated = restore(cache, req) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 76fc5f2945c5..7d6957120049 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2669,7 +2669,7 @@ def test_balanced_glm_boundary_capacity_includes_private_endpoints( block_bytes = kv_cache_utils._pool_bytes_per_block(groups) live_bytes = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) live_blocks = live_bytes // block_bytes - private_blocks = 3 * (len(groups) + 1) + private_blocks = 4 * (len(groups) + 1) monkeypatch.setattr( VllmConfig, "use_request_boundary_checkpoints", property(lambda self: True) ) @@ -4277,9 +4277,9 @@ def test_auto_fit_max_model_len_with_hybrid(): @pytest.mark.parametrize("dcp", [1, 2, 4]) -@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (70, 15)]) +@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (62, 0), (70, 7)]) def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_pages): - """Live MTP state alone does not cover instruction/prompt/response copies.""" + """Pool sizing includes private endpoints and the immutable restore source.""" config = SimpleNamespace( use_request_boundary_checkpoints=True, attention_config=SimpleNamespace(hisparse_config=None), @@ -4304,8 +4304,8 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ groups = [KVCacheGroupSpec([f"state-{i}"], state) for i in range(6)] groups.append(KVCacheGroupSpec(["attention"], attention)) block_bytes = kv_cache_utils._pool_bytes_per_block(groups) - # Six groups require five live/scratch states each; three endpoints - # additionally own eight blocks each. The null block is unavailable. + # Six groups require five live/scratch states each; three private endpoints + # and a restore source own eight blocks each. The null block is unavailable. available = (num_blocks - 1) * block_bytes estimate = kv_cache_utils._estimate_max_model_len_from_groups( config, groups, available @@ -4329,7 +4329,7 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups ) assert get_max_concurrency_for_kv_cache_config(config, cache) == ( - num_blocks / (30 + 24 + expected_pages) + num_blocks / (30 + 32 + expected_pages) ) # An aligned-policy control has no endpoint reservation. Its capacity diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 567cab0e4628..85f5c062ec82 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1077,7 +1077,7 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) -> int: - """Count private instruction, prompt, and response checkpoint storage.""" + """Count private endpoints and the immutable checkpoint used by a restore.""" if not num_groups or not vllm_config.use_request_boundary_checkpoints: return 0 @@ -1085,7 +1085,10 @@ def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) - # KVCacheManager.allocate_slots reserves one block per group and one # auxiliary block per endpoint, separately from the live sequence state. - return NUM_BOUNDARY_CHECKPOINT_SLOTS * (num_groups + 1) + # Restoring a checkpoint pins its source until the reader finishes, while + # partial-tail CoW creates writable state. One additional bundle covers + # these source pages and auxiliary state alongside the private endpoints. + return (NUM_BOUNDARY_CHECKPOINT_SLOTS + 1) * (num_groups + 1) def get_max_concurrency_for_kv_cache_config( From c1ef43814da82f4489dacc8b5e6fb50eb2272f40 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 17:39:31 +0000 Subject: [PATCH 093/365] test(cache): cover asynchronous decode at minimum restore capacity Exercise 128 MTP decode steps after a boundary restore in the smallest accepted shared pool, with copy-release fences and two in-flight batches. Assert complete block ownership release after the request. Three capacity cases pass; serving source is unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../v1/core/test_boundary_admission_horizon.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/v1/core/test_boundary_admission_horizon.py b/tests/v1/core/test_boundary_admission_horizon.py index 3a789f0679a8..b8a3b099fffa 100644 --- a/tests/v1/core/test_boundary_admission_horizon.py +++ b/tests/v1/core/test_boundary_admission_horizon.py @@ -10,6 +10,7 @@ from tests.v1.core import test_boundary_admission as base from tests.v1.core.test_boundary_admission import initialize_hash as initialize_hash from vllm.v1.core.kv_cache_utils import ( + KVCacheBlock, _estimate_max_model_len_from_groups, _pool_bytes_per_block, ) @@ -161,7 +162,24 @@ def test_small_pool_capacity_covers_instruction_checkpoint_restore(num_blocks): reader.recurrent_instruction_boundary = 6345 assert restore(cache, reader, lookahead=3) is not None base.drain(cache) + reader.append_output_token_ids([10000]) + pending = False + held: list[KVCacheBlock] = [] + for step in range(128): + cache.new_step_starts() + assert cache.allocate_slots(reader, 4, num_lookahead_tokens=3) is not None + reader.num_computed_tokens += 4 + reader.num_in_flight_tokens += 4 + copies = cache.take_kv_cache_block_copies()[1] + if pending: + reader.num_in_flight_tokens -= 4 + reader.append_output_token_ids([10001 + step] * 4) + cache.block_pool.free_blocks(held) + held = copies + pending = True + cache.cache_blocks(reader, min(reader.num_computed_tokens, reader.num_tokens)) cache.free(reader) + cache.block_pool.free_blocks(held) assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 From 5b39ccfb0495c1df56d1cd5855c18366f52972e8 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 18:37:42 +0000 Subject: [PATCH 094/365] perf(kda): reuse bounded prefill scratch for convolution outputs Borrow only the available workspace suffix after the exact B12X scratch and output layout. Keep mixed-batch selections disjoint, reject overlapping convolution outputs, and release consumed projection owners before the output GEMM. The allocating paths and arithmetic are unchanged; image and video admission are not restricted. Adapted from the GLM memory implementation in 7f4aecc66e857d093e012a2c57d03711bd23be4a by Martin Vit and Codex. Validation: 19 CPU/CUDA tests pass, including exact output/history parity and repeated graph replay with changed inputs; the parent fails seven selected ownership/interface controls. All applicable pre-commit hooks pass. Source-built full-model serving qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../kernels/mamba/test_kda_conv_workspace.py | 186 ++++++++++++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 100 +++++++++- .../layers/mamba/ops/causal_conv1d.py | 45 ++++- 3 files changed, 321 insertions(+), 10 deletions(-) create mode 100644 tests/kernels/mamba/test_kda_conv_workspace.py diff --git a/tests/kernels/mamba/test_kda_conv_workspace.py b/tests/kernels/mamba/test_kda_conv_workspace.py new file mode 100644 index 000000000000..7dbae25c39ac --- /dev/null +++ b/tests/kernels/mamba/test_kda_conv_workspace.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Convolution outputs occupy workspace disjoint from the KDA consumer.""" + +import weakref +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from vllm.model_executor.layers.mamba.gdn import kimi_gdn_linear_attn as kda +from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn +from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata +from vllm.v1.worker.workspace import WorkspaceManager + + +def _layer(): + layer: Any = object.__new__(kda.KimiGatedDeltaNetAttention) + torch.nn.Module.__init__(layer) + layer.kda_prefill_backend = "b12x" + layer.get_state_dtype = lambda: (torch.bfloat16, torch.float32) + layer._b12x_prefill_plan = SimpleNamespace( + scratch_specs=lambda: [SimpleNamespace(shape=(1024,), dtype=torch.uint8)] + ) + layer._b12x_prefill_max_tokens = 32 + layer.local_num_heads = 2 + layer.head_dim = 128 + layer.local_projection_size = 256 + layer.model_config = SimpleNamespace(dtype=torch.bfloat16) + return layer + + +@pytest.mark.parametrize("capacity", [1024, 100000]) +def test_convolution_suffix_preserves_kda_workspace(monkeypatch, capacity): + manager = WorkspaceManager(torch.device("cpu")) + manager.reserve_all(((capacity,), torch.uint8)) + manager.lock() + monkeypatch.setattr(kda, "current_workspace_manager", lambda: manager) + layer = _layer() + outputs = layer._borrow_b12x_prefill_transients(32, torch.bfloat16) + if capacity == 1024: + assert outputs is None + return + for index, output in enumerate(outputs): + output.fill_(index + 1) + scratch, target = layer._get_b12x_prefill_workspace() + scratch.fill_(255) + target.fill_(17) + for index, output in enumerate(outputs): + assert output.shape == (32, 256) and output.is_contiguous() + assert torch.all(output == index + 1) + + +@pytest.mark.parametrize( + "backend,dtype", [("flashkda", torch.bfloat16), ("b12x", torch.float32)] +) +def test_convolution_suffix_keeps_unsupported_paths(backend, dtype): + layer = _layer() + layer.kda_prefill_backend = backend + assert layer._borrow_b12x_prefill_transients(32, dtype) is None + + +def test_mixed_selection_survives_convolution_and_kda_writes(monkeypatch): + manager = WorkspaceManager(torch.device("cpu")) + manager.reserve_all(((200000,), torch.uint8)) + manager.lock() + monkeypatch.setattr(kda, "current_workspace_manager", lambda: manager) + layer = _layer() + source = torch.randn(32, 768, dtype=torch.bfloat16) + indices = torch.tensor([21, 2, 4, 0, 31, 17, 8, 7], dtype=torch.int64) + specs = (((8, 768), source.dtype),) + outputs = layer._borrow_b12x_prefill_transients(8, source.dtype, specs) + selected = outputs[3] + torch.index_select(source, 0, indices, out=selected) + conv = layer._borrow_b12x_prefill_transients(8, source.dtype) + for buffer in conv: + buffer.fill_(5) + scratch, target = layer._get_b12x_prefill_workspace() + scratch.fill_(255) + target.fill_(9) + torch.testing.assert_close(selected, source[indices], rtol=0, atol=0) + + +@pytest.mark.parametrize("full_rank_gate", [False, True]) +def test_consumed_kda_projections_release_before_output_gemm(full_rank_gate): + layer = _layer() + layer.local_num_heads = 2 + layer.head_dim = 4 + layer.local_projection_size = 8 + layer.use_full_rank_gate = full_rank_gate + layer.in_proj_padding = 2 if full_rank_gate else 0 + owners = [] + + def project(x, columns): + result = torch.ones(x.shape[0], columns) + owners.append(weakref.ref(result)) + return result, None + + layer.in_proj_qkvgfab = lambda x: project(x, 40 if full_rank_gate else 30) + layer.g_a_proj = lambda x: project(x, 4) + layer.g_b_proj = lambda x: project(x, 8) + layer.f_b_proj = lambda x: project(x, 8) + + def recurrence(*, mixed_qkv, g1, g2, beta, core_attn_out): + core_attn_out.copy_( + mixed_qkv[:, :8].view(1, -1, 2, 4) + + g1 + + g2.unsqueeze(0) + + beta.unsqueeze(-1) + ) + + layer._forward = recurrence + + def output_projection(x): + assert all(owner() is None for owner in owners) + return x * 2, None + + layer.o_proj = output_projection + output = torch.empty(17, 8) + layer.forward(torch.ones(17, 8), torch.arange(17), output) + torch.testing.assert_close(output, torch.full_like(output, 8), rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", [33, 2048, 3072]) +@pytest.mark.parametrize("speculative_slots", [0, 3]) +@pytest.mark.parametrize("capture", [False, True]) +def test_caller_owned_convolution_preserves_output_and_history( + rows, speculative_slots, capture +): + channels, width = 4096, 4 + storage = torch.empty(rows * 4 * channels, device="cuda", dtype=torch.bfloat16) + source = storage[rows * channels :].view(rows, 3 * channels) + source.normal_() + original = source.clone() + weight = torch.randn(channels, width, device="cuda", dtype=source.dtype) + starts = torch.tensor([0, rows // 2, rows], dtype=torch.int32, device="cuda") + slots = torch.tensor([1, 2], dtype=torch.int32, device="cuda") + initial = torch.tensor([False, True], device="cuda") + history = torch.randn( + 4, channels, width - 1 + speculative_slots, device="cuda", dtype=source.dtype + ) + expected_history = history.clone() + supplied = storage[: rows * channels].view(rows, channels) + value = source[:, channels : 2 * channels].transpose(0, 1) + nums_dict, batch_ptr, token_chunk_offset_ptr = compute_causal_conv1d_metadata( + starts.cpu(), device=starts.device + ) + kwargs = dict( + query_start_loc=starts, + cache_indices=slots, + has_initial_state=initial, + activation="silu", + metadata=SimpleNamespace( + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + ), + ) + expected = causal_conv1d_fn(value, weight, None, expected_history, **kwargs) + actual = causal_conv1d_fn( + value, weight, None, history, out=supplied.transpose(0, 1), **kwargs + ) + assert actual.data_ptr() == supplied.data_ptr() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(history, expected_history, rtol=0, atol=0) + torch.testing.assert_close(source, original, rtol=0, atol=0) + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + causal_conv1d_fn( + value, weight, None, history, out=supplied.transpose(0, 1), **kwargs + ) + for _ in range(3): + source.normal_() + original.copy_(source) + history.normal_() + expected_history.copy_(history) + expected = causal_conv1d_fn(value, weight, None, expected_history, **kwargs) + graph.replay() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(history, expected_history, rtol=0, atol=0) + torch.testing.assert_close(source, original, rtol=0, atol=0) + with pytest.raises(ValueError, match="disjoint"): + causal_conv1d_fn(value, weight, None, history, out=value, **kwargs) diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index e4700dd312aa..d17d5d61f9a1 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import replace +from math import prod from typing import TYPE_CHECKING, Any import torch @@ -820,14 +821,14 @@ def _get_b12x_kda_workspace(self) -> torch.Tensor: (scratch,) = get_b12x_scratch_buffers(plan) return scratch - def _get_b12x_prefill_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: + def _b12x_prefill_workspace_specs(self): plan = self._b12x_prefill_plan if plan is None: raise PreparationResourceUnavailableError( "b12x KDA prefill is not prepared for the current KV generation" ) (spec,) = plan.scratch_specs() - scratch, output = current_workspace_manager().get_simultaneous( + return ( (spec.shape, spec.dtype), ( ( @@ -838,8 +839,40 @@ def _get_b12x_prefill_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: self.model_config.dtype, ), ) + + def _get_b12x_prefill_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: + scratch, output = current_workspace_manager().get_simultaneous( + *self._b12x_prefill_workspace_specs() + ) return scratch, output + def _borrow_b12x_prefill_transients( + self, + tokens: int, + dtype: torch.dtype, + selection_specs: tuple[tuple[tuple[int, ...], torch.dtype], ...] = (), + ) -> list[torch.Tensor] | None: + """Borrow existing storage disjoint from KDA scratch and output. + + Convolution Q/K/V and mixed-batch selections remain live while KDA + executes. Only the spare suffix after its exact planned buffers may + hold these tensors; allocation and other prefill backends are unchanged. + """ + if self.kda_prefill_backend != "b12x" or dtype != self.get_state_dtype()[0]: + return None + specs = ( + self._b12x_prefill_workspace_specs() + + (((tokens, self.local_projection_size), dtype),) * 3 + + selection_specs + ) + required = sum( + ((prod(shape) * dt.itemsize + 255) // 256) * 256 for shape, dt in specs + ) + manager = current_workspace_manager() + if required > manager.available_bytes(): + return None + return manager.get_simultaneous(*specs)[2:] + def get_b12x_preparation_units( self, layer: torch.nn.Module, @@ -1402,6 +1435,11 @@ def forward( beta=beta, core_attn_out=core_attn_out, ) + # Recurrence has consumed these views. Do not retain their backing + # projections while the output GEMM and collective allocate results. + del mixed_qkv, g1, g2, beta, g_proj_states, f_a, projected_qkvgfab + if self.use_full_rank_gate: + del projected core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") output[:] = self.o_proj(core_attn_out)[0] @@ -1497,9 +1535,34 @@ def _forward( mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) g1_spec = g1.index_select(1, spec_token_indx) beta_spec = beta.index_select(1, spec_token_indx) - mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) - g1_ns = g1.index_select(1, non_spec_token_indx) - beta_ns = beta.index_select(1, non_spec_token_indx) + assert non_spec_token_indx is not None + n = non_spec_token_indx.numel() + selection = ( + self._borrow_b12x_prefill_transients( + n, + mixed_qkv.dtype, + ( + ((n, mixed_qkv.shape[1]), mixed_qkv.dtype), + ((g1.shape[0], n, *g1.shape[2:]), g1.dtype), + ((beta.shape[0], n, *beta.shape[2:]), beta.dtype), + ), + ) + if m.num_prefills > 0 + else None + ) + if selection is None: + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + else: + # Mixed speculative recurrence uses the allocating path; + # the workspace suffix stays live until prefill consumes it. + mixed_qkv_ns, g1_ns, beta_ns = selection[3:] + torch.index_select( + mixed_qkv, 0, non_spec_token_indx, out=mixed_qkv_ns + ) + torch.index_select(g1, 1, non_spec_token_indx, out=g1_ns) + torch.index_select(beta, 1, non_spec_token_indx, out=beta_ns) g2_spec = g2_ns = None else: mixed_qkv_spec = g1_spec = beta_spec = None @@ -1586,6 +1649,9 @@ def _forward( self.local_projection_size, dim=-1 ) prefill_mixed_qkv = mixed_qkv_ns + conv_outputs = self._borrow_b12x_prefill_transients( + mixed_qkv_ns.shape[0], mixed_qkv_ns.dtype + ) # Packed prefill conv would require copying V solely to make # it dense for KDA. Separate calls accept the strided inputs @@ -1596,7 +1662,9 @@ def _prefill_conv( x: torch.Tensor, state: torch.Tensor, weight: torch.Tensor, + out: torch.Tensor | None, ) -> torch.Tensor: + output_arg = {"out": out.transpose(0, 1)} if out is not None else {} return causal_conv1d_fn( x.transpose(0, 1), weight, @@ -1607,11 +1675,27 @@ def _prefill_conv( cache_indices=non_spec_state_indices_tensor, query_start_loc=non_spec_query_start_loc, metadata=m, + **output_arg, ).transpose(0, 1) - q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) - k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) - v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns = _prefill_conv( + q_ns, + q_conv_state, + q_conv_weight, + conv_outputs[0] if conv_outputs is not None else None, + ) + k_ns = _prefill_conv( + k_ns, + k_conv_state, + k_conv_weight, + conv_outputs[1] if conv_outputs is not None else None, + ) + v_ns = _prefill_conv( + v_ns, + v_conv_state, + v_conv_weight, + conv_outputs[2] if conv_outputs is not None else None, + ) q_ns, k_ns, v_ns = ( rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) for x in (q_ns, k_ns, v_ns) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f7776aada6f1..2757ec9d0efd 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -491,6 +491,27 @@ def _causal_conv1d_fwd_kernel( # continuous batching tl.store(o_ptrs, acc, mask=mask_1d) +def _byte_ranges_overlap(a: torch.Tensor, b: torch.Tensor) -> bool: + """Conservatively test strided spans, allowing disjoint workspace views.""" + if not a.numel() or not b.numel(): + return False + if a.untyped_storage().data_ptr() != b.untyped_storage().data_ptr(): + return False + + def span(tensor): + start = tensor.storage_offset() * tensor.element_size() + stop = ( + start + + (1 + sum((n - 1) * s for n, s in zip(tensor.shape, tensor.stride()))) + * tensor.element_size() + ) + return start, stop + + a_start, a_stop = span(a) + b_start, b_stop = span(b) + return a_start < b_stop and b_start < a_stop + + def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, @@ -509,6 +530,7 @@ def causal_conv1d_fn( block_size_to_align=0, metadata=None, validate_data=False, + out: torch.Tensor | None = None, ): """support varlen + continuous batching when x is 2D tensor @@ -559,7 +581,9 @@ def causal_conv1d_fn( The number of tokens already completed for each sequence block_size_to_align: int The block size to align the cached states to - out: same shape as `x` + out: Optional caller-owned output with the input shape, dtype and device. + Storage must be disjoint from inputs and state. Input and state dtypes + must match so the result does not require an output cast. """ if isinstance(activation, bool) and activation: activation = "silu" @@ -568,7 +592,24 @@ def causal_conv1d_fn( # Store original dtype to cast back at the end original_x_dtype = x.dtype x = x.to(conv_states.dtype) - out = torch.empty_like(x) + if out is None: + out = torch.empty_like(x) + elif ( + out.shape != x.shape + or out.dtype != x.dtype + or out.dtype != original_x_dtype + or out.device != x.device + ): + raise ValueError( + "Convolution output must match input/state shape, dtype and device" + ) + elif any( + source is not None and _byte_ranges_overlap(out, source) + for source in (x, weight, bias, conv_states) + ): + raise ValueError( + "Convolution output storage must be disjoint from inputs and state" + ) if metadata is not None: nums_dict = metadata.nums_dict args = nums_dict From 4aeec44d9765782dd012ddeed266903b12ea878c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 19:27:41 +0000 Subject: [PATCH 095/365] perf(kda): reuse consumed prefill gates and return GLM projection storage Prefill merges recurrence results into consumed gate storage, while decode retains a separate output buffer. GLM returns the output projection without allocating a duplicate hidden-state tensor; both projection schedules release consumed owners before the output GEMM. Mixed batches write directly into the merge destination. Kernels, precision and multimodal admission are unchanged. Validation: 45 focused CPU/CUDA tests pass, covering ownership, mixed decode/prefill, exact convolution history and changed-input CUDA graph replay. Full TP2 serving qualification remains pending. Mixed-output reuse also preserves the attributed intent of 7f4aecc66e857d093e012a2c57d03711bd23be4a. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .../kernels/mamba/test_kda_conv_workspace.py | 23 ++++++++++ tests/models/test_glm5next_kda_gate_stream.py | 46 +++++++++++++++++++ tests/models/test_glm5next_model.py | 28 ++++++++--- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 37 +++++++++------ vllm/models/glm5next/nvidia/kda.py | 24 ++++------ 5 files changed, 124 insertions(+), 34 deletions(-) diff --git a/tests/kernels/mamba/test_kda_conv_workspace.py b/tests/kernels/mamba/test_kda_conv_workspace.py index 7dbae25c39ac..1cc6f2312b55 100644 --- a/tests/kernels/mamba/test_kda_conv_workspace.py +++ b/tests/kernels/mamba/test_kda_conv_workspace.py @@ -15,10 +15,18 @@ from vllm.v1.worker.workspace import WorkspaceManager +@pytest.fixture(autouse=True) +def forward_metadata(monkeypatch): + context = SimpleNamespace(attn_metadata=None) + monkeypatch.setattr(kda, "get_forward_context", lambda: context) + return context + + def _layer(): layer: Any = object.__new__(kda.KimiGatedDeltaNetAttention) torch.nn.Module.__init__(layer) layer.kda_prefill_backend = "b12x" + layer.prefix = "attention" layer.get_state_dtype = lambda: (torch.bfloat16, torch.float32) layer._b12x_prefill_plan = SimpleNamespace( scratch_specs=lambda: [SimpleNamespace(shape=(1024,), dtype=torch.uint8)] @@ -31,6 +39,21 @@ def _layer(): return layer +@pytest.mark.parametrize("prefills", [None, 0, 1]) +def test_gate_storage_is_reused_only_after_prefill_consumers( + forward_metadata, prefills +): + layer = _layer() + gate = torch.ones(1, 32, 2, 128, dtype=torch.bfloat16) + if prefills is not None: + metadata = object.__new__(kda.GDNAttentionMetadata) + metadata.num_prefills = prefills + forward_metadata.attn_metadata = {layer.prefix: metadata} + output = layer._core_attn_buffer(gate) + assert output.shape == gate.shape and output.dtype == gate.dtype + assert (output.data_ptr() == gate.data_ptr()) == (prefills == 1) + + @pytest.mark.parametrize("capacity", [1024, 100000]) def test_convolution_suffix_preserves_kda_workspace(monkeypatch, capacity): manager = WorkspaceManager(torch.device("cpu")) diff --git a/tests/models/test_glm5next_kda_gate_stream.py b/tests/models/test_glm5next_kda_gate_stream.py index 3711f9597035..59882a23be0e 100644 --- a/tests/models/test_glm5next_kda_gate_stream.py +++ b/tests/models/test_glm5next_kda_gate_stream.py @@ -3,10 +3,14 @@ """The GLM-5.3 KDA gate side stream: same result as the sequential forward, gate projections off the main stream, graph-capturable.""" +import weakref +from types import SimpleNamespace + import pytest import torch from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.forward_context import override_forward_context from vllm.models.glm5next.nvidia import kda as kda_module from vllm.models.glm5next.nvidia.kda import Glm5NextLinearAttention from vllm.v1.worker import workspace @@ -15,6 +19,12 @@ PROJ = HEADS * HEAD_DIM +@pytest.fixture(autouse=True) +def forward_context(): + with override_forward_context(SimpleNamespace(attn_metadata=None)): + yield + + class _Linear: """Stands in for a vLLM linear layer: returns (y, None) and records the CUDA stream it ran on.""" @@ -217,3 +227,39 @@ def __call__(self, x): graph.replay() torch.accelerator.synchronize(device) torch.testing.assert_close(captured, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("overlap", [False, True]) +def test_projection_owners_release_before_glm_output_gemm(monkeypatch, overlap): + """GLM returns the output projection directly after releasing inputs.""" + from vllm.compilation import monitor + + device = torch.device("cuda", 0) + layer = _make_layer(device, [], torch.Generator().manual_seed(27)) + owners = [] + + class TrackedLinear(_Linear): + def __call__(self, x): + result, bias = super().__call__(x) + owners.append(weakref.ref(result)) + return result, bias + + for name in ("g_a_proj", "g_b_proj", "in_proj_qkvgfab", "f_b_proj"): + linear = getattr(layer, name) + setattr(layer, name, TrackedLinear(linear.name, linear.weight, linear.log)) + original = layer.o_proj + outputs = [] + + def output_projection(x): + assert all(owner() is None for owner in owners) + result = original(x) + outputs.append(result[0]) + return result + + layer.o_proj = output_projection + monkeypatch.setattr(kda_module, "_GATE_SIDE_STREAM", overlap) + monkeypatch.setattr(monitor, "is_cudagraph_capturing_enabled", lambda: False) + x = torch.randn(64, HIDDEN, device=device, dtype=torch.bfloat16) + actual = _run(layer, x) + assert actual is outputs[0] diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 790607dbdc20..de03a5fa3b78 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -709,21 +709,26 @@ def named_parameters(self): assert actual_shard_id == shard_id -def test_glm5next_kda_adapts_shared_out_buffer_forward(monkeypatch) -> None: +def test_glm5next_kda_returns_shared_projection_without_a_copy(monkeypatch) -> None: layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) torch.nn.Module.__init__(layer) layer.use_full_rank_gate = True hidden_states = torch.randn(2, 4) positions = torch.arange(2) - def fake_forward(self, hidden_states, positions, output) -> None: - output.copy_(hidden_states + positions[:, None]) + expected = hidden_states + positions[:, None] - monkeypatch.setattr(KimiGatedDeltaNetAttention, "forward", fake_forward) + def fake_forward(self, hidden_states, positions) -> torch.Tensor: + return expected + + monkeypatch.setattr( + KimiGatedDeltaNetAttention, "_forward_projections", fake_forward + ) actual = layer(hidden_states, positions) torch.testing.assert_close(actual, hidden_states + positions[:, None]) + assert actual is expected def test_glm5next_moe_applies_external_gate_once() -> None: @@ -857,8 +862,9 @@ def run(self, binding, *, lower_bound, max_live_tokens, max_live_seqs): @pytest.mark.parametrize("prefill_backend", ["triton", "flashkda", "b12x"]) +@pytest.mark.parametrize("reuse_gate", [False, True]) def test_glm5next_kda_splits_mixed_decode_prefill_batch( - monkeypatch, prefill_backend: str + monkeypatch, prefill_backend: str, reuse_gate: bool ) -> None: from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata @@ -903,6 +909,7 @@ def test_glm5next_kda_splits_mixed_decode_prefill_batch( layer._b12x_kda_plan = None layer._b12x_kda_scratch = None layer.model_config = SimpleNamespace(dtype=torch.float32) + layer.get_state_dtype = lambda: (torch.float32, torch.float32) layer._b12x_prefill_api = None layer._b12x_prefill_plan = None layer._b12x_prefill_max_tokens = 4 @@ -968,6 +975,9 @@ class FakeWorkspaceManager: def __init__(self) -> None: self.storage = torch.empty(1024, dtype=torch.uint8) + def available_bytes(self): + return self.storage.numel() + def get_simultaneous(self, *specs): outputs = [] offset = 0 @@ -1003,10 +1013,14 @@ def fake_gather(state, indices, has_initial_state): monkeypatch.setattr(kda_ops, "fused_recurrent_kda", fake_recurrent) monkeypatch.setattr(kda_ops, "chunk_kda_with_fused_gate", fake_chunk) - core_attn_out = torch.empty(1, 4, 1, 1) + gate = torch.ones(1, 4, 1, 1) + core_attn_out = ( + layer._core_attn_buffer(gate) if reuse_gate else torch.empty_like(gate) + ) + assert (core_attn_out.data_ptr() == gate.data_ptr()) == reuse_gate layer._forward( mixed_qkv=torch.arange(12, dtype=torch.float32).view(4, 3), - g1=torch.ones(1, 4, 1, 1), + g1=gate, g2=torch.ones(4, 1, 1), beta=torch.ones(1, 4, 1), core_attn_out=core_attn_out, diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index d17d5d61f9a1..b493d19baaf0 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1387,7 +1387,24 @@ def forward( positions: torch.Tensor, output: torch.Tensor, ) -> None: - num_tokens = hidden_states.size(0) + output.copy_(self._forward_projections(hidden_states, positions)) + + def _core_attn_buffer(self, gate: torch.Tensor) -> torch.Tensor: + """Reuse a consumed gate only when recurrence has separate outputs. + + Prefill and mixed batches finish every gate read, including checkpoint + writes, before merging recurrence outputs. Decode kernels can instead + write directly into the supplied output while reading the gate. + """ + metadata = get_forward_context().attn_metadata + m = metadata.get(self.prefix) if isinstance(metadata, dict) else None + if isinstance(m, GDNAttentionMetadata) and m.num_prefills > 0: + return gate + return torch.empty_like(gate) + + def _forward_projections( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] # Optional model-installed callback (e.g. GLM-5.3 L2 weight prefetch of # o_proj while the small projections and the recurrence run). @@ -1422,11 +1439,7 @@ def forward( g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) - core_attn_out = torch.empty( - (1, num_tokens, self.local_num_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + core_attn_out = self._core_attn_buffer(g1) self._forward( mixed_qkv=mixed_qkv, @@ -1441,7 +1454,7 @@ def forward( if self.use_full_rank_gate: del projected core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( @@ -1972,14 +1985,12 @@ def _prefill_conv( # ---------- merge spec and non-spec outputs ---------- if core_attn_out_spec is not None and core_attn_out_non_spec is not None: # Mixed batches require indexed placement in the original order. - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_spec.dtype, - device=core_attn_out_spec.device, - ) + # Both recurrence results are separate allocations in mixed mode. + # All gate reads are complete, so their consumed storage may be + # the destination without an additional full-token merge buffer. + merged = core_attn_out[:, :num_actual_tokens] merged.index_copy_(1, spec_token_indx, core_attn_out_spec) merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens] elif core_attn_out_non_spec is not None: core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ 0, :num_actual_tokens diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 938136ab242b..780c170ae51b 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -140,7 +140,6 @@ def forward( # type: ignore[override] hidden_states: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: - output = torch.empty_like(hidden_states) if _GATE_SIDE_STREAM and not self.use_full_rank_gate: # Reserve both branches even in serialized graph warmup, before # workspace growth is forbidden by capture. @@ -148,12 +147,10 @@ def forward( # type: ignore[override] hidden_states.size(0) ) if _gate_overlap_allowed(): - self._forward_gate_overlap( - hidden_states, output, gate_workspace, main_workspace + return self._forward_gate_overlap( + hidden_states, gate_workspace, main_workspace ) - return output - super().forward(hidden_states, positions, output) - return output + return self._forward_projections(hidden_states, positions) def _projection_workspaces( self, num_tokens: int @@ -176,10 +173,9 @@ def size(projection) -> int: def _forward_gate_overlap( self, hidden_states: torch.Tensor, - output: torch.Tensor, gate_workspace: torch.Tensor | None, main_workspace: torch.Tensor | None, - ) -> None: + ) -> torch.Tensor: num_tokens = hidden_states.size(0) main = torch.cuda.current_stream(hidden_states.device) side = _gate_stream(hidden_states.device) @@ -222,11 +218,7 @@ def _forward_gate_overlap( g_proj_states.record_stream(main) g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) - core_attn_out = torch.empty( - (1, num_tokens, self.local_num_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + core_attn_out = self._core_attn_buffer(g1) self._forward( mixed_qkv=mixed_qkv, g1=g1, @@ -234,8 +226,12 @@ def _forward_gate_overlap( beta=beta, core_attn_out=core_attn_out, ) + # The main stream has consumed both branches. Release every projection + # owner before the output GEMM; stream records still protect pending + # device reads when Python releases the last reference. + del mixed_qkv, g1, g2, beta, g_proj_states, f_a, projected_qkvgfab, g_a_states core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] __all__ = ["Glm5NextLinearAttention"] From d40c86fe8c3df2c3699a33605be9766be8b2e802 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 19:51:10 +0000 Subject: [PATCH 096/365] fix(sparse-mla): release preparation probes and unused prefill storage Declare cache-writer formats without token-sized CUDA allocations. Keep execution probes and cache-page backups inside synchronized preparation callbacks, not published plan owners. Generic prefill workspace is allocated only when a generic prefill backend consumes it. Cache precision, kernel selection and live-input graph bindings are unchanged. Validation: eight focused ownership/workspace cases pass, including native FP8/NVFP4 writers and changed-input CUDA graph replay. Sparse suite: 150 pass; six pre-existing query-discovery/indexer failures remain, with query-discovery identity failures confirmed on parent 4aeec44. Six memory negative controls fail on that parent. Six existing mypy findings remain; other applicable hooks pass. Full-model TP2/batch3072 qualification is pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 120 ++++++++++++++++++ .../layers/attention/sparse_mla_attention.py | 12 +- .../attention/backends/mla/b12x_mla_sparse.py | 41 +++--- 3 files changed, 154 insertions(+), 19 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 7850c95b4321..a2951b9c52be 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Behavior checks for the B12x sparse MLA adapters.""" +import gc import weakref from types import SimpleNamespace from typing import Any @@ -71,6 +72,68 @@ def get_simultaneous(self, *shapes_and_dtypes): return [torch.empty(shape, dtype=dtype) for shape, dtype in shapes_and_dtypes] +@pytest.mark.parametrize("has_prefill_backend", [False, True]) +@pytest.mark.parametrize("dcp", [1, 2]) +def test_sparse_chunked_workspace_requires_generic_prefill( + monkeypatch, has_prefill_backend, dcp +): + from vllm.model_executor.layers.attention import sparse_mla_attention as common + from vllm.v1.attention.ops.dcp import MLADCPManager + + gathers = [] + dcp_manager = object.__new__(MLADCPManager) + dcp_manager.init_kv_gather = lambda buf, rows: gathers.append((buf, rows)) + prefill = SimpleNamespace(clone=lambda: "generic-prefill") + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + index_topk=2048, + kv_lora_rank=512, + qk_rope_head_dim=0, + qk_nope_head_dim=128, + v_head_dim=128, + ), + model_arch_config=SimpleNamespace(total_num_attention_heads=64), + dtype=torch.bfloat16, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=3072, max_num_seqs=4), + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + compilation_config=SimpleNamespace( + static_forward_context={ + "layer": SimpleNamespace( + prefill_backend=prefill if has_prefill_backend else None, + dcp_manager=dcp_manager, + ), + } + ), + cache_config=SimpleNamespace(cache_dtype="fp8"), + ) + monkeypatch.setattr( + common, "get_dcp_group", lambda: SimpleNamespace(world_size=dcp) + ) + monkeypatch.setattr(common, "_is_masked_mha_available", lambda *args: False) + monkeypatch.setattr( + SparseMLACommonMetadataBuilder, + "determine_chunked_prefill_workspace_size", + staticmethod(lambda config: 8192), + ) + builder = SparseMLACommonMetadataBuilder( + SimpleNamespace(block_size=64), ["layer"], config, torch.device("cpu") + ) + rows = (8192 + (8192 // dcp if dcp > 1 else 0)) if has_prefill_backend else 0 + assert builder.chunked_prefill_workspace.shape == (rows, 512) + assert builder.chunked_prefill_workspace_size == ( + 8192 if has_prefill_backend else 0 + ) + assert len(gathers) == int(has_prefill_backend and dcp > 1) + if gathers: + assert gathers[0][0] is builder.chunked_prefill_workspace + assert gathers[0][1] == 8192 + + @pytest.mark.parametrize("grouped_backend", [False, True]) @pytest.mark.parametrize("explicit_builder", [False, True]) def test_sparse_backend_allocates_only_its_consumed_index_group( @@ -1537,6 +1600,63 @@ def bind(**kwargs): assert all(probe() is None for probe in probes), published_owners +@pytest.mark.parametrize("record_bytes", [528, 304]) +@pytest.mark.parametrize("rows", [4, 3072]) +@torch.inference_mode() +def test_cache_writer_preparation_releases_probes_and_rebinds_live_inputs( + record_bytes, rows +): + """Compilation owns formats; only the serving cache survives preparation.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("native B12X cache writer requires SM12x") + from b12x.attention import sparse_mla + from b12x.preparation import PreparationSession + + device = torch.device("cuda", torch.accelerator.current_device_index()) + cache = torch.full((2, 64, record_bytes), 17, dtype=torch.uint8, device=device) + original = cache.clone() + impl = object.__new__(B12xMLASparseImpl) + impl._bound_kv_cache = cache + impl._plans = {} + impl._is_glm_next = True + impl._max_tokens = rows + impl.kv_lora_rank = 512 + impl._module = sparse_mla + impl.dcp_world_size = 2 + impl._preparation_prefix = lambda: "test.cache-writer-lifetime" + gc.collect() + torch.accelerator.synchronize() + before = torch.accelerator.memory_allocated(device) + (unit,) = impl.get_b12x_preparation_units(impl, SimpleNamespace(token_counts=())) + # Retained declaration metadata must not allocate token-sized GPU inputs. + assert torch.accelerator.memory_allocated(device) == before + with PreparationSession(device=device, autotune=False) as session: + session.prepare(unit.requests) + session.freeze() + torch.accelerator.synchronize() + assert impl._cache_writer_plan.prepared.owners == () + assert torch.accelerator.memory_allocated(device) == before + torch.testing.assert_close(cache, original, rtol=0, atol=0) + + keys = torch.randn((4, 512), dtype=torch.bfloat16, device=device) + slots = torch.tensor([1, 4, 9, 63], dtype=torch.int64, device=device) + state = impl._cache_writer_plan.prepared.state + state.run(keys, cache, slots) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + state.run(keys, cache, slots) + for offset in (0, 1, 64): + keys.normal_() + slots.copy_(torch.tensor([0, 3, 8, 62], device=device) + offset) + cache.copy_(original) + state.run(keys, cache, slots) + expected = cache.clone() + cache.copy_(original) + graph.replay() + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + assert not torch.equal(cache, original) + + @pytest.mark.parametrize( "mode,rows,heads", [("decode", 4, 16), ("extend", 64, 16), ("extend", 4096, 64)], diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 9e6664913ee6..cf2af081eb75 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -204,8 +204,16 @@ def __init__( self.dcp_local_block_size = self.cp_kv_cache_interleave_size self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size + attention_layer = vllm_config.compilation_config.static_forward_context[ + layer_names[0] + ] + layer_prefill_backend = attention_layer.prefill_backend + # Backends with native sparse prefill use their own caller-owned scratch + # and never build generic chunked-context metadata. self.chunked_prefill_workspace_size = ( self.determine_chunked_prefill_workspace_size(vllm_config) + if layer_prefill_backend is not None + else 0 ) workspace_head_size = ( self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim @@ -235,10 +243,6 @@ def __init__( dtype=torch.int32, device=device, ) - attention_layer = vllm_config.compilation_config.static_forward_context[ - layer_names[0] - ] - layer_prefill_backend = attention_layer.prefill_backend self.dcp_manager: MLADCPManager | None = None if self.dcp_world_size > 1: self.dcp_manager = getattr(attention_layer, "dcp_manager", None) diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 6776c452d1ee..efe7aa017960 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1520,28 +1520,37 @@ def get_b12x_preparation_units( ) if self._is_glm_next: from b12x.preparation import PreparedCall - - kv_probe = torch.empty( - (self._max_tokens, self.kv_lora_rank), - dtype=torch.bfloat16, - device=self._bound_kv_cache.device, - ) - ignored_slots = torch.full( - (self._max_tokens,), - -1, - dtype=torch.int64, - device=self._bound_kv_cache.device, - ) + from torch._subclasses.fake_tensor import FakeTensorMode + + # The declaration retains tensor formats for compilation, not + # live inputs. Allocate execution probes only inside each call. + with FakeTensorMode(): + kv_probe = torch.empty( + (self._max_tokens, self.kv_lora_rank), + dtype=torch.bfloat16, + device=self._bound_kv_cache.device, + ) + ignored_slots = torch.empty( + (self._max_tokens,), + dtype=torch.int64, + device=self._bound_kv_cache.device, + ) writer_plan = self._module.plan_cache_writer( kv_probe, self._bound_kv_cache, ignored_slots ) self._cache_writer_plan = writer_plan def prepare_writer(state): - probe = torch.empty_like(kv_probe) + probe = torch.empty( + (self._max_tokens, self.kv_lora_rank), + dtype=torch.bfloat16, + device=state.kv_cache.device, + ) # The writer touches only slot zero; retain that single page # for restoration instead of cloning the serving KV pool. - slots = torch.zeros_like(ignored_slots) + slots = torch.zeros( + (self._max_tokens,), dtype=torch.int64, device=state.kv_cache.device + ) saved_page = state.kv_cache[0].clone() def produce(): @@ -1554,7 +1563,9 @@ def restore(): run=lambda: state.run(probe, state.kv_cache, slots), produce=produce, restore=restore, - owners=(probe, slots, saved_page), + # Closures own these tensors until preparation completes. + # Published launchers rebind live inputs and must not keep + # the probes or the temporary cache-page backup alive. ) requests.append( From ef1aeaf080879865febd27a92c5644233d157987 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 19:54:03 +0000 Subject: [PATCH 097/365] fix(sparse-mla): release shadowed index buffers after MTP sharing When an indexer owns the selection buffer, retain only the indexer reference. A redundant explicit buffer otherwise pins draft storage after the indexer switches to target-owned storage. Explicit-buffer-only layers and setter precedence are unchanged. Validation: four lazy-read, owner-release and explicit-override cases pass; all applicable pre-commit hooks pass. The qualified TP2 allocation census identifies a 25202688-byte draft buffer retained only by the unused fallback field. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_sparse_mla_backends.py | 9 +++++++-- .../layers/attention/sparse_mla_attention.py | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index ab883594aed8..30d641baffc0 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3,6 +3,7 @@ """Unit tests for the sparse MLA backends and utilities.""" import math +import weakref from collections import deque from types import MethodType, SimpleNamespace from unittest.mock import MagicMock @@ -3768,7 +3769,8 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): assert ensure_args[1] == 4 -def test_sparse_impl_observes_repointed_indexer_buffer(): +@pytest.mark.parametrize("explicit_initial_buffer", [False, True]) +def test_sparse_impl_observes_repointed_indexer_buffer(explicit_initial_buffer): """The MTP proposer repoints the draft's indexer at the target model's buffer after the backend impl is built, so the impl must resolve the buffer per read. Snapshotting it in __init__ leaves the layer reading indices nothing writes.""" @@ -3776,13 +3778,16 @@ def test_sparse_impl_observes_repointed_indexer_buffer(): own = torch.zeros(4, 8, dtype=torch.int32) target = torch.ones(4, 8, dtype=torch.int32) indexer = SimpleNamespace(topk_indices_buffer=own) - impl.init_topk_indices_buffer(indexer, None) + impl.init_topk_indices_buffer(indexer, own if explicit_initial_buffer else None) assert impl.topk_indices_buffer is own indexer.topk_indices_buffer = target assert impl.topk_indices_buffer is target + released = weakref.ref(own) + del own + assert released() is None def test_explicit_topk_buffer_supersedes_indexer(): diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index cf2af081eb75..16a751edb106 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -631,7 +631,7 @@ def init_topk_indices_buffer( topk_indices_buffer: torch.Tensor | None, ) -> None: self._indexer = indexer - self._topk_indices_buffer = topk_indices_buffer + self._topk_indices_buffer = topk_indices_buffer if indexer is None else None @property def topk_indices_buffer(self) -> torch.Tensor | None: From 69bf428f3d50e273ca167678f751202acb12aca8 Mon Sep 17 00:00:00 2001 From: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:10:50 -0700 Subject: [PATCH 098/365] Cut transient memory in NVFP4 Marlin scale-factor computation `_nvfp4_compute_scale_factor` materialized an FP32 copy of the whole scale tensor, a boolean mask, and a gathered copy of the positive entries just to find the largest scale. For fused MoE parameters the helper runs on the scales of every expert at once, so on a 288-expert model (GLM-5.3-Flash NVFP4) one call touched about 300M elements and allocated several GiB of transient device memory during weight loading. On unified-memory machines with little headroom at load time (DGX Spark GB10, see #50925) that spike is enough to fail the load. The factor only depends on the largest scale, so reduce the input tensor directly. The result is identical: a max commutes with the exact multiplication by 2^7, non-positive scales never win the max when a positive one exists, and an all-non-positive tensor still yields 1.0. An empty tensor also yields 1.0 as before. A NaN scale now raises instead of being silently skipped. Measured on CPU with a bf16 tensor of 37.7M elements (one eighth of the production fused w13 scales): peak RSS growth 1011 MiB before, 3 MiB after, same factor (32768.0). The same shape on CUDA is covered by the new test, which asserts no tensor-sized temporaries. Signed-off-by: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Co-Authored-By: Claude Fable 5.1 (cherry picked from commit 941bfa989ca050aba3627e455ab8bd70e45a49ed) (cherry picked from commit 3e9d78044375b2bc078d9a864cb5d5e5a9bd01dc) --- .../test_nvfp4_marlin_scale_factor.py | 100 ++++++++++++++++++ .../quantization/utils/marlin_utils_fp4.py | 20 ++-- 2 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py new file mode 100644 index 000000000000..f99eb06317b6 --- /dev/null +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the NVFP4 Marlin power-of-2 scale factor helper. + +The helper must return the same factor as the previous mask-and-gather +implementation while allocating no tensor-sized temporaries. +""" + +import pytest +import torch + +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + _nvfp4_compute_scale_factor, +) + +DTYPES = [torch.float16, torch.bfloat16] +SHAPES = [(64,), (128, 256), (8, 64, 128)] +UPPER = 448 * (2**7) + + +def _reference_scale_factor( + marlin_scales: torch.Tensor, a_dtype: torch.dtype | None = None +) -> float: + """The previous implementation, kept verbatim as the oracle.""" + if a_dtype is not None and a_dtype == torch.half: + return 1.0 + ws_float = marlin_scales.float() * (2**7) + nonzero_mask = ws_float > 0 + if nonzero_mask.any(): + max_val = ws_float[nonzero_mask].max() + if max_val < UPPER: + sf = (UPPER / max_val).log2().floor().exp2() + return sf.item() + return 1.0 + + +def _cases(shape: tuple[int, ...], dtype: torch.dtype) -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(0) + small = torch.rand(shape, generator=gen) * 0.01 + with_zeros = small.clone() + with_zeros.flatten()[::3] = 0.0 + with_negatives = small.clone() + with_negatives.flatten()[::5] *= -1.0 + normalized = torch.rand(shape, generator=gen) * 400.0 + 48.0 + return { + "small": small.to(dtype), + "with_zeros": with_zeros.to(dtype), + "with_negatives": with_negatives.to(dtype), + "all_zeros": torch.zeros(shape, dtype=dtype), + "all_negative": (-small - 0.001).to(dtype), + "already_normalized": normalized.to(dtype), + "single_large": torch.full(shape, 0.5, dtype=dtype), + } + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("shape", SHAPES) +@pytest.mark.parametrize("a_dtype", [None, torch.float16, torch.bfloat16]) +def test_scale_factor_matches_reference( + dtype: torch.dtype, shape: tuple[int, ...], a_dtype: torch.dtype | None +): + for name, scales in _cases(shape, dtype).items(): + expected = _reference_scale_factor(scales, a_dtype) + actual = _nvfp4_compute_scale_factor(scales, a_dtype) + assert actual == expected, name + assert actual >= 1.0 + assert actual == 2.0 ** round(torch.tensor(actual).log2().item()) + + +def test_scale_factor_rescales_into_range(): + scales = torch.full((32, 64), 2**-10, dtype=torch.bfloat16) + sf = _nvfp4_compute_scale_factor(scales) + rescaled = scales.float().max() * (2**7) * sf + assert 2.0 <= rescaled < UPPER + + +def test_scale_factor_empty_tensor(): + assert _nvfp4_compute_scale_factor(torch.empty(0, dtype=torch.bfloat16)) == 1.0 + + +def test_scale_factor_rejects_nan(): + scales = torch.rand(16, 16, dtype=torch.bfloat16) + scales[3, 3] = float("nan") + with pytest.raises(ValueError, match="NaN"): + _nvfp4_compute_scale_factor(scales) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_scale_factor_allocates_no_tensor_temporaries(): + device = torch.device("cuda") + scales = torch.rand((288, 512, 256), dtype=torch.bfloat16, device=device) * 0.01 + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + baseline = torch.cuda.memory_allocated(device) + sf = _nvfp4_compute_scale_factor(scales) + torch.cuda.synchronize(device) + transient = torch.cuda.max_memory_allocated(device) - baseline + assert sf > 1.0 + # The previous implementation needed ~4.5x the tensor size here. + assert transient < scales.numel() * scales.element_size() // 64 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 3b6990101ed7..1e5ad5660186 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -47,14 +47,20 @@ def _nvfp4_compute_scale_factor( # no rescaling is applied here if active dtype is half. if a_dtype is not None and a_dtype == torch.half: return 1.0 + if marlin_scales.numel() == 0: + return 1.0 - ws_float = marlin_scales.float() * (2**7) - nonzero_mask = ws_float > 0 - if nonzero_mask.any(): - max_val = ws_float[nonzero_mask].max() - if max_val < 448 * (2**7): - sf = (448 * (2**7) / max_val).log2().floor().exp2() - return sf.item() + # The factor depends only on the largest scale, so reduce the input + # directly. Materializing an FP32 copy, a boolean mask, and a gathered + # copy of every scale costs up to 9x the tensor size in transient device + # memory; for fused MoE parameters (all experts at once) that is several + # GiB during weight loading. + max_val = marlin_scales.max().float() * (2**7) + if torch.isnan(max_val): + raise ValueError("NVFP4 Marlin scales contain NaN") + if 0 < max_val < 448 * (2**7): + sf = (448 * (2**7) / max_val).log2().floor().exp2() + return sf.item() return 1.0 From 173fb9306317bc246509bcc62ffe593d1774d94d Mon Sep 17 00:00:00 2001 From: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:38:18 -0700 Subject: [PATCH 099/365] Shorten the NVFP4 Marlin scale-factor comment Keep the helper comment to two lines, per AGENTS.md guidance on brevity. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Elevatormusic <22101396+Elevatormusic@users.noreply.github.com> (cherry picked from commit 78722557d8eeab820555e04147256d847b0a8756) (cherry picked from commit 9a7206c69fc09bbb1b6e78862d5613b0afc150ac) --- .../layers/quantization/utils/marlin_utils_fp4.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 1e5ad5660186..90d833bbc0c3 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -50,11 +50,8 @@ def _nvfp4_compute_scale_factor( if marlin_scales.numel() == 0: return 1.0 - # The factor depends only on the largest scale, so reduce the input - # directly. Materializing an FP32 copy, a boolean mask, and a gathered - # copy of every scale costs up to 9x the tensor size in transient device - # memory; for fused MoE parameters (all experts at once) that is several - # GiB during weight loading. + # Reduce the input directly: the FP32 copy, mask, and gathered copy of the + # old path cost several times the tensor size in transient memory. max_val = marlin_scales.max().float() * (2**7) if torch.isnan(max_val): raise ValueError("NVFP4 Marlin scales contain NaN") From 9e4225e9c47de0df183c8b254408a34b0e614577 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:31:00 +0000 Subject: [PATCH 100/365] test(marlin): qualify bounded NVFP4 scale reduction on CUDA Cover finite E4M3 scale values in contiguous and strided GPU tensors, preserve exact scale-factor parity, and record scalar-reduction peak memory. Use portable accelerator memory APIs. The imported upstream implementation and author attribution remain intact. Validation: 26 native tests pass on RTX PRO 6000; mask-and-gather reference uses 1,243,946,496 transient bytes versus 2,560 for the reduction on 37,748,736 BF16 scales. Applicable pre-commit hooks pass. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 25fec51de35363e7f565fec39237dc3db9adc59b) --- .../test_nvfp4_marlin_scale_factor.py | 36 ++++++++++++++----- .../quantization/utils/marlin_utils_fp4.py | 4 +-- 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py index f99eb06317b6..f2521c35f518 100644 --- a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the NVFP4 Marlin power-of-2 scale factor helper. -The helper must return the same factor as the previous mask-and-gather +The helper must return the same factor as the mask-and-gather implementation while allocating no tensor-sized temporaries. """ @@ -21,7 +21,7 @@ def _reference_scale_factor( marlin_scales: torch.Tensor, a_dtype: torch.dtype | None = None ) -> float: - """The previous implementation, kept verbatim as the oracle.""" + """Reference reduction over positive scales in FP32.""" if a_dtype is not None and a_dtype == torch.half: return 1.0 ws_float = marlin_scales.float() * (2**7) @@ -86,15 +86,33 @@ def test_scale_factor_rejects_nan(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_scale_factor_allocates_no_tensor_temporaries(): +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_scale_factor_matches_fp8_scale_domain_on_cuda(dtype, contiguous): + """All nonnegative finite E4M3 scales retain their factor after reduction.""" + values = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn) + scales = values.to(dtype=dtype, device="cuda").repeat(4, 1).T + if contiguous: + scales = scales.contiguous() + for multiplier in (1.0, 2**-7, 2**-10): + sample = scales * multiplier + assert _nvfp4_compute_scale_factor(sample, torch.bfloat16) == ( + _reference_scale_factor(sample, torch.bfloat16) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_scale_factor_allocates_no_tensor_temporaries(record_property): device = torch.device("cuda") scales = torch.rand((288, 512, 256), dtype=torch.bfloat16, device=device) * 0.01 - torch.cuda.synchronize(device) - torch.cuda.reset_peak_memory_stats(device) - baseline = torch.cuda.memory_allocated(device) + torch.accelerator.synchronize(device) + torch.accelerator.reset_peak_memory_stats(device) + baseline = torch.accelerator.memory_allocated(device) sf = _nvfp4_compute_scale_factor(scales) - torch.cuda.synchronize(device) - transient = torch.cuda.max_memory_allocated(device) - baseline + torch.accelerator.synchronize(device) + transient = torch.accelerator.max_memory_allocated(device) - baseline + record_property("transient_bytes", transient) + record_property("input_bytes", scales.numel() * scales.element_size()) assert sf > 1.0 - # The previous implementation needed ~4.5x the tensor size here. + # Scalar reduction must not allocate a scale-tensor-sized temporary. assert transient < scales.numel() * scales.element_size() // 64 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 90d833bbc0c3..863f4a6d6ac7 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -50,8 +50,8 @@ def _nvfp4_compute_scale_factor( if marlin_scales.numel() == 0: return 1.0 - # Reduce the input directly: the FP32 copy, mask, and gathered copy of the - # old path cost several times the tensor size in transient memory. + # Reduce before conversion to keep temporary storage independent of the + # number of experts. Power-of-two multiplication preserves the maximum. max_val = marlin_scales.max().float() * (2**7) if torch.isnan(max_val): raise ValueError("NVFP4 Marlin scales contain NaN") From 3d2b2e53e6740e480514442edea43c9e1424df10 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 11:52:24 +0000 Subject: [PATCH 101/365] Bound NVFP4 Marlin expert-scale preparation memory Reduce one expert at a time to select the shared scale factor, then write each converted expert directly into the output bank. Avoid retaining a full half-precision bank and concatenating a second full output. Finite scale bytes, shared global factors, input ownership, gated and ungated padding remain unchanged; inference kernels are untouched. Validation: 35 CPU/CUDA tests pass on SM120, including bitwise whole-bank reference parity and a memory regression. For 27 MiB output scales, peak additional memory is 28,775,936 bytes versus 75,503,616 bytes before this change. All applicable pre-commit hooks pass. Whole-model TP2 qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 4a8904910ab41785dc66353924a3dcb6fa5a4b32) --- .../test_nvfp4_marlin_scale_factor.py | 137 ++++++++++++++++++ .../quantization/utils/marlin_utils_fp4.py | 42 ++++-- 2 files changed, 164 insertions(+), 15 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py index f2521c35f518..a1fa6726c07b 100644 --- a/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py +++ b/tests/kernels/quantization/test_nvfp4_marlin_scale_factor.py @@ -6,11 +6,20 @@ implementation while allocating no tensor-sized temporaries. """ +from types import SimpleNamespace + import pytest import torch +from vllm.model_executor.layers.quantization.utils import marlin_utils_fp4 +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + marlin_permute_scales, +) from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( _nvfp4_compute_scale_factor, + nvfp4_marlin_process_global_scale, + nvfp4_marlin_process_scales, + prepare_nvfp4_moe_layer_for_marlin, ) DTYPES = [torch.float16, torch.bfloat16] @@ -116,3 +125,131 @@ def test_scale_factor_allocates_no_tensor_temporaries(record_property): assert sf > 1.0 # Scalar reduction must not allocate a scale-tensor-sized temporary. assert transient < scales.numel() * scales.element_size() // 64 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("intermediate", [64, 96]) +@pytest.mark.parametrize("gated", [False, True]) +def test_moe_scales_match_whole_bank_reference(dtype, intermediate, gated): + """Per-expert conversion preserves scale bytes, padding and global factors.""" + experts, hidden = 4, 128 + shards = 2 if gated else 1 + padded = (intermediate + 63) // 64 * 64 + layer = SimpleNamespace( + num_experts=experts, + hidden_size=hidden, + intermediate_size_per_partition=intermediate, + params_dtype=dtype, + ) + + def scales(shape): + count = 1 + for extent in shape: + count *= extent + # Include zero and finite positive E4M3 scales at several exponents. + values = ( + (torch.arange(count, device="cuda") % 120) + .to(torch.uint8) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + factors = torch.tensor([0.125, 0.25, 0.5, 1.0], device="cuda") + return (values.float() * factors[:, None, None]).to(torch.float8_e4m3fn) + + s13 = scales((experts, shards * intermediate, hidden // 16)) + s2 = scales((experts, hidden, intermediate // 16)) + g = torch.arange(1, experts + 1, dtype=torch.float32, device="cuda") + inputs = (s13.view(torch.uint8).clone(), s2.view(torch.uint8).clone()) + result = prepare_nvfp4_moe_layer_for_marlin( + layer, + torch.zeros( + (experts, shards * intermediate, hidden // 2), + dtype=torch.uint8, + device="cuda", + ), + s13, + g, + torch.zeros( + (experts, hidden, intermediate // 2), dtype=torch.uint8, device="cuda" + ), + s2, + g, + gated, + ) + for source, actual, actual_global, is_gate in ( + (s13, result[1], result[2], True), + (s2, result[4], result[5], False), + ): + source = source.to(dtype) + if is_gate: + source = source.reshape(experts, shards, intermediate, -1) + source = torch.nn.functional.pad(source, (0, 0, 0, padded - intermediate)) + source = source.reshape(experts, shards * padded, -1) + size_k, size_n = hidden, shards * padded + else: + source = torch.nn.functional.pad(source, (0, (padded - intermediate) // 16)) + size_k, size_n = padded, hidden + factor = _reference_scale_factor(source, dtype) + expected = [] + for expert in source: + permuted = marlin_permute_scales(expert.T, size_k, size_n, 16) + converted, _ = nvfp4_marlin_process_scales(permuted, factor, dtype) + expected.append(converted) + torch.testing.assert_close( + actual.view(torch.uint8), + torch.stack(expected).view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual_global, + nvfp4_marlin_process_global_scale(g, dtype) / factor, + rtol=0, + atol=0, + ) + torch.testing.assert_close(s13.view(torch.uint8), inputs[0], rtol=0, atol=0) + torch.testing.assert_close(s2.view(torch.uint8), inputs[1], rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_moe_scale_conversion_avoids_bank_temporaries(monkeypatch, record_property): + """Additional storage contains output scales and one expert's conversion.""" + experts, hidden, intermediate = 288, 1024, 512 + layer = SimpleNamespace( + num_experts=experts, + hidden_size=hidden, + intermediate_size_per_partition=intermediate, + params_dtype=torch.bfloat16, + ) + # Weight packing is independently tested; isolate scale preparation's peak. + monkeypatch.setattr(marlin_utils_fp4, "_repack_marlin_experts", lambda w, *a: w) + weight = torch.zeros((), dtype=torch.uint8, device="cuda") + s13 = torch.ones( + (experts, 2 * intermediate, hidden // 16), + dtype=torch.float8_e4m3fn, + device="cuda", + ) + s2 = torch.ones( + (experts, hidden, intermediate // 16), dtype=torch.float8_e4m3fn, device="cuda" + ) + global_scale = torch.ones(experts, dtype=torch.float32, device="cuda") + torch.accelerator.synchronize() + torch.accelerator.reset_peak_memory_stats() + before = torch.accelerator.memory_allocated() + result = prepare_nvfp4_moe_layer_for_marlin( + layer, + weight.expand(experts, 2 * intermediate, hidden // 2), + s13, + global_scale, + weight.expand(experts, hidden, intermediate // 2), + s2, + global_scale, + True, + ) + torch.accelerator.synchronize() + extra = torch.accelerator.max_memory_allocated() - before + output_bytes = sum(t.numel() * t.element_size() for t in (result[1], result[4])) + record_property("peak_extra_bytes", extra) + record_property("output_scale_bytes", output_bytes) + assert extra < output_bytes + 4 * 1024 * 1024 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 863f4a6d6ac7..a94823434cbf 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -379,9 +379,10 @@ def pad_w13(x: torch.Tensor) -> torch.Tensor: tensor to padded_N rows.""" if padded_N == N: return x - x = x.view(E, num_shards, N, x.size(-1)) + experts = x.size(0) + x = x.view(experts, num_shards, N, x.size(-1)) x = torch.nn.functional.pad(x, (0, 0, 0, padded_N - N)) - return x.reshape(E, num_shards * padded_N, -1) + return x.reshape(experts, num_shards * padded_N, -1) def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: """Zero-pad the packed N (last) dim of a (E, K, N / packing) @@ -423,24 +424,29 @@ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: def permute_scales( scales: torch.Tensor, g_scales: torch.Tensor, name: str ) -> tuple[torch.Tensor, torch.Tensor]: - scales = scales.to(param_dtype) - - tensor_list = [] if "w13" in name: - scales = pad_w13(scales) size_n, size_k = padded_N * num_shards, K else: - scales = pad_w2(scales, packing=GROUP_SIZE) size_n, size_k = K, padded_N - # All experts share one global_scale, so compute the max - # scale_factor across all experts first, then apply uniformly. - combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype) + # A shared factor must cover all experts. Convert one expert at a time + # so loading never retains a half-precision copy of the entire bank. + combined_scale_factor = 1.0 + if param_dtype != torch.half: + maxima = torch.empty(E, device=scales.device, dtype=param_dtype) + for i in range(E): + torch.amax(scales[i].to(param_dtype), out=maxima[i]) + combined_scale_factor = _nvfp4_compute_scale_factor(maxima, param_dtype) + output = None for i in range(E): - scale = scales[i].T + scale = scales[i : i + 1].to(param_dtype) + if "w13" in name: + scale = pad_w13(scale) + else: + scale = pad_w2(scale, packing=GROUP_SIZE) marlin_scales = marlin_permute_scales( - s=scale, + s=scale[0].T, size_k=size_k, size_n=size_n, group_size=GROUP_SIZE, @@ -449,12 +455,18 @@ def permute_scales( marlin_scales, _ = nvfp4_marlin_process_scales( marlin_scales, scale_factor=combined_scale_factor, a_dtype=param_dtype ) - tensor_list.append(marlin_scales) + if output is None: + output = torch.empty( + (E, *marlin_scales.shape), + device=marlin_scales.device, + dtype=marlin_scales.dtype, + ) + output[i].copy_(marlin_scales) - scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + assert output is not None g_scales = nvfp4_marlin_process_global_scale(g_scales, param_dtype) g_scales = g_scales / combined_scale_factor - return scales, g_scales + return output, g_scales w13_scale, w13_scale_2 = permute_scales(w13_scale, w13_scale_2, "w13") w2_scale, w2_scale_2 = permute_scales(w2_scale, w2_scale_2, "w2") From d6a88c611d56885365d9b29c85d5a25c146208fa Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:16:11 +0000 Subject: [PATCH 102/365] Bound GLM MTP vocabulary-head scale reduction memory Reduce and sanitize FP32 blocks without materializing a full vocabulary-sized copy. Preserve exact scale values, packed NVFP4 bytes, and the target BF16 head. CUDA parity and allocation tests pass; full TP2 startup is validated separately. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 59aa0765cdb341e80b6048773f1a001d9bafc1ed) --- tests/models/test_glm5next_model.py | 76 +++++++++++++++++++ vllm/models/glm5next/nvidia/mtp_draft_head.py | 15 +++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 39e845f70a04..2491b3b9f645 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -308,6 +308,82 @@ def test_glm5next_mtp_draft_head_rejects_unknown_mode(monkeypatch) -> None: mtp_draft_head.configured_draft_head_mode() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("kind", ["finite", "zeros", "nonfinite", "bf16_bits"]) +def test_glm5next_mtp_head_scale_preserves_fp32_reduction(kind: str) -> None: + if kind == "bf16_bits": + weight = torch.arange(65536, dtype=torch.int32, device="cuda").to(torch.int16) + weight = weight.view(torch.bfloat16).view(256, 256) + else: + weight = torch.linspace(-8, 6, 512, device="cuda", dtype=torch.bfloat16) + weight = weight.view(16, 32) + if kind == "zeros": + weight.zero_() + elif kind == "nonfinite": + weight[0, :3] = torch.tensor( + [float("nan"), float("inf"), -float("inf")], device="cuda" + ) + original_bits = weight.view(torch.int16).clone() + expected = 2688.0 / weight.float().abs().nan_to_num().max() + + actual = mtp_draft_head._nvfp4_weight_global_scale(weight) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(weight.view(torch.int16), original_bits, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_glm5next_mtp_head_scale_bounds_conversion_memory(record_property) -> None: + weight = torch.ones((1024, 16384), dtype=torch.bfloat16, device="cuda") + weight[-1, -1] = -8 + mtp_draft_head._nvfp4_weight_global_scale(weight[:1]) + torch.accelerator.synchronize() + torch.accelerator.reset_peak_memory_stats() + allocated = torch.accelerator.memory_allocated() + + actual = mtp_draft_head._nvfp4_weight_global_scale(weight) + + torch.accelerator.synchronize() + peak = torch.accelerator.max_memory_allocated() - allocated + record_property("peak_additional_bytes", peak) + torch.testing.assert_close( + actual, torch.tensor(336.0, device="cuda"), rtol=0, atol=0 + ) + assert weight[-1, -1].item() == -8 + assert peak < 17 * 1024 * 1024 + + +@pytest.mark.skipif( + not current_platform.is_device_capability(120), reason="requires SM120" +) +def test_glm5next_mtp_head_quantization_preserves_packed_weights() -> None: + import flashinfer + + source = torch.nn.Linear(256, 256, bias=False, dtype=torch.bfloat16).cuda() + source.shard_indices = SimpleNamespace() + original = source.weight.detach().clone() + scale = 2688.0 / original.float().abs().nan_to_num().max() + weight, scales = flashinfer.nvfp4_quantize( + original, + scale, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=False, + backend="cuda", + ) + expected = flashinfer.prepare_bf16_fp4_weights( + weight.view(torch.uint8), + scales, + scale.reciprocal().reshape(1), + backend="cute-dsl", + ) + + actual = mtp_draft_head.QuantizedDraftHead(source, "nvfp4") + + for observed, reference in zip(actual.buffers(), expected, strict=True): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + torch.testing.assert_close(source.weight, original, rtol=0, atol=0) + + def test_glm5next_mtp_prepares_configured_draft_head(monkeypatch) -> None: source_head = torch.nn.Linear(4, 8, bias=False) quantized_head = torch.nn.Linear(4, 8, bias=False) diff --git a/vllm/models/glm5next/nvidia/mtp_draft_head.py b/vllm/models/glm5next/nvidia/mtp_draft_head.py index ae662622e9eb..03a4f1f34937 100644 --- a/vllm/models/glm5next/nvidia/mtp_draft_head.py +++ b/vllm/models/glm5next/nvidia/mtp_draft_head.py @@ -30,6 +30,17 @@ _SUPPORTED_MODES = frozenset(("bf16", "nvfp4")) +def _nvfp4_weight_global_scale(weight: torch.Tensor) -> torch.Tensor: + """Compute the BF16 head's FP32 scale with at most 16 MiB of conversion scratch.""" + # Each float() result owns its storage; sanitization must not alter the + # BF16 target head. Chunk maxima preserve the full FP32 reduction exactly. + maxima = [ + chunk.float().abs_().nan_to_num_().amax() + for chunk in weight.view(-1).split(4 * 1024 * 1024) + ] + return _NVFP4_GLOBAL_MAX / torch.stack(maxima).amax() + + def configured_draft_head_mode() -> str: """Return the validated GLM MTP draft-head storage mode.""" mode = os.getenv("VLLM_GLM53_MTP_DRAFT_HEAD", "bf16").strip().lower() @@ -91,9 +102,7 @@ def __init__(self, source_head: nn.Module, mode: str) -> None: ) import flashinfer - weight_global_scale = ( - _NVFP4_GLOBAL_MAX / weight.float().abs().nan_to_num().max() - ) + weight_global_scale = _nvfp4_weight_global_scale(weight) weight_fp4, weight_sf = flashinfer.nvfp4_quantize( weight, weight_global_scale, From 02dc2740c795815356f81c812d8db43fb4d1159b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 11 Sep 2026 21:51:36 +0000 Subject: [PATCH 103/365] Account for recurrent endpoint storage when sizing the KV pool Reserve private instruction, prompt, and response checkpoint blocks in boundary-policy layout selection, maximum-context admission, and capacity reporting. Keep aligned-policy accounting unchanged. This prevents a small hybrid pool from passing startup while every request waits for unavailable checkpoint blocks. Validation: 131 KV sizing and contiguous-packing CPU tests passed; the seven-group TP2 MTP3 reproducer requires 41 blocks without an instruction checkpoint and 49 with one, exceeding its 37 usable blocks. Ruff and git diff --check pass. GPU serving qualification is pending. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit 94febe5c35530d630b161b5abf5617e484b70fbe) (cherry picked from commit fd047ddfc7f923df1f411a689a8c4e56a6d7e87d) --- tests/v1/core/test_contiguous_kv_packing.py | 1 + tests/v1/core/test_kv_cache_utils.py | 104 ++++++++++++++++++++ vllm/v1/core/kv_cache_utils.py | 35 ++++++- 3 files changed, 139 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index dc56526e2329..15888fe56db2 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -72,6 +72,7 @@ def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5): def _mock_vllm_config(layout: str | None): config = MagicMock() + config.use_request_boundary_checkpoints = False config.cache_config = CacheConfig() config.cache_config.num_gpu_blocks_override = None config.cache_config.kv_cache_layout = layout diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1da2be251f9a..76fc5f2945c5 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2653,6 +2653,42 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): ) +@pytest.mark.parametrize("with_tail", [False, True]) +def test_balanced_glm_boundary_capacity_includes_private_endpoints( + monkeypatch, with_tail +): + """The GLM slot-sharing layout still reserves distinct endpoint block IDs.""" + config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + specs = ( + _glm5_like_kv_cache_spec_with_tail() + if with_tail + else _glm5_like_kv_cache_spec()[0] + ) + groups = get_kv_cache_groups(config, specs) + assert kv_cache_utils._glm5_next_tensor_layout(groups) is not None + block_bytes = kv_cache_utils._pool_bytes_per_block(groups) + live_bytes = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) + live_blocks = live_bytes // block_bytes + private_blocks = 3 * (len(groups) + 1) + monkeypatch.setattr( + VllmConfig, "use_request_boundary_checkpoints", property(lambda self: True) + ) + + required = (live_blocks + private_blocks) * block_bytes + assert ( + kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) == required + ) + assert ( + kv_cache_utils._get_kv_cache_group_allocation_cost(config, groups) == required + ) + pool = kv_cache_utils.get_kv_cache_config_from_groups(config, groups, required) + assert get_max_concurrency_for_kv_cache_config(config, pool) == 1 + assert ( + kv_cache_utils._estimate_max_model_len_from_groups(config, groups, live_bytes) + < config.model_config.max_model_len + ) + + def test_glm5_kpool_tail_does_not_drag_hash_block_size(): """The tail's kpool-sized scratch block (4 tokens) must not constrain the prefix-cache hash granularity: participating groups alone decide it.""" @@ -3243,6 +3279,7 @@ def test_glm5next_split_cache_preserves_physical_pages( def _glm5next_split_config() -> SimpleNamespace: return SimpleNamespace( + use_request_boundary_checkpoints=False, scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), speculative_config=None, model_config=SimpleNamespace(max_model_len=202_752), @@ -3380,6 +3417,7 @@ def test_glm5next_nvfp4_auto_geometry_capacity() -> None: prefix_cache_retention_interval=4096, ) vllm_config = SimpleNamespace( + use_request_boundary_checkpoints=False, model_config=SimpleNamespace(max_model_len=202_752), parallel_config=SimpleNamespace(decode_context_parallel_size=4), cache_config=SimpleNamespace(mamba_cache_mode="align"), @@ -4238,6 +4276,72 @@ def test_auto_fit_max_model_len_with_hybrid(): assert vllm_config.model_config.max_model_len == 1024 +@pytest.mark.parametrize("dcp", [1, 2, 4]) +@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (70, 15)]) +def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_pages): + """Live MTP state alone does not cover instruction/prompt/response copies.""" + config = SimpleNamespace( + use_request_boundary_checkpoints=True, + attention_config=SimpleNamespace(hisparse_config=None), + model_config=SimpleNamespace(max_model_len=1048576), + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), + cache_config=SimpleNamespace(mamba_cache_mode="align"), + ) + attention = MLAAttentionSpec( + block_size=2048, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + model_version="glm5_next", + ) + state = MambaSpec( + block_size=2048, + shapes=((1,),), + dtypes=(torch.uint8,), + mamba_cache_mode="align", + num_speculative_blocks=3, + ) + groups = [KVCacheGroupSpec([f"state-{i}"], state) for i in range(6)] + groups.append(KVCacheGroupSpec(["attention"], attention)) + block_bytes = kv_cache_utils._pool_bytes_per_block(groups) + # Six groups require five live/scratch states each; three endpoints + # additionally own eight blocks each. The null block is unavailable. + available = (num_blocks - 1) * block_bytes + estimate = kv_cache_utils._estimate_max_model_len_from_groups( + config, groups, available + ) + assert estimate == expected_pages * 2048 * dcp + assert config.model_config.max_model_len == 1048576 + + if expected_pages == 0: + with pytest.raises(ValueError, match="even a single token"): + kv_cache_utils._auto_fit_max_model_len(config, [groups], [available]) + else: + kv_cache_utils._auto_fit_max_model_len(config, [groups], [available]) + assert config.model_config.max_model_len == estimate + assert kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) == ( + available + ) + assert kv_cache_utils._get_kv_cache_group_allocation_cost(config, groups) == ( + available + ) + cache = KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ) + assert get_max_concurrency_for_kv_cache_config(config, cache) == ( + num_blocks / (30 + 24 + expected_pages) + ) + + # An aligned-policy control has no endpoint reservation. Its capacity + # calculation must keep the ordinary live-state contract. + config.use_request_boundary_checkpoints = False + config.model_config.max_model_len = 1048576 + control = kv_cache_utils._estimate_max_model_len_from_groups( + config, groups, available + ) + assert control == (num_blocks - 1 - 30) * 2048 * dcp + + def test_auto_fit_max_model_len_not_triggered(): """Test that auto-fit is not triggered when original_max_model_len is not -1.""" model_config = ModelConfig(max_model_len=16) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 490273b796a4..567cab0e4628 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1005,6 +1005,14 @@ def check_enough_kv_cache_memory( if groups else available_memory ) + if vllm_config.use_request_boundary_checkpoints: + _check_enough_kv_cache_memory( + check_memory, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + return _check_enough_kv_cache_memory( check_memory, lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), @@ -1068,6 +1076,18 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: return True +def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) -> int: + """Count private instruction, prompt, and response checkpoint storage.""" + if not num_groups or not vllm_config.use_request_boundary_checkpoints: + return 0 + + from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS + + # KVCacheManager.allocate_slots reserves one block per group and one + # auxiliary block per endpoint, separately from the live sequence state. + return NUM_BOUNDARY_CHECKPOINT_SLOTS * (num_groups + 1) + + def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1095,6 +1115,9 @@ def get_max_concurrency_for_kv_cache_config( host_blocks_per_request += required else: num_blocks_per_request += required + num_blocks_per_request += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_config.kv_cache_groups) + ) limits = [kv_cache_config.num_blocks / num_blocks_per_request] if host_blocks_per_request: assert kv_cache_config.hisparse_host_num_blocks is not None @@ -1748,6 +1771,9 @@ def _get_kv_cache_group_allocation_cost( ) for group in kv_cache_groups ) + num_blocks_per_request += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_groups) + ) return _get_kv_cache_bytes_per_block(kv_cache_groups) * num_blocks_per_request @@ -2776,7 +2802,10 @@ def _max_memory_usage_bytes_from_groups( return 0 if vllm_config.attention_config.hisparse_config is not None: - return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + ( + _request_boundary_reserve_blocks(vllm_config, len(kv_cache_groups)) + * _pool_bytes_per_block(kv_cache_groups) + ) if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( @@ -2800,6 +2829,9 @@ def _max_memory_usage_bytes_from_groups( ) if tail_names: total_blocks += 1 + total_blocks += _request_boundary_reserve_blocks( + vllm_config, len(kv_cache_groups) + ) return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page) bytes_per_block = _pool_bytes_per_block(kv_cache_groups) @@ -2814,6 +2846,7 @@ def _max_memory_usage_bytes_from_groups( spec.page_size_bytes, ) + total_blocks += _request_boundary_reserve_blocks(vllm_config, len(kv_cache_groups)) return bytes_per_block * total_blocks From 8cb89ee1d3b84cad3e2d5531d7e70dd8d3dc2490 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 13:18:34 +0000 Subject: [PATCH 104/365] fix(glm): bound and prepare the DCP profiling cache Measure the single-request attention peak with one temporary cache block rather than a block for every configured request slot. Prepare pool-bound B12X plans before the dummy forward and release profile state on success or failure. Real serving capacity and CUDA-graph batch profiling are unchanged. Validation: 23 focused profile, state-lifecycle and sparse-group tests pass. Negative controls on the installed unmodified wheel fail. Existing GPUModelRunner typing and allocator-cleanup API hook findings reproduce on the unchanged parent; no new findings remain. Whole-model TP2/DCP2 qualification follows in a source-built wheel. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit b7d22fb3772d6c93e67cf624f7bb9ed95b53af00) --- tests/v1/worker/test_gpu_model_runner.py | 38 +++++++++++++++------ tests/v1/worker/test_gpu_model_runner_v2.py | 38 +++++++++++++++------ tests/v1/worker/test_gpu_worker.py | 14 ++++++-- vllm/v1/worker/gpu/model_runner.py | 12 ++++--- vllm/v1/worker/gpu_model_runner.py | 11 +++--- vllm/v1/worker/gpu_worker.py | 7 +++- 6 files changed, 89 insertions(+), 31 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 042b378010b4..1e3827bb02df 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1830,10 +1830,10 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): runner.initialize_kv_cache(kv_cache_config) -@pytest.mark.parametrize("dummy_run_fails", [False, True]) +@pytest.mark.parametrize("failure_stage", [None, "init", "prepare", "forward"]) def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( monkeypatch: pytest.MonkeyPatch, - dummy_run_fails: bool, + failure_stage: str | None, ): runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_config = SimpleNamespace( @@ -1844,11 +1844,22 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( runner.max_num_tokens = 4096 events: list[object] = [] - runner._init_minimal_kv_cache_for_profiling = lambda: events.append("init-kv") + def initialize(*, num_blocks): + assert num_blocks == 1 + events.append("init-kv") + if failure_stage == "init": + raise RuntimeError("expected DCP profile failure") + + def prepare(): + events.append("prepare") + if failure_stage == "prepare": + raise RuntimeError("expected DCP profile failure") + + runner._init_minimal_kv_cache_for_profiling = initialize def dummy_run(*args, **kwargs): events.append(("dummy-run", args, kwargs)) - if dummy_run_fails: + if failure_stage == "forward": raise RuntimeError("expected DCP profile failure") return torch.empty(1), torch.empty(1) @@ -1862,14 +1873,22 @@ def dummy_run(*args, **kwargs): lambda _: nullcontext(), ) - if dummy_run_fails: + if failure_stage is not None: with pytest.raises(RuntimeError, match="expected DCP profile failure"): - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) else: - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) assert events[0] == "init-kv" - assert events[1] == ( + assert events[-1] == "cleanup" + if failure_stage == "init": + assert events == ["init-kv", "cleanup"] + return + assert events[1] == "prepare" + if failure_stage == "prepare": + assert events == ["init-kv", "prepare", "cleanup"] + return + assert events[2] == ( "dummy-run", (4096,), { @@ -1879,8 +1898,7 @@ def dummy_run(*args, **kwargs): "single_request_prefill": True, }, ) - assert events[-1] == "cleanup" - if not dummy_run_fails: + if failure_stage is None: assert events[-2] == "sync" diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 35d9f2f5d491..806d8e068e13 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -396,10 +396,10 @@ def test_boundary_logits_only_dispatches_pending_cache_tasks( assert input_batch.positions.item() == 3 -@pytest.mark.parametrize("dummy_run_fails", [False, True]) +@pytest.mark.parametrize("failure_stage", [None, "init", "prepare", "forward"]) def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( monkeypatch: pytest.MonkeyPatch, - dummy_run_fails: bool, + failure_stage: str | None, ): runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_config = SimpleNamespace( @@ -410,10 +410,21 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( runner.max_num_tokens = 4096 events: list[object] = [] + def initialize(_, *, num_blocks): + assert num_blocks == 1 + events.append("init-kv") + if failure_stage == "init": + raise RuntimeError("expected DCP profile failure") + + def prepare(): + events.append("prepare") + if failure_stage == "prepare": + raise RuntimeError("expected DCP profile failure") + monkeypatch.setattr( model_runner_module, "_init_minimal_kv_cache_for_profiling", - lambda _: events.append("init-kv"), + initialize, ) monkeypatch.setattr( model_runner_module, @@ -423,21 +434,29 @@ def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up( def dummy_run(*args, **kwargs): events.append(("dummy-run", args, kwargs)) - if dummy_run_fails: + if failure_stage == "forward": raise RuntimeError("expected DCP profile failure") return torch.empty(1), torch.empty(1) runner._dummy_run = dummy_run monkeypatch.setattr(torch.accelerator, "synchronize", lambda: events.append("sync")) - if dummy_run_fails: + if failure_stage is not None: with pytest.raises(RuntimeError, match="expected DCP profile failure"): - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) else: - runner.profile_glm_dcp_attention() + runner.profile_glm_dcp_attention(prepare) assert events[0] == "init-kv" - assert events[1] == ( + assert events[-1] == "cleanup" + if failure_stage == "init": + assert events == ["init-kv", "cleanup"] + return + assert events[1] == "prepare" + if failure_stage == "prepare": + assert events == ["init-kv", "prepare", "cleanup"] + return + assert events[2] == ( "dummy-run", (4096,), { @@ -448,8 +467,7 @@ def dummy_run(*args, **kwargs): "profile_all_kv_cache_groups": True, }, ) - assert events[-1] == "cleanup" - if not dummy_run_fails: + if failure_stage is None: assert events[-2] == "sync" diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index cda0764b458a..2133d437d7ee 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -326,10 +326,18 @@ def profile_cudagraph_memory(prepare_profile_state): events.append("profile_cudagraph_memory") return graph_estimate + def profile_run(prepare_profile_state): + prepare_profile_state() + events.append("profile_run") + + def profile_glm_dcp_attention(prepare_profile_state): + prepare_profile_state() + events.append("profile_glm_dcp_attention") + model_runner = SimpleNamespace( model_memory_usage=0, - profile_run=lambda prepare: (prepare(), events.append("profile_run")), - profile_glm_dcp_attention=lambda: events.append("profile_glm_dcp_attention"), + profile_run=profile_run, + profile_glm_dcp_attention=profile_glm_dcp_attention, profile_cudagraph_memory=profile_cudagraph_memory, ) profile_result = SimpleNamespace( @@ -408,7 +416,9 @@ def final_snapshot(**_kwargs): "prepare_b12x_profile_state", "profile_run", "release_b12x_profile_state", + "prepare_b12x_profile_state", "profile_glm_dcp_attention", + "release_b12x_profile_state", "prepare_b12x_profile_state", "profile_cudagraph_memory", "release_b12x_profile_state", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43da8b0e4646..0fb2540b8fb6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1094,15 +1094,17 @@ def _profile_deepseek_v4_attention( _teardown_profiling_state(self) @torch.inference_mode() - def profile_glm_dcp_attention(self) -> None: + def profile_glm_dcp_attention( + self, prepare_profile_state: Callable[[], None] | None = None + ) -> None: """Measure the GLM DCP query-gather peak before KV cache sizing. The general activation profile deliberately skips attention and splits its token budget across many requests. A GLM sparse-MLA prefill can put the complete scheduler budget in one request, causing the DCP query all-gather to require substantially more temporary memory. Bind a - minimal split cache, execute that shape, then release all temporary - cache and backend state before production cache allocation. + one-block split cache for that single request, prepare its bound kernel + plans, then release the temporary cache before production allocation. """ if ( self.model_config.architecture != "Glm5NextForConditionalGeneration" @@ -1110,8 +1112,10 @@ def profile_glm_dcp_attention(self) -> None: ): return - _init_minimal_kv_cache_for_profiling(self) try: + _init_minimal_kv_cache_for_profiling(self, num_blocks=1) + if prepare_profile_state is not None: + prepare_profile_state() self._dummy_run( self.max_num_tokens, context_len=self.dcp_size * self.cp_interleave, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1fc4022327a5..8e890bfc9e94 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6644,7 +6644,9 @@ def _profile_deepseek_v4_attention( self._cleanup_profiling_kv_cache() @torch.inference_mode() - def profile_glm_dcp_attention(self) -> None: + def profile_glm_dcp_attention( + self, prepare_profile_state: Callable[[], None] | None = None + ) -> None: """Profile GLM split-cache DCP attention before KV cache sizing. The generic activation profile omits attention metadata and spreads the @@ -6659,11 +6661,12 @@ def profile_glm_dcp_attention(self) -> None: ): return - with set_current_vllm_config(self.vllm_config): - self._init_minimal_kv_cache_for_profiling() - model_output: tuple[torch.Tensor, torch.Tensor] | None = None try: + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling(num_blocks=1) + if prepare_profile_state is not None: + prepare_profile_state() model_output = self._dummy_run( self.max_num_tokens, force_attention=True, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index e2af61a2e389..bc7d8949cf07 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -641,7 +641,12 @@ def determine_available_memory(self) -> int: self.model_runner.profile_run(self._prepare_b12x_profile_state) finally: self._release_b12x_profile_state() - self.model_runner.profile_glm_dcp_attention() + try: + self.model_runner.profile_glm_dcp_attention( + self._prepare_b12x_profile_state + ) + finally: + self._release_b12x_profile_state() # Profile CUDA graph memory if graphs will be captured. # ROCm is included: #44825 moved the profiler to From 515941cfe8640a96500bbe8c7b1e79346e9207d6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 16:59:50 +0000 Subject: [PATCH 105/365] fix(cudagraph): reserve retained native initialization once Measure non-Torch memory across the V2 capture span and reconcile it after backend-plan cleanup. Exclude only overlap already charged as late persistent memory, bounded by the measured graph delta. Bootstrap-only and cleanup-only allocations, transient native graph storage, activation headroom, and graph coverage remain reserved. Validation: 98 worker/profile tests pass, including native CUDA graph and cudaMalloc coverage. Eight worker overlap cases fail against the parent implementation. Ruff, formatting and type hooks pass. The CUDA API lint finding is unchanged at gpu_worker.py:592 in the parent revision; skip only that hook. Whole-model TP2 admission remains a separate serving gate. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 951279a0999d32c927b313b98d7e9bba34f4e8be) --- ...gpu_model_runner_v2_cudagraph_profiling.py | 80 +++++++++++++++++-- tests/v1/worker/test_gpu_worker.py | 25 +++++- vllm/v1/worker/gpu/cudagraph_utils.py | 15 ++++ vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu_worker.py | 31 +++++-- 5 files changed, 134 insertions(+), 18 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index 99f57ce4b20c..33a27e9272a8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py @@ -141,6 +141,7 @@ def _fake_set_current_vllm_config(_cfg): # retained; default to a constant (nothing retained). monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None) monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr(cgu.torch.accelerator, "memory_reserved", lambda: 0) monkeypatch.setattr( cgu.torch.accelerator, "get_memory_info", lambda: (1 << 30, 1 << 30) ) @@ -214,6 +215,49 @@ def test_profile_cudagraph_memory_piecewise_only_returns_measured(monkeypatch): assert result == captured_bytes +@pytest.mark.parametrize( + "init_native,capture_native,retained_native", + [ + (0, 256, 256), + (512, 256, 64), + (0, 256, 0), + (512, 0, 0), + (0, 0, 256), + (0, 256, 512), + ], +) +def test_profile_cudagraph_memory_records_native_growth_inside_capture_only( + monkeypatch, init_native, capture_native, retained_native +): + """Worker reconciles capture growth after all prepared plans are released.""" + _patch_module(monkeypatch) + mib = 1 << 20 + runner = _make_profiling_runner( + CUDAGraphMode.FULL_AND_PIECEWISE, + piecewise_only=True, + captured_bytes=1000 * mib, + ) + + def memory_info(): + native = init_native + if "teardown" in runner.events: + native += retained_native + elif "capture" in runner.events: + native += capture_native + # Torch-owned memory is not native initialization and is fully reserved. + return ((4096 - native - 512) * mib, 4096 * mib) + + monkeypatch.setattr(cgu.torch.accelerator, "get_memory_info", memory_info) + monkeypatch.setattr(cgu.torch.accelerator, "memory_reserved", lambda: 512 * mib) + + assert cgu.profile_cudagraph_memory(runner) == 1000 * mib + assert runner.cudagraph_native_memory_profile == ( + init_native * mib, + (init_native + capture_native) * mib, + 1000 * mib, + ) + + def test_profile_cudagraph_memory_tears_down_on_capture_error(monkeypatch): _patch_module(monkeypatch) runner = _make_profiling_runner(CUDAGraphMode.FULL) @@ -367,6 +411,8 @@ def test_profile_cudagraph_memory_clears_captured_graphs(monkeypatch): # CUDA graph executables must be destroyed and synchronized before their # B12X channel checkpoints and tensor workspaces are released. assert lifecycle == [ + "synchronize", # Native-memory measurement before capture. + "synchronize", # Native-memory measurement after capture. "synchronize", "reset-piecewise", "reset-breakable", @@ -495,8 +541,10 @@ def _capture_model(*, profile_only: bool = False) -> int: @create_new_process_for_each_test("spawn") @pytest.mark.skipif(not cgu.current_platform.is_cuda(), reason="requires CUDA") -def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch): - """Profiling graph memory must be freed before teardown completes.""" +@pytest.mark.parametrize("native_bytes", [0, 256 << 20]) +def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch, native_bytes): + """Release profiling graphs and avoid reserving persistent cudaMalloc twice.""" + from cuda.bindings import runtime @contextlib.contextmanager def _fake_set_current_vllm_config(_cfg): @@ -513,6 +561,7 @@ def _fake_set_current_vllm_config(_cfg): runner.maybe_remove_all_loras = lambda _: None allocation_bytes = 64 << 20 memory: dict[str, int] = {} + native_pointer = None torch.accelerator.synchronize() gc.collect() @@ -533,6 +582,11 @@ def _init(r): ) def _capture_model(*, profile_only: bool = False) -> int: + nonlocal native_pointer + before_free = torch.accelerator.get_memory_info()[0] + if native_bytes: + error, native_pointer = runtime.cudaMalloc(native_bytes) + assert error == runtime.cudaError_t.cudaSuccess for owner in ( runner.cudagraph_manager, runner.speculator.cudagraph_manager, @@ -545,17 +599,27 @@ def _capture_model(*, profile_only: bool = False) -> int: owner.graph_capture_resources["profile"] = [output] torch.accelerator.synchronize() memory["captured"] = torch.accelerator.memory_reserved() - return memory["captured"] - memory["before"] + memory["measured"] = before_free - torch.accelerator.get_memory_info()[0] + return memory["measured"] monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config) monkeypatch.setattr(cgu, "_init_minimal_kv_cache_for_profiling", _init) runner.capture_model = _capture_model - cgu.profile_cudagraph_memory(runner) - memory["after"] = torch.accelerator.memory_reserved() - - assert memory["captured"] - memory["before"] >= 3 * allocation_bytes - assert memory["after"] == memory["before"] + try: + graph_estimate = cgu.profile_cudagraph_memory(runner) + memory["after"] = torch.accelerator.memory_reserved() + + assert memory["captured"] - memory["before"] >= 3 * allocation_bytes + assert memory["after"] == memory["before"] + native_before, native_after, measured = runner.cudagraph_native_memory_profile + assert measured == graph_estimate == memory["measured"] + assert native_after - native_before >= native_bytes + finally: + if native_pointer is not None: + assert ( + runtime.cudaFree(native_pointer)[0] == runtime.cudaError_t.cudaSuccess + ) def test_teardown_profiling_state_clears_mamba_align_metadata(monkeypatch): diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 2133d437d7ee..b8ef5dc843fa 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -311,12 +311,27 @@ def run_model(scheduler_output, intermediate_tensors): ) @pytest.mark.parametrize("graph_estimate", [0, 4]) @pytest.mark.parametrize("estimate_graphs", [False, True]) +@pytest.mark.parametrize( + "native_profile,final_native,overlap", + [ + (None, 0, 0), # Runners without native capture measurements retain the budget. + ((100, 103, 4), 103, 3), # Persistent capture initialization. + ((100, 103, 4), 101, 1), # Prepared-plan release frees part of the growth. + ((100, 103, 4), 100, 0), # All native capture storage was temporary. + ((103, 103, 4), 107, 0), # Bootstrap and cleanup growth are not capture cost. + ((100, 103, 1), 107, 1), # The measured capture delta bounds the overlap. + ((100, 103, 4), 107, 3), # Cleanup-only growth cannot increase the discount. + ], +) def test_cudagraph_memory_profile_prepares_and_releases_b12x_state( monkeypatch, final_free_memory, expected_available_memory, graph_estimate, estimate_graphs, + native_profile, + final_native, + overlap, ): """KV admission counts retained allocations, not released profiling blocks.""" events: list[object] = [] @@ -335,6 +350,7 @@ def profile_glm_dcp_attention(prepare_profile_state): events.append("profile_glm_dcp_attention") model_runner = SimpleNamespace( + cudagraph_native_memory_profile=native_profile, model_memory_usage=0, profile_run=profile_run, profile_glm_dcp_attention=profile_glm_dcp_attention, @@ -381,7 +397,9 @@ def fake_memory_profiling(*args, **kwargs): def final_snapshot(**_kwargs): assert events[-3:] == ["collect", "synchronize", "empty_cache"] events.append("final_snapshot") - return SimpleNamespace(free_memory=final_free_memory) + return SimpleNamespace( + free_memory=final_free_memory, non_torch_memory=final_native + ) monkeypatch.setattr(gpu_worker.gc, "collect", lambda: events.append("collect")) monkeypatch.setattr( @@ -427,6 +445,7 @@ def final_snapshot(**_kwargs): "empty_cache", "final_snapshot", ] + graph_estimate -= min(overlap, graph_estimate, 90 - final_free_memory) applied_graph_estimate = graph_estimate if estimate_graphs else 0 assert available == expected_available_memory - applied_graph_estimate assert worker.peak_activation_memory == 5 @@ -488,8 +507,8 @@ def test_post_capture_recommendation_counts_measured_graph_memory_once( monkeypatch.setattr(gpu_worker, name, lambda *args: None) monkeypatch.setattr("vllm.utils.jit_monitor.activate", lambda **kwargs: None) - worker._compile_or_warm_up_model_after_preparation = ( - lambda: gpu_worker.Worker._compile_or_warm_up_model_after_preparation(worker) + worker._compile_or_warm_up_model_after_preparation = lambda: ( + gpu_worker.Worker._compile_or_warm_up_model_after_preparation(worker) ) worker._b12x_session = None worker._get_cudagraph_capture_context = nullcontext diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 7f5069b47f08..da2549011d80 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1181,6 +1181,7 @@ def profile_cudagraph_memory( partition reclaims the storages of earlier cudagraph recordings once the real capture records new ones, leading to use-after-free crashes). """ + runner.cudagraph_native_memory_profile = None if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: return 0 @@ -1253,7 +1254,14 @@ def profile_cudagraph_memory( mem_samples: list[int] = [] manager._capture_mem_samples = mem_samples + native_before_capture = _non_torch_memory_bytes() measured = int(runner.capture_model(profile_only=True)) + native_after_capture = _non_torch_memory_bytes() + runner.cudagraph_native_memory_profile = ( + native_before_capture, + native_after_capture, + measured, + ) # The measured delta covers PIECEWISE, encoder and speculator graphs # plus the sampled FULL graphs; swap the sampled FULL cost for the @@ -1303,6 +1311,13 @@ def profile_cudagraph_memory( platform_cls._global_graph_pool = saved_global_pool +def _non_torch_memory_bytes() -> int: + """Measure device memory not owned by the Torch caching allocator.""" + torch.accelerator.synchronize() + free, total = torch.accelerator.get_memory_info() + return total - free - torch.accelerator.memory_reserved() + + def _extrapolate_full_graph_memory( mem_samples: list[int], graph_descs: list[BatchExecutionDescriptor], diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0fb2540b8fb6..0038575c3096 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -201,6 +201,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): + self.cudagraph_native_memory_profile: tuple[int, int, int] | None = None self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index bc7d8949cf07..fa308b65fe87 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -669,13 +669,6 @@ def determine_available_memory(self) -> int: finally: self._release_b12x_profile_state() - # Respect the opt-in flag as originally designed. - cudagraph_memory_estimate_applied = ( - cudagraph_memory_estimate - if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS - else 0 - ) - # Prepared plans can retain profiling tensors beyond graph teardown. # Reclaim their released allocator blocks before measuring persistence. gc.collect() @@ -692,6 +685,30 @@ def determine_available_memory(self) -> int: - final_profile_snapshot.free_memory, 0, ) + native_profile = getattr( + self.model_runner, "cudagraph_native_memory_profile", None + ) + if native_profile is not None: + native_before, native_after, measured = native_profile + # Count capture-initialized modules/communication storage once. + # Sample retention only after prepared plans have been released; + # bootstrap-only and cleanup-only growth do not discount graphs. + retained_native_capture = max( + min(native_after, final_profile_snapshot.non_torch_memory) + - native_before, + 0, + ) + cudagraph_memory_estimate -= min( + retained_native_capture, + late_persistent_memory, + max(measured, 0), + cudagraph_memory_estimate, + ) + cudagraph_memory_estimate_applied = ( + cudagraph_memory_estimate + if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS + else 0 + ) init_free_memory = self.init_snapshot.free_memory free_gpu_memory = final_profile_snapshot.free_memory rocm_fallback = maybe_rocm_profiling_fallback(profile_result) From 1c90e4a7fb281cd08c18e85ccdc20d2ff41e872e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 17:35:57 +0000 Subject: [PATCH 106/365] fix(cache): reserve immutable boundary restore sources in KV capacity Count checkpoint source ownership and partial-tail copy-on-write alongside private endpoints when sizing request-boundary cache pools. Reject pools that admit a cold prompt but cannot restore its checkpoint; retain the existing aligned-policy budget and allocator invariants. Validation: 396 KV-cache and boundary tests pass, including the 62-block deadlock reproducer and 65-block successful restore. Five environment-dependent tests are rerun with their model configs and GPU visibility. Nine capacity cases reject the parent implementation. Full-model TP2 qualification is pending. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit 09e7a0b4aeaf0829ed06bb75b909efa8b5ae9652) --- .../core/test_boundary_admission_horizon.py | 80 ++++++++++++++++++- tests/v1/core/test_kv_cache_utils.py | 12 +-- vllm/v1/core/kv_cache_utils.py | 7 +- 3 files changed, 89 insertions(+), 10 deletions(-) diff --git a/tests/v1/core/test_boundary_admission_horizon.py b/tests/v1/core/test_boundary_admission_horizon.py index 18faf1c33b79..3a789f0679a8 100644 --- a/tests/v1/core/test_boundary_admission_horizon.py +++ b/tests/v1/core/test_boundary_admission_horizon.py @@ -2,18 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Per-reader attention horizons, low-pressure concurrency, and release fences.""" +from types import SimpleNamespace + import pytest import torch from tests.v1.core import test_boundary_admission as base from tests.v1.core.test_boundary_admission import initialize_hash as initialize_hash +from vllm.v1.core.kv_cache_utils import ( + _estimate_max_model_len_from_groups, + _pool_bytes_per_block, +) from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, MambaSpec, MLAAttentionSpec, ) -from vllm.v1.request import RequestStatus +from vllm.v1.request import Request, RequestStatus pytestmark = pytest.mark.cpu_test @@ -89,12 +95,82 @@ def restore(cache, req, lookahead=4, defer=True, pending=()): return allocated +@pytest.mark.parametrize("num_blocks", [62, 65, 80]) +def test_small_pool_capacity_covers_instruction_checkpoint_restore(num_blocks): + """Accepted pool geometry must admit both cold prompts and their restores.""" + attention = MLAAttentionSpec( + block_size=2048, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + model_version="glm5_next", + ) + recurrent = MambaSpec( + block_size=256, + shapes=((1,),), + dtypes=(torch.uint8,), + mamba_cache_mode="align", + num_speculative_blocks=3, + ) + groups = [KVCacheGroupSpec([str(i)], recurrent) for i in range(6)] + groups.append(KVCacheGroupSpec(["attention"], attention)) + config = SimpleNamespace( + use_request_boundary_checkpoints=True, + attention_config=SimpleNamespace(hisparse_config=None), + model_config=SimpleNamespace(max_model_len=1048576), + parallel_config=SimpleNamespace(decode_context_parallel_size=2), + cache_config=SimpleNamespace(mamba_cache_mode="align"), + ) + max_model_len = _estimate_max_model_len_from_groups( + config, groups, (num_blocks - 1) * _pool_bytes_per_block(groups) + ) + if num_blocks == 62: + # Cold admission fits, but restore needs 64 non-null blocks. Startup + # must reject this pool instead of leaving an idle scheduler waiting. + assert max_model_len == 0 + return + assert max_model_len >= 6357 + cache = base.make_kv_cache_manager( + KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ), + max_model_len=max_model_len, + max_in_flight_tokens=8192, + enable_caching=True, + use_eagle=True, + num_prefill_lookahead=1, + dcp_world_size=2, + scheduler_block_size=4096, + hash_block_size=256, + enable_boundary_checkpoints=True, + ) + producer = request("producer", length=6357, cap=1024) + producer.recurrent_instruction_boundary = 6345 + cache.get_computed_blocks(producer) + for count in (4096, 2249): + cache.new_step_starts() + assert cache.allocate_slots(producer, count, num_lookahead_tokens=3) + producer.num_computed_tokens += count + base.drain(cache) + checkpoint = cache.publish_boundary_checkpoint(producer, 6345, kind="instruction") + assert checkpoint is not None + cache.free(producer) + assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 + + reader = request("reader", length=6357, cap=1024) + reader.recurrent_instruction_boundary = 6345 + assert restore(cache, reader, lookahead=3) is not None + base.drain(cache) + cache.free(reader) + assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 + + def test_all_16_short_warm_readers_admit_without_checkpoint_loss(): cache = manager() for index in range(16): seed(cache, str(index)) published = set(cache.boundary_checkpoints._entries) - active = [] + active: list[Request] = [] for index in range(16): req = request("warm-" + str(index), str(index)) allocated = restore(cache, req) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 76fc5f2945c5..7d6957120049 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2669,7 +2669,7 @@ def test_balanced_glm_boundary_capacity_includes_private_endpoints( block_bytes = kv_cache_utils._pool_bytes_per_block(groups) live_bytes = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) live_blocks = live_bytes // block_bytes - private_blocks = 3 * (len(groups) + 1) + private_blocks = 4 * (len(groups) + 1) monkeypatch.setattr( VllmConfig, "use_request_boundary_checkpoints", property(lambda self: True) ) @@ -4277,9 +4277,9 @@ def test_auto_fit_max_model_len_with_hybrid(): @pytest.mark.parametrize("dcp", [1, 2, 4]) -@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (70, 15)]) +@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (62, 0), (70, 7)]) def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_pages): - """Live MTP state alone does not cover instruction/prompt/response copies.""" + """Pool sizing includes private endpoints and the immutable restore source.""" config = SimpleNamespace( use_request_boundary_checkpoints=True, attention_config=SimpleNamespace(hisparse_config=None), @@ -4304,8 +4304,8 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ groups = [KVCacheGroupSpec([f"state-{i}"], state) for i in range(6)] groups.append(KVCacheGroupSpec(["attention"], attention)) block_bytes = kv_cache_utils._pool_bytes_per_block(groups) - # Six groups require five live/scratch states each; three endpoints - # additionally own eight blocks each. The null block is unavailable. + # Six groups require five live/scratch states each; three private endpoints + # and a restore source own eight blocks each. The null block is unavailable. available = (num_blocks - 1) * block_bytes estimate = kv_cache_utils._estimate_max_model_len_from_groups( config, groups, available @@ -4329,7 +4329,7 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups ) assert get_max_concurrency_for_kv_cache_config(config, cache) == ( - num_blocks / (30 + 24 + expected_pages) + num_blocks / (30 + 32 + expected_pages) ) # An aligned-policy control has no endpoint reservation. Its capacity diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 567cab0e4628..85f5c062ec82 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1077,7 +1077,7 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) -> int: - """Count private instruction, prompt, and response checkpoint storage.""" + """Count private endpoints and the immutable checkpoint used by a restore.""" if not num_groups or not vllm_config.use_request_boundary_checkpoints: return 0 @@ -1085,7 +1085,10 @@ def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) - # KVCacheManager.allocate_slots reserves one block per group and one # auxiliary block per endpoint, separately from the live sequence state. - return NUM_BOUNDARY_CHECKPOINT_SLOTS * (num_groups + 1) + # Restoring a checkpoint pins its source until the reader finishes, while + # partial-tail CoW creates writable state. One additional bundle covers + # these source pages and auxiliary state alongside the private endpoints. + return (NUM_BOUNDARY_CHECKPOINT_SLOTS + 1) * (num_groups + 1) def get_max_concurrency_for_kv_cache_config( From 05f5ecd928055aa329b585f9ec7e54593e76f442 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 17:39:31 +0000 Subject: [PATCH 107/365] test(cache): cover asynchronous decode at minimum restore capacity Exercise 128 MTP decode steps after a boundary restore in the smallest accepted shared pool, with copy-release fences and two in-flight batches. Assert complete block ownership release after the request. Three capacity cases pass; serving source is unchanged. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit c1ef43814da82f4489dacc8b5e6fb50eb2272f40) --- .../v1/core/test_boundary_admission_horizon.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/v1/core/test_boundary_admission_horizon.py b/tests/v1/core/test_boundary_admission_horizon.py index 3a789f0679a8..b8a3b099fffa 100644 --- a/tests/v1/core/test_boundary_admission_horizon.py +++ b/tests/v1/core/test_boundary_admission_horizon.py @@ -10,6 +10,7 @@ from tests.v1.core import test_boundary_admission as base from tests.v1.core.test_boundary_admission import initialize_hash as initialize_hash from vllm.v1.core.kv_cache_utils import ( + KVCacheBlock, _estimate_max_model_len_from_groups, _pool_bytes_per_block, ) @@ -161,7 +162,24 @@ def test_small_pool_capacity_covers_instruction_checkpoint_restore(num_blocks): reader.recurrent_instruction_boundary = 6345 assert restore(cache, reader, lookahead=3) is not None base.drain(cache) + reader.append_output_token_ids([10000]) + pending = False + held: list[KVCacheBlock] = [] + for step in range(128): + cache.new_step_starts() + assert cache.allocate_slots(reader, 4, num_lookahead_tokens=3) is not None + reader.num_computed_tokens += 4 + reader.num_in_flight_tokens += 4 + copies = cache.take_kv_cache_block_copies()[1] + if pending: + reader.num_in_flight_tokens -= 4 + reader.append_output_token_ids([10001 + step] * 4) + cache.block_pool.free_blocks(held) + held = copies + pending = True + cache.cache_blocks(reader, min(reader.num_computed_tokens, reader.num_tokens)) cache.free(reader) + cache.block_pool.free_blocks(held) assert cache.block_pool.get_num_free_blocks() == num_blocks - 1 From 23622426046054e43c8f7fa91bc28a439b3d02f8 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 11 Sep 2026 23:46:40 +0000 Subject: [PATCH 108/365] Share GLM pooled-indexer temporary storage across sequential target layers Keep query, scale, score and block-table buffers under one model-owned scratch module. Preserve full token-row capacity, per-layer recurrent tails and independent MTP storage. Rebinding equal table geometry preserves graph addresses. Five focused scratch/capacity tests pass, including live-input CUDA graph replay and independent recurrent/draft ownership. Existing model and attention suites are recorded in local qualification artifacts. Co-authored-by: Codex (cherry picked from commit d5b19d25eeedec5a1e127531d39938e11de364e2) Signed-off-by: Martin Vit (cherry picked from commit b47c3d049a0b9d0a6ba893c7a18e5c4fa008167a) --- tests/models/test_glm5next_pooled_indexer.py | 154 +++++++++++++++++- vllm/models/glm5next/nvidia/attention.py | 4 +- vllm/models/glm5next/nvidia/model.py | 11 +- vllm/models/glm5next/nvidia/pooled_indexer.py | 152 ++++++++--------- 4 files changed, 236 insertions(+), 85 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 5b6d3ca555bc..21248d24b3b2 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -10,6 +10,7 @@ from torch import nn from vllm.models.deepseek_v4.nvidia.b12x_indexer import _flatten_index_cache +from vllm.models.glm5next.nvidia import pooled_indexer as indexer_module from vllm.models.glm5next.nvidia.ops import glm_kpool from vllm.models.glm5next.nvidia.ops.glm_kpool import ( expand_c4_block_table, @@ -19,7 +20,10 @@ prepare_c4_decode_metadata, update_decode_pools, ) -from vllm.models.glm5next.nvidia.pooled_indexer import Glm5NextPooledIndexer +from vllm.models.glm5next.nvidia.pooled_indexer import ( + Glm5NextIndexerScratch, + Glm5NextPooledIndexer, +) from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( @@ -659,7 +663,7 @@ def test_glm53_physical_selection_provider_is_explicit() -> None: indexer.dcp_world_size = 1 indexer._emit_physical_selection = True indexer.topk_indices_buffer = torch.empty((8, 2051), dtype=torch.int32) - indexer._physical_active_counts = torch.empty(8, dtype=torch.int32) + indexer.scratch = Glm5NextIndexerScratch(8, 1, torch.device("cpu")) selected = indexer.get_b12x_physical_selection( num_tokens=3, @@ -718,8 +722,7 @@ def test_glm53_missing_physical_plan_does_not_mutate_recurrent_cache( index_kpool_compress_gate=torch.zeros((128, 128)), index_kpool_compress_ape=torch.zeros((4, 128)), _project_head_weights=lambda value: torch.zeros((1, 32)), - _q_fp8=torch.empty((1, 32, 128), dtype=torch.float8_e4m3fn), - _q_scale=torch.empty((1, 32)), + scratch=Glm5NextIndexerScratch(1, 1, torch.device("cpu")), main_layer_name="attention", _index_cache=torch.zeros((1, 64, 132), dtype=torch.uint8), _tail=torch.zeros((1, 2, 4, 128)), @@ -815,7 +818,7 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( indexer._physical_selection_plan = object() indexer.bind_main_kv_cache(main) - assert indexer.indexer_op._num_q_heads == indexer._q_fp8.shape[1] == 32 + assert indexer.indexer_op._num_q_heads == indexer.scratch.q_fp8.shape[1] == 32 assert indexer._physical_selection_plan is None indexer._physical_selection_plan = object() indexer.unbind_main_kv_cache() @@ -862,14 +865,19 @@ def publish(*args, **kwargs): indexer.indexer_op = SimpleNamespace( max_model_len=4096 // 4, set_b12x_index_cache=publish ) + indexer.scratch = Glm5NextIndexerScratch( + indexer.max_tokens, indexer.max_seqs, device + ) indexer.bind_main_kv_cache(main) - assert indexer._decode_block_table.shape[0] == indexer.max_tokens - assert indexer._decode_block_table.shape[0] > indexer.max_seqs * (5 + 1) + assert indexer.scratch.decode_block_table.shape[0] == indexer.max_tokens + assert indexer.scratch.decode_block_table.shape[0] > indexer.max_seqs * (5 + 1) assert indexer.indexer_op.max_model_len == 4096 // 4 assert published["num_q_heads"] == 32 - assert published["max_page_table_width"] == indexer._pool_block_table.shape[1] + assert ( + published["max_page_table_width"] == indexer.scratch.pool_block_table.shape[1] + ) def test_b12x_c4_declaration_matches_glm_query_and_page_table_layout() -> None: @@ -1025,7 +1033,8 @@ def test_glm53_physical_selection_prepares_before_resolution_freeze( indexer.max_tokens = 32 indexer.prefix = "model.layers.3.indexer" indexer.topk_indices_buffer = output - indexer._physical_active_counts = counts + indexer.scratch = Glm5NextIndexerScratch(32, 1, device) + indexer.scratch.physical_active_counts = counts request = indexer.get_b12x_physical_selection_preparation_request() session = PreparationSession(device=device, autotune=autotune, compile_workers=0) session.prepare((request,)) @@ -1072,6 +1081,133 @@ def test_glm53_selector_capacity_tracks_auto_fit_max_model_len() -> None: assert indexer.indexer_op.max_model_len == 497 +def test_glm53_indexer_scratch_rebind_preserves_same_geometry_storage() -> None: + scratch = Glm5NextIndexerScratch(128, 4, torch.device("cpu")) + scratch.bind_tables(32, torch.device("cpu")) + pointers = {name: buffer.data_ptr() for name, buffer in scratch.named_buffers()} + scratch.decode_block_table.fill_(53) + scratch.bind_tables(32, torch.device("cpu")) + assert { + name: buffer.data_ptr() for name, buffer in scratch.named_buffers() + } == pointers + assert torch.all(scratch.decode_block_table == 53) + scratch.bind_tables(16, torch.device("cpu")) + assert scratch.decode_block_table.shape == (128, 16) + assert scratch.pool_block_table.shape == (4, 16) + assert scratch.q_fp8.data_ptr() == pointers["q_fp8"] + assert not scratch.state_dict() + + +def test_glm53_shared_scratch_keeps_layer_checkpoints_and_draft_selection_private( + monkeypatch, +) -> None: + monkeypatch.setattr( + indexer_module, "ReplicatedLinear", lambda *a, **k: nn.Identity() + ) + monkeypatch.setattr( + indexer_module, "B12xC4SparseIndexer", lambda *a, **k: SimpleNamespace() + ) + monkeypatch.setattr( + indexer_module, + "get_b12x_sparse_mla", + lambda: SimpleNamespace(expand_pooled_topk_to_physical_slots=lambda *a: None), + ) + config = SimpleNamespace( + index_topk=2048, + index_n_heads=32, + index_head_dim=128, + index_kpool=4, + qk_rope_head_dim=0, + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=128, max_num_seqs=4), + model_config=SimpleNamespace(max_model_len=4096), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, cp_kv_cache_interleave_size=4 + ), + speculative_config=SimpleNamespace(num_speculative_tokens=3), + ) + scratch = Glm5NextIndexerScratch(128, 4, torch.device("cpu")) + target_ids = torch.zeros((128, 2051), dtype=torch.int32) + target_pools = torch.zeros((128, 512), dtype=torch.int32) + + def make(name, shared): + return Glm5NextPooledIndexer( + vllm_config, + config, + 8, + 8, + None, + SimpleNamespace(block_size=2048), + target_ids if shared else torch.full_like(target_ids, 42), + target_pools if shared else torch.full_like(target_pools, 42), + main_layer_name=name, + prefix=name, + scratch=scratch if shared else None, + ) + + first, second, draft = ( + make("target.0", True), + make("target.1", True), + make("mtp", False), + ) + assert first.scratch is second.scratch + assert draft.scratch is not scratch + for indexer, value in ((first, 1), (second, 2), (draft, 3)): + indexer._tail.fill_(value) + indexer.snapshot_speculative_interval_starts() + scratch.q_fp8.view(torch.uint8).fill_(255) + target_ids.fill_(-1) + first._tail.zero_() + first.restore_speculative_interval_starts() + for indexer, value in ((first, 1), (second, 2), (draft, 3)): + assert torch.all(indexer._tail == value) + assert torch.all(indexer._tail_snapshot == value) + assert torch.all(draft.topk_indices_buffer == 42) + assert torch.all(draft.pool_topk_indices_buffer == 42) + + +def test_glm53_shared_indexer_scratch_graph_consumes_each_live_query() -> None: + device = _require_glm_gpu() + scratch = Glm5NextIndexerScratch(8, 1, device) + inputs = [ + torch.randn((8, 32, 128), device=device, dtype=torch.bfloat16) for _ in range(2) + ] + outputs = [torch.empty_like(scratch.q_fp8) for _ in range(2)] + scales = [torch.empty_like(scratch.q_scale) for _ in range(2)] + + def run() -> None: + for query, output, scale in zip(inputs, outputs, scales): + glm_kpool.fwht128_quant_fp8( + query.view(-1, 128), + scratch.q_fp8.view(-1, 128), + scratch.q_scale.view(-1), + ) + output.copy_(scratch.q_fp8) + scale.copy_(scratch.q_scale) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for _ in range(3): + for query in inputs: + query.normal_() + scratch.q_fp8.view(torch.uint8).fill_(255) + scratch.q_scale.fill_(float("nan")) + graph.replay() + for query, output, scale in zip(inputs, outputs, scales): + reference_q = torch.empty_like(output) + reference_scale = torch.empty_like(scale) + glm_kpool.fwht128_quant_fp8( + query.view(-1, 128), + reference_q.view(-1, 128), + reference_scale.view(-1), + ) + torch.testing.assert_close(output, reference_q, rtol=0, atol=0) + torch.testing.assert_close(scale, reference_scale, rtol=0, atol=0) + + def test_glm53_parent_table_width_tracks_dcp_sharding() -> None: max_model_len = 524288 block_size = 2304 diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py index f3a093a14ff5..fe8714b7c0e6 100644 --- a/vllm/models/glm5next/nvidia/attention.py +++ b/vllm/models/glm5next/nvidia/attention.py @@ -22,7 +22,7 @@ from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.registry import AttentionBackendEnum -from .pooled_indexer import Glm5NextPooledIndexer +from .pooled_indexer import Glm5NextIndexerScratch, Glm5NextPooledIndexer def _select_sparse_backend( @@ -63,6 +63,7 @@ def __init__( skip_rope: bool | None = False, is_mtp_layer: bool = False, attn_backend: type | None = None, + indexer_scratch: Glm5NextIndexerScratch | None = None, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -171,6 +172,7 @@ def __init__( prefix=f"{prefix}.indexer", # MTP compacts and reuses request-relative selections. emit_physical_selection=not is_mtp_layer, + scratch=indexer_scratch, ) else: self.indexer = None diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index a1f5adecf675..8107a0d9d8f1 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -96,7 +96,7 @@ from . import l2_prefetch as _l2pf from .attention import Glm5NextMLAAttention from .kda import Glm5NextLinearAttention -from .pooled_indexer import Glm5NextPooledIndexer +from .pooled_indexer import Glm5NextIndexerScratch, Glm5NextPooledIndexer logger = init_logger(__name__) @@ -333,6 +333,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None = None, pool_topk_indices_buffer: torch.Tensor | None = None, is_mtp_layer: bool = False, + indexer_scratch: Glm5NextIndexerScratch | None = None, **kwargs, ) -> None: super().__init__() @@ -389,6 +390,7 @@ def __init__( pool_topk_indices_buffer=pool_topk_indices_buffer, skip_rope=getattr(config, "mla_nope", False), is_mtp_layer=is_mtp_layer, + indexer_scratch=indexer_scratch, ) # MTP layers sit past the base model's hidden layers (layer_idx >= @@ -960,10 +962,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dtype=torch.int32, device=self.device, ) + indexer_scratch = Glm5NextIndexerScratch( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.scheduler_config.max_num_seqs, + torch.device(self.device), + ) else: # Full-MLA config (no kpool sparse indexer): no topk buffer. topk_indices_buffer = None pool_topk_indices_buffer = None + indexer_scratch = None if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -983,6 +991,7 @@ def get_layer(prefix: str): prefix=prefix, topk_indices_buffer=topk_indices_buffer, pool_topk_indices_buffer=pool_topk_indices_buffer, + indexer_scratch=indexer_scratch, ) self.start_layer, self.end_layer, self.layers = make_layers( diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 63ee0d63546f..9e94d04e6455 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -51,6 +51,55 @@ _INDEX_PAGE_BYTES = _INDEX_PAGE_SIZE * _INDEX_CACHE_WIDTH +class Glm5NextIndexerScratch(nn.Module): + """Temporary storage for sequential pooled-indexer calls on one model stream. + + A model owns one instance; its draft model owns another. KV views, recurrent + tails and selected token IDs are not scratch and remain separately owned. + Bind tables before graph capture; calls must finish consuming the scratch + before another layer overwrites it. + """ + + q_fp8: torch.Tensor + q_scale: torch.Tensor + pool_scores: torch.Tensor + pool_seq_lens: torch.Tensor + physical_active_counts: torch.Tensor + pool_block_table: torch.Tensor + decode_block_table: torch.Tensor + + def __init__(self, max_tokens: int, max_seqs: int, device: torch.device) -> None: + super().__init__() + self.max_tokens = max_tokens + self.max_seqs = max_seqs + for name, shape, dtype in ( + ("q_fp8", (max_tokens, _INDEX_HEADS, _INDEX_HEAD_DIM), torch.float8_e4m3fn), + ("q_scale", (max_tokens, _INDEX_HEADS), torch.float32), + ("pool_scores", (max_tokens, _POOL_TOPK), torch.float32), + ("pool_seq_lens", (max_tokens,), torch.int32), + ("physical_active_counts", (max_tokens,), torch.int32), + ("pool_block_table", (max_seqs, 1), torch.int32), + ("decode_block_table", (max_tokens, 1), torch.int32), + ): + self.register_buffer( + name, torch.empty(shape, dtype=dtype, device=device), persistent=False + ) + + def bind_tables(self, width: int, device: torch.device) -> None: + if self.pool_block_table.shape[1] == width and self.q_fp8.device == device: + return + if self.q_fp8.device != device: + raise ValueError("GLM indexer scratch and KV cache must share a device") + self.pool_block_table = torch.empty( + (self.max_seqs, width), dtype=torch.int32, device=device + ) + # Short prefills can use decode metadata, so rows need the token budget, + # not merely max_seqs times the speculative width. + self.decode_block_table = torch.empty( + (self.max_tokens, width), dtype=torch.int32, device=device + ) + + class Glm5NextPooledIndexer(nn.Module): """Produce GLM C4 pools, select them with b12x, and expand token IDs.""" @@ -68,6 +117,7 @@ def __init__( main_layer_name: str, prefix: str, emit_physical_selection: bool = True, + scratch: Glm5NextIndexerScratch | None = None, ) -> None: super().__init__() if cache_config is None: @@ -132,6 +182,17 @@ def __init__( raise ValueError("GLM pool selection buffer has the wrong row capacity") device = topk_indices_buffer.device + self.scratch = ( + scratch + if scratch is not None + else Glm5NextIndexerScratch(self.max_tokens, self.max_seqs, device) + ) + if ( + self.scratch.max_tokens != self.max_tokens + or self.scratch.max_seqs != self.max_seqs + or self.scratch.q_fp8.device != device + ): + raise ValueError("GLM indexer scratch has incompatible capacity or device") b12x_sparse_mla = get_b12x_sparse_mla() if b12x_sparse_mla is None or not hasattr( b12x_sparse_mla, "expand_pooled_topk_to_physical_slots" @@ -219,57 +280,6 @@ def __init__( torch.empty_like(self._tail), persistent=False, ) - self.register_buffer( - "_q_fp8", - torch.empty( - (self.max_tokens, _INDEX_HEADS, _INDEX_HEAD_DIM), - dtype=torch.float8_e4m3fn, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_physical_active_counts", - torch.empty(self.max_tokens, dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_q_scale", - torch.empty( - (self.max_tokens, _INDEX_HEADS), - dtype=torch.float32, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_pool_seq_lens", - torch.empty(self.max_tokens, dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_pool_scores", - torch.empty( - (self.max_tokens, _POOL_TOPK), - dtype=torch.float32, - device=device, - ), - persistent=False, - ) - self.register_buffer( - "_pool_block_table", - torch.empty((self.max_seqs, 1), dtype=torch.int32, device=device), - persistent=False, - ) - self.register_buffer( - "_decode_block_table", - torch.empty( - (self.max_tokens, 1), - dtype=torch.int32, - device=device, - ), - persistent=False, - ) self._weights_proj_fp32: torch.Tensor | None = None self._index_cache: torch.Tensor | None = None self._parent_table_width = 0 @@ -370,15 +380,7 @@ def bind_main_kv_cache(self, main_cache: torch.Tensor) -> None: self.dcp_world_size, ) pool_table_width = parent_table_width * subpages - device = main_cache.device - self._pool_block_table = torch.empty( - (self.max_seqs, pool_table_width), dtype=torch.int32, device=device - ) - self._decode_block_table = torch.empty( - (self.max_tokens, pool_table_width), - dtype=torch.int32, - device=device, - ) + self.scratch.bind_tables(pool_table_width, main_cache.device) self._index_cache = index_cache self._parent_table_width = parent_table_width self._subpages_per_parent = subpages @@ -450,7 +452,7 @@ def get_b12x_physical_selection_preparation_request(self): def prepare(state): return self.make_b12x_physical_selection_prepare_call( torch.empty_like(self.topk_indices_buffer), - torch.empty_like(self._physical_active_counts), + torch.empty_like(self.scratch.physical_active_counts), state=state, ) @@ -485,7 +487,7 @@ def get_b12x_physical_selection( return None return ( self.topk_indices_buffer[:num_tokens], - self._physical_active_counts[:num_tokens], + self.scratch.physical_active_counts[:num_tokens], ) def _project_head_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -522,8 +524,8 @@ def forward( gate = F.linear(hidden_states, self.index_kpool_compress_gate) weights = self._project_head_weights(hidden_states) - q_fp8 = self._q_fp8[:rows] - q_scale = self._q_scale[:rows] + q_fp8 = self.scratch.q_fp8[:rows] + q_scale = self.scratch.q_scale[:rows] fwht128_quant_fp8( query.contiguous().view(-1, _INDEX_HEAD_DIM), q_fp8.view(-1, _INDEX_HEAD_DIM), @@ -595,8 +597,8 @@ def forward( parent_stride_pages=self._parent_stride_pages, ) parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] - seq_lens = self._pool_seq_lens[:live_rows] - decode_table = self._decode_block_table[:decode_rows] + seq_lens = self.scratch.pool_seq_lens[:live_rows] + decode_table = self.scratch.decode_block_table[:decode_rows] if decode_only: prepare_c4_decode_metadata( parent_table, @@ -613,7 +615,7 @@ def forward( else: expand_c4_block_table( parent_table, - self._pool_block_table, + self.scratch.pool_block_table, rows=num_reqs, subpages_per_parent=self._subpages_per_parent, parent_stride_pages=self._parent_stride_pages, @@ -626,12 +628,14 @@ def forward( pool_interleave=self.pool_interleave, ) pool_ids = self.pool_topk_indices_buffer[:rows] - pool_scores = self._pool_scores[:rows] if self.dcp_world_size > 1 else None + pool_scores = ( + self.scratch.pool_scores[:rows] if self.dcp_world_size > 1 else None + ) if decode_rows: if not decode_only: gather_c4_block_table_rows( - self._pool_block_table, + self.scratch.pool_block_table, main_metadata.req_id_per_token[:decode_rows], decode_table, ) @@ -671,19 +675,19 @@ def forward( active_pages = self._active_index_page_count( int(request_seq_lens_cpu[local_request]) ) - if active_pages > int(self._pool_block_table.shape[1]): + if active_pages > int(self.scratch.pool_block_table.shape[1]): raise RuntimeError( "GLM selector visible C4 pages exceed table capacity: " f"active={active_pages}, " - f"capacity={int(self._pool_block_table.shape[1])}" + f"capacity={int(self.scratch.pool_block_table.shape[1])}" ) - request_table = self._pool_block_table[ + request_table = self.scratch.pool_block_table[ request : request + 1, :active_pages ] else: # DCP pool ownership is interleaved across ranks, so a global # sequence length does not define a contiguous local prefix. - request_table = self._pool_block_table[request : request + 1] + request_table = self.scratch.pool_block_table[request : request + 1] shared_table = request_table.expand(int(query_len), -1) self.indexer_op.run_paged_topk( q=q_fp8[row_start:row_end], @@ -723,7 +727,7 @@ def forward( main_metadata.req_id_per_token[:live_rows], main_metadata.block_table, output[:live_rows], - self._physical_active_counts[:live_rows], + self.scratch.physical_active_counts[:live_rows], pool_size=_POOL_SIZE, block_size=self.block_size, block_stride_rows=self.block_size, From 0ce2e7badd9b65f18a22d15f30ec05d022f1a5bc Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:26:17 +0000 Subject: [PATCH 109/365] Reuse GLM indexer storage across serial target and MTP execution Port the storage-sharing part of PR #739 commit 7f4aecc66e857d093e012a2c57d03711bd23be4a to Karmic loading. Require one draft layer, matching geometry, PP1 and no overlapped microbatches. Keep parameters, recurrent tails and KV views independent. Eleven CPU/CUDA ownership tests pass, including live-input graph replay. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit a311102cd031b02b5fac9cfa3ed2dc2af01c767a) --- tests/models/test_glm5next_indexer_storage.py | 194 ++++++++++++++++++ vllm/models/glm5next/nvidia/mtp.py | 71 +++++++ vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 9 +- 3 files changed, 273 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_glm5next_indexer_storage.py diff --git a/tests/models/test_glm5next_indexer_storage.py b/tests/models/test_glm5next_indexer_storage.py new file mode 100644 index 000000000000..9b872bffb211 --- /dev/null +++ b/tests/models/test_glm5next_indexer_storage.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ownership checks for serial GLM target/MTP indexer storage.""" + +import weakref +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.glm5next.nvidia.mtp import Glm5NextMTP +from vllm.models.glm5next.nvidia.pooled_indexer import ( + Glm5NextIndexerScratch, + Glm5NextPooledIndexer, +) + + +def make_indexer(device="cpu", shared=None): + indexer = Glm5NextPooledIndexer.__new__(Glm5NextPooledIndexer) + nn.Module.__init__(indexer) + for name, value in dict( + max_tokens=16, + max_seqs=4, + max_model_len=1048576, + block_size=2048, + dcp_world_size=2, + dcp_rank=0, + pool_interleave=1, + ).items(): + setattr(indexer, name, value) + indexer.scratch = ( + shared.scratch + if shared + else Glm5NextIndexerScratch(16, 4, torch.device(device)) + ) + for name, width in ( + ("topk_indices_buffer", 2051), + ("pool_topk_indices_buffer", 512), + ): + value = ( + getattr(shared, name) + if shared + else torch.empty((16, width), dtype=torch.int32, device=device) + ) + setattr(indexer, name, value) + indexer.indexer_op = nn.Module() + indexer.indexer_op.topk_indices_buffer = indexer.pool_topk_indices_buffer + indexer._index_cache = torch.ones(32, dtype=torch.uint8, device=device) + indexer.register_buffer("_tail", torch.ones(32, device=device), persistent=False) + indexer.register_parameter("weight", nn.Parameter(torch.ones(4, device=device))) + return indexer + + +def make_pair(device="cpu"): + target = nn.Module() + first = make_indexer(device) + target.layers = nn.ModuleList([first, make_indexer(device, shared=first)]) + draft = Glm5NextMTP.__new__(Glm5NextMTP) + nn.Module.__init__(draft) + draft.model = nn.Module() + draft.model.num_mtp_layers = 1 + draft.model.indexer = make_indexer(device) + draft.model.attention = nn.Module() + draft.model.attention.topk_indices_buffer = draft.model.indexer.topk_indices_buffer + draft.model.attention.impl = SimpleNamespace( + topk_indices_buffer=draft.model.indexer.topk_indices_buffer + ) + return target, draft + + +def test_share_indexer_storage_preserves_roles_and_persistent_state(): + target, draft = make_pair() + source, indexer = target.layers[0], draft.model.indexer + scratch_ref = weakref.ref(indexer.scratch) + token_ref = weakref.ref(indexer.topk_indices_buffer) + pool_ref = weakref.ref(indexer.pool_topk_indices_buffer) + cache, tail, weight = indexer._index_cache, indexer._tail, indexer.weight + assert draft.share_target_indexer_storage(target) + assert scratch_ref() is token_ref() is pool_ref() is None + assert indexer.scratch is source.scratch + assert indexer.topk_indices_buffer is source.topk_indices_buffer + assert indexer.pool_topk_indices_buffer is source.pool_topk_indices_buffer + assert indexer.indexer_op.topk_indices_buffer is source.pool_topk_indices_buffer + assert draft.model.attention.topk_indices_buffer is source.topk_indices_buffer + assert draft.model.attention.impl.topk_indices_buffer is source.topk_indices_buffer + assert indexer._index_cache is cache + assert indexer._tail is tail + assert indexer.weight is weight + source.scratch.bind_tables(2048, torch.device("cpu")) + pointer = source.scratch.decode_block_table.data_ptr() + indexer.scratch.bind_tables(2048, torch.device("cpu")) + assert indexer.scratch.decode_block_table.data_ptr() == pointer + assert draft.share_target_indexer_storage(target) # Idempotent ownership. + + +@pytest.mark.parametrize( + "difference", + [ + "layers", + "geometry", + "capacity", + "dtype", + "scratch", + "target_storage", + ], +) +def test_incompatible_indexer_storage_remains_independent(difference): + target, draft = make_pair() + indexer = draft.model.indexer + if difference == "layers": + draft.model.num_mtp_layers = 2 + elif difference == "geometry": + indexer.dcp_world_size = 1 + elif difference == "capacity": + indexer.topk_indices_buffer = indexer.topk_indices_buffer[:8] + elif difference == "dtype": + indexer.pool_topk_indices_buffer = indexer.pool_topk_indices_buffer.long() + elif difference == "scratch": + indexer.scratch.bind_tables(2, torch.device("cpu")) + else: + target.layers[1].scratch = Glm5NextIndexerScratch(16, 4, torch.device("cpu")) + scratch, tokens, pools = ( + indexer.scratch, + indexer.topk_indices_buffer, + indexer.pool_topk_indices_buffer, + ) + assert not draft.share_target_indexer_storage(target) + assert indexer.scratch is scratch + assert indexer.topk_indices_buffer is tokens + assert indexer.pool_topk_indices_buffer is pools + + +@pytest.mark.parametrize( + "pp,dbo,expected", [(1, False, True), (2, False, False), (1, True, False)] +) +def test_loader_shares_only_serial_indexers(monkeypatch, pp, dbo, expected): + from vllm.v1.worker.gpu.spec_decode.eagle import utils + + inner, draft = make_pair() + target = nn.Module() + target.model = inner + config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=None, + moe_backend=None, + kv_cache_dtype=None, + attention_backend=None, + ), + parallel_config=SimpleNamespace(enable_dbo=dbo), + load_config=SimpleNamespace(load_format="auto"), + ) + monkeypatch.setattr(utils, "_make_eagle_draft_vllm_config", lambda value: value) + monkeypatch.setattr(utils, "get_model", lambda **_: draft) + monkeypatch.setattr(utils, "get_pp_group", lambda: SimpleNamespace(world_size=pp)) + monkeypatch.setattr(utils, "get_pp_safe_draft_load_config", lambda value: value) + assert utils.load_eagle_model(target, config) is draft + assert (draft.model.indexer.scratch is inner.layers[0].scratch) is expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph ownership") +def test_shared_indexer_storage_serial_graph_replay(): + target, draft = make_pair("cuda") + assert draft.share_target_indexer_storage(target) + source, indexer = target.layers[0], draft.model.indexer + value = torch.ones((), device="cuda") + + def run(owner, factor): + owner.scratch.pool_scores.copy_(value * factor) + owner.topk_indices_buffer.fill_(factor) + owner.pool_topk_indices_buffer.fill_(factor + 1) + return ( + owner.scratch.pool_scores.clone(), + owner.topk_indices_buffer.clone(), + owner.pool_topk_indices_buffer.clone(), + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run(source, 2) + run(indexer, 3) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + target_result = run(source, 2) + draft_result = run(indexer, 3) + for scalar in (1, 5, 9): + value.fill_(scalar) + graph.replay() + for result, factor in ((target_result, 2), (draft_result, 3)): + assert torch.equal(result[0], torch.full_like(result[0], scalar * factor)) + assert torch.equal(result[1], torch.full_like(result[1], factor)) + assert torch.equal(result[2], torch.full_like(result[2], factor + 1)) diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index 3bf518b5d636..423a3416f4b6 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -336,6 +336,77 @@ def prepare_draft_lm_head(self, source_head: nn.Module) -> None: """Create a draft-only quantized copy of the shared target head.""" self.model.prepare_draft_lm_head(source_head) + def share_target_indexer_storage(self, target: nn.Module) -> bool: + """Reuse temporary indexer storage under serial target/MTP execution. + + The V2 loader calls this before profiling or graph capture, without + pipeline parallelism or overlapped microbatches. One MTP layer consumes + its selections before target execution resumes. KV views, pool tails + and model parameters retain independent ownership. + """ + if self.model.num_mtp_layers != 1: + return False + targets = [ + module + for module in target.modules() + if isinstance(module, Glm5NextPooledIndexer) + ] + drafts = [ + module + for module in self.model.modules() + if isinstance(module, Glm5NextPooledIndexer) + ] + if not targets or len(drafts) != 1: + return False + source, draft = targets[0], drafts[0] + geometry = ( + "max_tokens", + "max_seqs", + "max_model_len", + "block_size", + "dcp_world_size", + "dcp_rank", + "pool_interleave", + ) + if any(getattr(source, key) != getattr(draft, key) for key in geometry): + return False + names = ("topk_indices_buffer", "pool_topk_indices_buffer") + for module in targets: + if module.scratch is not source.scratch or any( + getattr(module, name) is not getattr(source, name) for name in names + ): + return False + for name in names: + src, dst = getattr(source, name), getattr(draft, name) + if (src.shape, src.dtype, src.device) != (dst.shape, dst.dtype, dst.device): + return False + if any( + (src.shape, src.dtype, src.device) != (dst.shape, dst.dtype, dst.device) + for (_, src), (_, dst) in zip( + source.scratch.named_buffers(), + draft.scratch.named_buffers(), + strict=True, + ) + ): + return False + + replacements = { + id(getattr(draft, name)): getattr(source, name) for name in names + } + # The native indexer holds 512 pool IDs; MLA consumes 2051 token IDs. + # Rebind by identity, not by attribute name, including non-Module impls. + for module in self.model.modules(): + for owner in (module, getattr(module, "impl", None)): + if owner is None: + continue + for name in names: + value = getattr(owner, name, None) + replacement = replacements.get(id(value)) + if replacement is not None: + setattr(owner, name, replacement) + draft.scratch = source.scratch + return True + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 3d42451a166a..89ed39d14f2b 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -182,7 +182,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod # every module in the draft that holds a buffer reference so that # the per-layer indexer and sparse-attention backends all point to # the target's buffer. - if hasattr(target_inner, "topk_indices_buffer"): + share_indexer_storage = getattr(eagle_model, "share_target_indexer_storage", None) + if share_indexer_storage is not None: + if ( + get_pp_group().world_size == 1 + and not vllm_config.parallel_config.enable_dbo + ): + share_indexer_storage(target_inner) + elif hasattr(target_inner, "topk_indices_buffer"): target_buffer = target_inner.topk_indices_buffer if target_buffer is not None: for _, module in draft_inner.named_modules(): From 4135bd116cb1395868b0964829e9ff7e3c7b7806 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 18 Sep 2026 12:37:31 +0000 Subject: [PATCH 110/365] Bound GLM vision projection memory without partitioning attention Port the vision and workspace portions of 7f4aecc66e857 and the rotary-bounded profiling canvas from ab51db9b4801a551e5c5fc4a6e0b4e20fb37bbf4 to the common GLM model module. Preserve the Karmic processor and execution-lane contracts. Token-local projections use bounded chunks; image attention retains complete Q/K/V and sequence boundaries. Validation: 30 focused tests pass. Native 24-layer BF16 vision on TP2 produces bit-identical output on both ranks; additional peak allocation falls from 662164480 to 221254656 bytes, or 179295232 bytes when borrowing preallocated worker storage. Whole-model serving validation is separate. Co-authored-by: Codex Signed-off-by: Martin Vit (cherry picked from commit adc85c66df33e94f4c8e6eb704d1e7d3cfbbfb6b) --- tests/models/test_glm5next_model.py | 162 +++++++++++++++++++ tests/v1/worker/test_workspace.py | 4 + vllm/models/glm5next/common/multimodal.py | 184 +++++++++++++++++++++- vllm/v1/worker/workspace.py | 17 ++ 4 files changed, 360 insertions(+), 7 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 39e845f70a04..cf713d199242 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -5,6 +5,7 @@ import math import sys import types +import weakref from dataclasses import dataclass from types import SimpleNamespace @@ -2339,6 +2340,167 @@ def fake_from_pretrained(model, **kwargs): ] +def test_glm5next_vision_attention_releases_projection_and_norm_inputs(monkeypatch): + from vllm.models.glm5next.common import multimodal + + owners = {} + + def project(value): + packed = value.repeat(1, 1, 3) + owners["packed"] = weakref.ref(packed) + return packed, None + + def normalize(q, k, *args): + owners["q_norm_input"] = weakref.ref(q) + owners["k_norm_input"] = weakref.ref(k) + return q.clone(), k.clone() + + def attention(query, key, value, **kwargs): + assert all(owner() is None for owner in owners.values()) + owners.update(q=weakref.ref(query), k=weakref.ref(key), v=weakref.ref(value)) + return value.clone() + + def output_project(value): + assert all(owner() is None for owner in owners.values()) + return value, None + + layer = SimpleNamespace( + head_dim=2, + num_attention_heads_per_partition=2, + hidden_size_per_attention_head=2, + qkv=project, + q_norm=SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + k_norm=SimpleNamespace(weight=torch.ones(2)), + apply_rotary_emb=lambda value, cos, sin: value.clone(), + attn=attention, + proj=output_project, + ) + layer.split_qkv = lambda value: multimodal.Glm5NextVisionAttention.split_qkv( + layer, value + ) + layer._project_qkv = lambda *args: multimodal.Glm5NextVisionAttention._project_qkv( + layer, *args + ) + monkeypatch.setattr(multimodal, "fused_q_kv_rmsnorm", normalize) + x = torch.arange(28, dtype=torch.float32).view(7, 1, 4) + original = x.clone() + output = multimodal.Glm5NextVisionAttention.forward( + layer, x, torch.tensor([0, 7]), torch.ones(7, 1), torch.zeros(7, 1), 7 + ) + torch.testing.assert_close(output, original, rtol=0, atol=0) + torch.testing.assert_close(x, original, rtol=0, atol=0) + + +@pytest.mark.parametrize("reserved_bytes", [0, 256, 4096]) +def test_glm5next_chunked_projection_keeps_whole_image_attention( + monkeypatch, reserved_bytes +): + from vllm.models.glm5next.common import multimodal + from vllm.v1.worker import workspace + + manager = workspace.WorkspaceManager(torch.device("cpu")) + manager.reserve_all(((reserved_bytes,), torch.uint8)) + manager.lock() + monkeypatch.setattr(workspace, "_manager", manager) + monkeypatch.setattr(multimodal, "_VISION_PROJECTION_CHUNK_ROWS", 4) + x = torch.randn(11, 1, 8, generator=torch.Generator().manual_seed(72)) + cu_seqlens = torch.tensor([0, 3, 11]) + cos, sin = torch.ones(11, 2), torch.zeros(11, 2) + projected_rows = [] + attention_calls = [] + + def project(value, cos, sin): + projected_rows.append(len(value)) + assert cos.shape[0] == sin.shape[0] == len(value) + value = value.reshape(1, len(value), 2, 4) + return value, value * 0.5, value * 2 + + def attention(query, key, value, **metadata): + assert metadata["cu_seqlens"] is cu_seqlens + attention_calls.append(query.shape[1]) + outputs = [] + for start, stop in zip(cu_seqlens[:-1], cu_seqlens[1:]): + q, k, v = (t[:, start:stop].transpose(1, 2) for t in (query, key, value)) + outputs.append( + torch.nn.functional.scaled_dot_product_attention(q, k, v).transpose( + 1, 2 + ) + ) + return torch.cat(outputs, dim=1) + + layer = SimpleNamespace( + num_attention_heads_per_partition=2, + hidden_size_per_attention_head=4, + _project_qkv=project, + _borrow_qkv_workspace=True, + attn=attention, + ) + norm = torch.nn.RMSNorm(8) + q, k, v = project(norm(x), cos, sin) + expected = attention(q, k, v, cu_seqlens=cu_seqlens).reshape_as(x) + projected_rows.clear() + attention_calls.clear() + actual = multimodal.Glm5NextVisionAttention.context_with_chunked_projection( + layer, x, norm, cu_seqlens, cos, sin, 8 + ) + torch.testing.assert_close(actual, expected) + assert projected_rows == [4, 4, 3] + assert attention_calls == [11] + + +@pytest.mark.parametrize( + "token_budget,expand,expected_tokens", + [ + (8000, 1, 8000), + (8192, 1, 8192), + (8000, 2, 8000), + (8191, 1, 8190), + (8191, 2, 8188), + ], +) +def test_glm5next_image_budget_covers_rectangular_canvases( + token_budget, expand, expected_tokens +): + from vllm.models.glm5next.common.multimodal import Glm5NextProcessingInfo + + processor = SimpleNamespace( + patch_size=14, merge_size=2, temporal_patch_size=2, patch_expand_factor=expand + ) + info = Glm5NextProcessingInfo.__new__(Glm5NextProcessingInfo) + info.get_hf_processor = lambda **kwargs: SimpleNamespace(image_processor=processor) + info.get_hf_config = lambda: SimpleNamespace( + vision_config=SimpleNamespace( + patch_size=14, spatial_merge_size=2, temporal_patch_size=2 + ) + ) + info._get_image_max_pixels = lambda: token_budget * 2 * 28**2 + canvas = info.get_image_size_with_most_features() + assert canvas.width % (28 * expand) == 0 + assert canvas.height % (28 * expand) == 0 + assert canvas.width // processor.patch_size <= 8192 + assert canvas.height // processor.patch_size <= 8192 + assert canvas.width * canvas.height // 28**2 == expected_tokens + assert info.get_max_image_tokens() == expected_tokens + assert ( + info.get_num_image_tokens(image_width=3082, image_height=2048) <= token_budget + ) + if token_budget == 8000 and expand == 1: + assert (canvas.width, canvas.height) == (2800, 2240) + assert canvas.width * canvas.height // 28**2 > 7957 + + +def test_glm5next_image_pixel_budget_respects_token_override(): + from vllm.models.glm5next.common.multimodal import Glm5NextProcessingInfo + + processor = SimpleNamespace(patch_size=14, merge_size=2, temporal_patch_size=2) + info = Glm5NextProcessingInfo.__new__(Glm5NextProcessingInfo) + info.get_hf_processor = lambda **kwargs: SimpleNamespace(image_processor=processor) + info.ctx = SimpleNamespace( + get_merged_mm_kwargs=lambda kwargs: {"max_image_tokens": 2048} + ) + assert info._get_image_max_pixels() == 2048 * 2 * 28**2 + + def test_glm5next_processor_counts_video_only_tokens() -> None: class FakeVideoProcessor: merge_size = 2 diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index e5df32783a15..341c9359a446 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -50,12 +50,16 @@ def test_workspace_lanes_do_not_alias_and_restore_context(monkeypatch) -> None: ) assert manager._current_workspaces == [None, None, None, None] + assert manager.available_bytes() == 0 (target,) = manager.get_simultaneous(((512,), torch.uint8)) + assert manager.available_bytes() == 512 with workspace.use_workspace_lane(1): (draft,) = manager.get_simultaneous(((256,), torch.uint8)) (draft_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + assert manager.available_bytes() == 256 (target_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + assert manager.available_bytes() == 512 assert manager._current_workspaces[0].numel() == 512 # type: ignore[union-attr] assert manager._current_workspaces[1].numel() == 256 # type: ignore[union-attr] diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py index aa2db6da23a4..a0fe790a4993 100644 --- a/vllm/models/glm5next/common/multimodal.py +++ b/vllm/models/glm5next/common/multimodal.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from functools import cached_property, partial +from math import isqrt, prod import numpy as np import torch @@ -24,6 +25,7 @@ MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope @@ -44,6 +46,13 @@ from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.multimodal.parse import ImageSize, MultiModalDataItems from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +_VISION_PROJECTION_CHUNK_ROWS = 4096 +_VISION_ROPE_MAX_POSITION = 8192 class Glm5NextVisionPatchEmbed(nn.Module): @@ -169,6 +178,10 @@ def __init__( prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + self._borrow_qkv_workspace = ( + isinstance(self.qkv.quant_method, UnquantizedLinearMethod) + and self.attn.attn_backend == AttentionBackendEnum.FLASH_ATTN + ) def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]: seq_len, bs, _ = qkv.shape @@ -190,6 +203,25 @@ def forward( rotary_pos_emb_sin: torch.Tensor, max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: + q, k, v = self._project_qkv(x, rotary_pos_emb_cos, rotary_pos_emb_sin) + context_layer = self.attn( + query=q, + key=k, + value=v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + del q, k, v + context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() + output, _ = self.proj(context_layer) + return output + + def _project_qkv( + self, + x: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor | None, + rotary_pos_emb_sin: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: x, _ = self.qkv(x) q, k, v = self.split_qkv(x) @@ -204,10 +236,14 @@ def forward( self.k_norm.weight, self.q_norm.variance_epsilon, ) + del q_flat, k_flat q = q.view(q_shape) k = k.view(k_shape) q, k, v = (rearrange(t, "s b ... -> b s ...").contiguous() for t in (q, k, v)) + # Contiguous attention inputs own their storage; the packed projection + # and norm input copies must not stay live through attention/output GEMM. + del x if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: qk_concat = torch.cat([q, k], dim=0) qk_rotated = self.apply_rotary_emb( @@ -216,18 +252,64 @@ def forward( rotary_pos_emb_sin, ) q, k = torch.chunk(qk_rotated, 2, dim=0) + del qk_concat, qk_rotated - context_layer = self.attn( + return q, k, v + + def context_with_chunked_projection( + self, + x: torch.Tensor, + input_norm: nn.Module, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None, + ) -> torch.Tensor: + """Bound projection scratch without partitioning image attention. + + Normalization, QKV projection and rotary embedding are token-local. + Attention still consumes the complete Q/K/V and original image/frame + sequence boundaries. No cache or tensor dtype changes are involved. + """ + rows, batch_size, _ = x.shape + shape = ( + batch_size, + rows, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ) + manager = ( + current_workspace_manager() + if self._borrow_qkv_workspace and is_workspace_manager_initialized() + else None + ) + qkv_bytes = 3 * ((prod(shape) * x.element_size() + 255) // 256) * 256 + if manager is not None and manager.available_bytes() >= qkv_bytes: + # The BF16 projections, rotary operation and FlashAttention do not + # borrow worker scratch. These Q/K/V views expire before the output + # projection; the text model executes after the encoder returns. + q, k, v = manager.get_simultaneous(*((shape, x.dtype),) * 3) + else: + q, k, v = (x.new_empty(shape) for _ in range(3)) + for start in range(0, rows, _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, rows) + cq, ck, cv = self._project_qkv( + input_norm(x[start:stop]), + rotary_pos_emb_cos[start:stop], + rotary_pos_emb_sin[start:stop], + ) + q[:, start:stop].copy_(cq) + k[:, start:stop].copy_(ck) + v[:, start:stop].copy_(cv) + del cq, ck, cv + context = self.attn( query=q, key=k, value=v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, ) - context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() - - output, _ = self.proj(context_layer) - return output + return rearrange(context, "b s h d -> s b (h d)").contiguous() class Glm5NextVisionBlock(nn.Module): @@ -270,6 +352,24 @@ def forward( rotary_pos_emb_sin: torch.Tensor, max_seqlen: int | None = None, ) -> torch.Tensor: + if x.shape[0] > _VISION_PROJECTION_CHUNK_ROWS: + context = self.attn.context_with_chunked_projection( + x, + self.norm1, + cu_seqlens, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + max_seqlen, + ) + for start in range(0, x.shape[0], _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, x.shape[0]) + x_attn, _ = self.attn.proj(context[start:stop]) + normalized, residual = self.norm2(x[start:stop], residual=x_attn) + # Each token's residual is consumed before its owned input + # storage is reused. Attention has already read the full image. + x[start:stop].copy_(residual + self.mlp(normalized)) + del normalized, residual, x_attn + return x x_attn = self.attn( self.norm1(x), cu_seqlens=cu_seqlens, @@ -330,6 +430,7 @@ def forward(self, x: torch.Tensor): x = self.extra_activation_func(self.post_projection_norm(x)) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) + del gate_up x, _ = self.down_proj(x) return x @@ -390,7 +491,7 @@ def __init__( head_dim = self.hidden_size // self.num_heads self.rotary_pos_emb = get_rope( head_size=head_dim, - max_position=8192, + max_position=_VISION_ROPE_MAX_POSITION, is_neox_style=True, rope_parameters={"partial_rotary_factor": 0.5}, ) @@ -599,6 +700,38 @@ def forward( max_seqlen=max_seqlen, ) + # The GLM merger preserves the number of elements in each spatial + # group. Reuse the owned transformer output only after that group's + # normalization, convolution and merger have consumed its values. + if ( + x.shape[0] > _VISION_PROJECTION_CHUNK_ROWS + and self.out_hidden_size == self.spatial_merge_size**2 * x.shape[-1] + ): + hidden_size = x.shape[-1] + for start in range(0, x.shape[0], _VISION_PROJECTION_CHUNK_ROWS): + stop = min(start + _VISION_PROJECTION_CHUNK_ROWS, x.shape[0]) + x[start:stop].copy_(self.post_layernorm(x[start:stop])) + x = x.view(-1, self.out_hidden_size) + groups_per_chunk = max( + 1, _VISION_PROJECTION_CHUNK_ROWS // self.spatial_merge_size**2 + ) + for start in range(0, x.shape[0], groups_per_chunk): + stop = min(start + groups_per_chunk, x.shape[0]) + patches = ( + x[start:stop] + .view( + -1, + self.spatial_merge_size, + self.spatial_merge_size, + hidden_size, + ) + .permute(0, 3, 1, 2) + ) + merged = self.downsample(patches).view(-1, self.out_hidden_size) + x[start:stop].copy_(self.merger(merged)) + del patches, merged + return x + # adapter x = self.post_layernorm(x) x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1]) @@ -652,7 +785,44 @@ def _get_image_max_pixels(self) -> int: mm_kwargs = self.ctx.get_merged_mm_kwargs({}) if (override := mm_kwargs.get("max_pixels")) is not None: return int(override) - return self._processor_pixel_budget(self.get_hf_processor().image_processor)[1] + processor = self.get_hf_processor().image_processor + if (tokens := mm_kwargs.get("max_image_tokens")) is not None: + return ( + int(tokens) + * processor.temporal_patch_size + * (processor.patch_size * processor.merge_size) ** 2 + ) + return self._processor_pixel_budget(processor)[1] + + def get_image_size_with_most_features(self) -> ImageSize: + """Maximize aligned canvas area within the token and vision RoPE limits. + + A square canvas can underfill the budget: 89 by 89 merged patches + produce 7921 features, while a valid 80 by 100 canvas produces 8000. + Encoder admission and profiling must cover the rectangular maximum. + """ + processor = self.get_hf_processor().image_processor + factor = ( + processor.patch_size * processor.merge_size * processor.patch_expand_factor + ) + cells = self._get_image_max_pixels() // ( + processor.temporal_patch_size * factor**2 + ) + if cells < 1: + raise ValueError("Image budget must fit at least one aligned canvas cell") + max_side = _VISION_ROPE_MAX_POSITION // ( + processor.merge_size * processor.patch_expand_factor + ) + if max_side < 1: + raise ValueError("One aligned canvas cell exceeds the vision RoPE grid") + width = height = 1 + # Include every feasible shorter side. Maximum area prevents memory + # underprofiling; equal-area ties prefer the least elongated canvas. + for candidate_height in range(1, min(isqrt(cells), max_side) + 1): + candidate_width = min(cells // candidate_height, max_side) + if candidate_width * candidate_height >= width * height: + width, height = candidate_width, candidate_height + return ImageSize(width=width * factor, height=height * factor) def _get_video_max_pixels(self) -> int: mm_kwargs = self.ctx.get_merged_mm_kwargs({}) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index ddde9f8b1df5..b44225340491 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -178,6 +178,23 @@ def is_locked(self) -> bool: """Check if workspace is locked.""" return self._locked + def available_bytes(self) -> int: + """Capacity of the active execution slot, without allocating storage. + + A borrower must ensure that no nested or concurrent operation consumes + this slot while its views are live. + """ + ubatch_id = dbo_current_ubatch_id() + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + return self._workspace_size_bytes( + self._current_workspaces[ubatch_id * self._num_lanes + lane] + ) + def get_simultaneous( self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] ) -> list[torch.Tensor]: From a5bab53f7bfe005823bbea57e074708c84bbafdc Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 02:45:08 +0000 Subject: [PATCH 111/365] docs(quantization): define NVFP4 Marlin rescaling errors Document FP16 bypass, empty and non-positive inputs, the returned factor and NaN rejection. Executable AST is unchanged; Ruff passes. AI assistance: OpenAI Codex. Signed-off-by: Martin Vit --- .../quantization/utils/marlin_utils_fp4.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index a94823434cbf..217aeb32c3a1 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -39,9 +39,20 @@ def _nvfp4_compute_scale_factor( marlin_scales: torch.Tensor, a_dtype: torch.dtype | None = None, ) -> float: - """Compute the power-of-2 scale_factor needed so that all non-zero - values in marlin_scales * 2^7 are >= 2 after rescaling. - Returns a Python float (power of 2, >= 1.0).""" + """Compute the power-of-two rescaling factor for NVFP4 Marlin scales. + + Args: + marlin_scales: Weight scales in the Marlin layout. + a_dtype: Activation dtype. FP16 disables rescaling. + + Returns: + A Python power-of-two factor greater than or equal to one, derived + from the largest positive scale. Empty or non-positive inputs and + FP16 activations return one. + + Raises: + ValueError: Scales contain NaN and rescaling is enabled. + """ # Since half has a smaller dynamic range compared to bfloat16, # no rescaling is applied here if active dtype is half. From eb4ec03df6a1cb9fed33ba2f6e2981cade413edd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 10:12:50 +0000 Subject: [PATCH 112/365] fix(cache): retain a bounded checkpoint before long prompt tails Preserve one budgeted recurrent checkpoint one to two prefill chunks before a long prompt ends. Tail edits can restore a strictly matching earlier prefix without retaining every intermediate state or relaxing identity checks. Account for the optional bundle in KV admission and automatic context sizing, including short-context capacity. Validation: 371 prefix/admission tests, 16 capacity cases, and 5 GPU capture/restore cases pass. TP2/DCP2 MTP3 restores 516096 shared tokens after a 520k prompt-tail edit in 1.77s versus 70.07s without the checkpoint. Factual GPU/RAM and restart-disk restores pass; an earlier edit correctly misses. The additional bundle reduces the tested fixed-pool maximum context by 27648 tokens. Co-authored-by: Codex Signed-off-by: Martin Vit --- .../test_partial_prefix_cache_hits.py | 106 ++++++++++++++++++ .../core/test_boundary_admission_horizon.py | 1 + tests/v1/core/test_kv_cache_utils.py | 23 ++-- tests/v1/worker/test_gpu_input_batch_v2.py | 49 ++++---- vllm/config/cache.py | 3 +- vllm/v1/core/boundary_checkpoint.py | 28 ++++- vllm/v1/core/kv_cache_manager.py | 35 ++++-- vllm/v1/core/kv_cache_utils.py | 22 +++- vllm/v1/core/sched/output.py | 2 + vllm/v1/core/sched/scheduler.py | 47 ++++++-- vllm/v1/outputs.py | 4 +- vllm/v1/request.py | 3 + vllm/v1/worker/gpu/boundary_checkpoint.py | 34 +++++- vllm/v1/worker/gpu/input_batch.py | 6 +- 14 files changed, 300 insertions(+), 63 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 6aaa2f1cfe52..7a40f9372ed7 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -16,6 +16,7 @@ from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request from vllm.distributed.kv_transfer.kv_connector.v1.base import SupportsHMA from vllm.utils.hashing import sha256 +from vllm.v1.core.boundary_checkpoint import get_prefill_tail_checkpoint_position from vllm.v1.core.kv_cache_utils import ( KVCacheBlockCopy, get_block_hash, @@ -2700,6 +2701,111 @@ def test_request_boundaries_split_prefill_at_instruction_endpoint(): assert Scheduler._mamba_block_aligned_split(scheduler, request, 4096) == 4096 +@pytest.mark.parametrize( + "prompt,budget,instruction,expected", + [ + (8191, 4096, None, None), + (8192, 4096, None, 4096), + (10000, 4096, None, 4096), + (520013, 3072, None, 516096), + (12000, 4096, 5000, None), + (16000, 4096, 5000, 8192), + ], +) +def test_prefill_tail_leaves_a_bounded_suffix_without_duplicating_instructions( + prompt, budget, instruction, expected +): + assert get_prefill_tail_checkpoint_position(prompt, budget, instruction) == expected + + +def test_request_boundaries_stop_at_instruction_and_prefill_tail_once(): + request = make_request("chat", [0] * 9000, 32, sha256) + request.use_boundary_checkpoints = True + request.recurrent_instruction_boundary = 3000 + request.recurrent_prefill_tail_boundary = 4096 + scheduler = SimpleNamespace() + for start, expected in ((0, 3000), (3000, 1096), (4096, 4096)): + request.num_computed_tokens = start + assert ( + Scheduler._mamba_block_aligned_split(scheduler, request, 4096) == expected + ) + request.num_computed_tokens = 0 + assert ( + Scheduler._mamba_block_aligned_split( + scheduler, request, 4096, num_external_computed_tokens=4096 + ) + == 4096 + ) + + +@pytest.mark.parametrize("instruction", [None, 5]) +@pytest.mark.parametrize("tail", [None, 12]) +def test_optional_boundary_slots_allocate_only_active_bundles(instruction, tail): + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=4, + num_blocks=64, + enable_boundary_checkpoints=True, + ) + free = manager.block_pool.get_num_free_blocks() + request = make_request("producer", list(range(17)), 4, sha256) + request.recurrent_instruction_boundary = instruction + request.recurrent_prefill_tail_boundary = tail + manager.get_computed_blocks(request) + assert manager.allocate_slots(request, 17) is not None + slots = request.boundary_checkpoint_blocks + assert slots is not None + width = manager.num_kv_cache_groups + 1 + assert ( + sum(block != 0 for slot in slots for block in slot) + == (2 + int(instruction is not None) + int(tail is not None)) * width + ) + if tail is not None and instruction is None: + assert slots[2] == (0,) * width + assert all(slots[3]) + manager.free(request) + assert manager.reset_prefix_cache() + assert manager.block_pool.get_num_free_blocks() == free + + +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_prefill_tail_recovers_suffix_edits_but_never_earlier_divergence(dcp): + manager = make_full_mamba_manager( + dcp_world_size=dcp, + hash_block_size=4, + num_blocks=128, + enable_boundary_checkpoints=True, + use_eagle=True, + ) + tokens = list(range(29)) + producer = make_request("producer", tokens, 4, sha256) + producer.recurrent_prefill_tail_boundary = 17 + manager.get_computed_blocks(producer) + assert manager.allocate_slots(producer, 17) is not None + tail = manager.publish_boundary_checkpoint(producer, 17, kind="prefill_tail") + assert tail is not None and tail.kind == "prefill_tail" + producer.num_computed_tokens = 17 + assert manager.allocate_slots(producer, 12) is not None + assert manager.publish_boundary_checkpoint(producer, 29, kind="prompt") is not None + manager.free(producer) + + for name, suffix, expected in ( + ("extended", tokens + [80, 81], 29), + ("tail-edit", tokens[:27] + [80, 81, 82], 17), + ("earlier-edit", tokens[:16] + [80] + tokens[17:], 0), + ): + consumer = make_request(name, suffix, 4, sha256) + blocks, hit, _ = manager.get_computed_blocks(consumer) + assert hit == expected + assert ( + manager.allocate_slots(consumer, len(suffix) - hit, hit, blocks) is not None + ) + manager.free(consumer) + _, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + assert manager.reset_prefix_cache() + + def test_internal_checkpoint_requires_block_aligned_start(): hash_block_size = 2 mamba_block_size = 16 diff --git a/tests/v1/core/test_boundary_admission_horizon.py b/tests/v1/core/test_boundary_admission_horizon.py index b8a3b099fffa..3dbc49e63a53 100644 --- a/tests/v1/core/test_boundary_admission_horizon.py +++ b/tests/v1/core/test_boundary_admission_horizon.py @@ -117,6 +117,7 @@ def test_small_pool_capacity_covers_instruction_checkpoint_restore(num_blocks): groups.append(KVCacheGroupSpec(["attention"], attention)) config = SimpleNamespace( use_request_boundary_checkpoints=True, + scheduler_config=SimpleNamespace(max_num_scheduled_tokens=4096), attention_config=SimpleNamespace(hisparse_config=None), model_config=SimpleNamespace(max_model_len=1048576), parallel_config=SimpleNamespace(decode_context_parallel_size=2), diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 7d6957120049..36316cbeb558 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2654,11 +2654,13 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): @pytest.mark.parametrize("with_tail", [False, True]) +@pytest.mark.parametrize("chunk_budget,private_bundles", [(8192, 4), (4096, 5)]) def test_balanced_glm_boundary_capacity_includes_private_endpoints( - monkeypatch, with_tail + monkeypatch, with_tail, chunk_budget, private_bundles ): """The GLM slot-sharing layout still reserves distinct endpoint block IDs.""" config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + config.scheduler_config.max_num_scheduled_tokens = chunk_budget specs = ( _glm5_like_kv_cache_spec_with_tail() if with_tail @@ -2669,7 +2671,7 @@ def test_balanced_glm_boundary_capacity_includes_private_endpoints( block_bytes = kv_cache_utils._pool_bytes_per_block(groups) live_bytes = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) live_blocks = live_bytes // block_bytes - private_blocks = 4 * (len(groups) + 1) + private_blocks = private_bundles * (len(groups) + 1) monkeypatch.setattr( VllmConfig, "use_request_boundary_checkpoints", property(lambda self: True) ) @@ -4277,11 +4279,17 @@ def test_auto_fit_max_model_len_with_hybrid(): @pytest.mark.parametrize("dcp", [1, 2, 4]) -@pytest.mark.parametrize("num_blocks,expected_pages", [(38, 0), (62, 0), (70, 7)]) -def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_pages): +@pytest.mark.parametrize( + "num_blocks,expected_pages,chunk_budget", + [(38, 0, 1048576), (62, 0, 1048576), (70, 7, 1048576), (78, 7, 256)], +) +def test_boundary_capacity_includes_private_endpoints( + dcp, num_blocks, expected_pages, chunk_budget +): """Pool sizing includes private endpoints and the immutable restore source.""" config = SimpleNamespace( use_request_boundary_checkpoints=True, + scheduler_config=SimpleNamespace(max_num_scheduled_tokens=chunk_budget), attention_config=SimpleNamespace(hisparse_config=None), model_config=SimpleNamespace(max_model_len=1048576), parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), @@ -4304,8 +4312,9 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ groups = [KVCacheGroupSpec([f"state-{i}"], state) for i in range(6)] groups.append(KVCacheGroupSpec(["attention"], attention)) block_bytes = kv_cache_utils._pool_bytes_per_block(groups) - # Six groups require five live/scratch states each; three private endpoints - # and a restore source own eight blocks each. The null block is unavailable. + # Six groups require five live/scratch states each. Three semantic endpoints + # and a restore source own eight blocks each; enabling the prefill-tail + # endpoint costs another eight. The null block is unavailable. available = (num_blocks - 1) * block_bytes estimate = kv_cache_utils._estimate_max_model_len_from_groups( config, groups, available @@ -4329,7 +4338,7 @@ def test_boundary_capacity_includes_private_endpoints(dcp, num_blocks, expected_ num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups ) assert get_max_concurrency_for_kv_cache_config(config, cache) == ( - num_blocks / (30 + 32 + expected_pages) + num_blocks / (30 + (40 if chunk_budget == 256 else 32) + expected_pages) ) # An aligned-policy control has no endpoint reservation. Its capacity diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py index 70e8c7d2e161..4b107955affb 100644 --- a/tests/v1/worker/test_gpu_input_batch_v2.py +++ b/tests/v1/worker/test_gpu_input_batch_v2.py @@ -178,8 +178,8 @@ def test_boundary_auxiliary_restore_survives_slot_reuse_and_virtual_attention_pa state.blocks.copy_( torch.tensor( [ - [[4, 8], [5, 9], [12, 13]], - [[6, auxiliary_block], [7, 11], [14, 15]], + [[4, 8], [5, 9], [12, 13], [0, 0]], + [[6, auxiliary_block], [7, 11], [14, 15], [0, 0]], ], device=device, ) @@ -187,9 +187,9 @@ def test_boundary_auxiliary_restore_survives_slot_reuse_and_virtual_attention_pa idx = torch.tensor([1, 0], dtype=torch.int32, device=device) capture = torch.tensor( [ - [[7, 0, 0], [0, 11, 0]], - [[0, 0, 0], [0, 0, 0]], - [[0, -1, 0], [-1, 1, 0]], + [[7, 0, 0, 0], [0, 11, 0, 0]], + [[0, 0, 0, 0], [0, 0, 0, 0]], + [[0, -1, 0, 0], [-1, 1, 0, 0]], ], dtype=torch.int32, device=device, @@ -229,11 +229,13 @@ def test_boundary_auxiliary_restore_survives_slot_reuse_and_virtual_attention_pa lora_request=None, boundary_checkpoint=BoundaryCheckpoint(1, 7, ((6,),), (auxiliary_block,)), boundary_checkpoint_blocks=((12, 13), (14, 15)), + recurrent_prefill_tail_boundary=4, ) state.add_request(restore_slot, request) torch.testing.assert_close(raw_state[restore_slot], expected_state) assert anchors[restore_slot].item() == 6 assert accepted[restore_slot].item() == 1 + assert state.seen[restore_slot, 3].item() == 1 torch.testing.assert_close(state.get_hidden_states(auxiliary_block), hidden[:1]) torch.testing.assert_close( state.get_hidden_states(auxiliary_block, draft=True), spec_hidden[:1] @@ -251,15 +253,15 @@ def t(data): state = SimpleNamespace( metadata=t( [ - [1, 7, 0, 7, 8, 1], - [1, 7, 0, 7, 30, 1], - [0, 7, 0, 7, 30, 1], - [1, 7, 0, 12, 30, 1], - [1, 9, 4, 9, 30, 1], + [1, 7, 0, 7, 8, 1, 0], + [1, 7, 0, 7, 30, 1, 0], + [0, 7, 0, 7, 30, 1, 0], + [1, 7, 0, 12, 30, 1, 0], + [1, 9, 4, 9, 30, 1, 4], ] ), stop_tokens=torch.full((5, 128), 99, dtype=torch.int32, device="cuda"), - seen=t([[0, 0, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0], [0, 0, 0]]), + seen=t([[0, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0]]), ) sampled = t( [ @@ -276,7 +278,7 @@ def t(data): rejected = t([0, 0, 0, 0, 0]) last_sampled = t([0, 0, 0, 0, 0]) all_tokens = torch.zeros((5, 32), dtype=torch.int32, device="cuda") - capture = torch.empty((3, 5, 3), dtype=torch.int32, device="cuda") + capture = torch.empty((3, 5, 4), dtype=torch.int32, device="cuda") post_update( t([0, 1, 2, 3, 4]), computed, @@ -292,27 +294,28 @@ def t(data): capture, ) assert capture[0].tolist() == [ - [7, 7, 0], - [0, 10, 0], - [0, 0, 0], - [0, 11, 0], - [0, 0, 4], + [7, 7, 0, 0], + [0, 10, 0, 0], + [0, 0, 0, 0], + [0, 11, 0, 0], + [0, 0, 4, 4], ] assert capture[1].tolist() == [ - [0, 0, 0], - [0, 1, 0], - [0, 3, 0], - [0, 2, 0], - [0, 0, 0], + [0, 0, 0, 0], + [0, 1, 0, 0], + [0, 3, 0, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], ] assert num_sampled.tolist() == [1, 2, 4, 3, 0] assert computed.tolist() == [7, 10, 12, 11, 4] assert total.tolist() == [8, 11, 13, 12, 9] assert last_sampled.tolist() == [10, 99, 13, 99, 0] - assert capture[2][0].tolist() == [6, 6, -1] + assert capture[2][0].tolist() == [6, 6, -1, -1] assert capture[2][1, 1].item() == 8 assert capture[2][3, 1].item() == 17 assert capture[2][4, 2].item() == 22 + assert capture[2][4, 3].item() == 22 @pytest.mark.parametrize( diff --git a/vllm/config/cache.py b/vllm/config/cache.py index a874be06494e..acbc9e44ed7d 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -165,7 +165,8 @@ class CacheConfig: Applies only to sliding-window and Mamba cache groups.""" recurrent_checkpoint_policy: RecurrentCheckpointPolicy = "auto" """Retention policy for reusable recurrent state. ``request_boundaries`` - retains a verified leading chat-instruction prefix, the completed prompt, + retains a verified leading chat-instruction prefix, one prefill checkpoint + before the final chunk of a long prompt, the completed prompt, and the committed response endpoint, disabling arbitrary intermediate checkpoints. ``aligned`` preserves block-aligned retention. ``auto`` selects request boundaries for supported configurations and aligned retention diff --git a/vllm/v1/core/boundary_checkpoint.py b/vllm/v1/core/boundary_checkpoint.py index 1be917a21808..cf5595009ce4 100644 --- a/vllm/v1/core/boundary_checkpoint.py +++ b/vllm/v1/core/boundary_checkpoint.py @@ -14,14 +14,34 @@ MAX_BOUNDARY_STOP_TOKENS = 128 -# Prompt and response retain their original slots so cached request state and -# worker metadata remain compatible when an instruction checkpoint is absent. +# Optional checkpoints have fixed slots; absent slots contain no KV blocks. PROMPT_CHECKPOINT_SLOT = 0 RESPONSE_CHECKPOINT_SLOT = 1 INSTRUCTION_CHECKPOINT_SLOT = 2 -NUM_BOUNDARY_CHECKPOINT_SLOTS = 3 +PREFILL_TAIL_CHECKPOINT_SLOT = 3 +NUM_BOUNDARY_CHECKPOINT_SLOTS = 4 -BoundaryCheckpointKind = Literal["instruction", "prompt", "response"] +BoundaryCheckpointKind = Literal["instruction", "prompt", "response", "prefill_tail"] + + +def get_prefill_tail_checkpoint_position( + prompt_tokens: int, chunk_budget: int, instruction_boundary: int | None +) -> int | None: + """Leave one to two prefill chunks after a bounded suffix-recovery point.""" + if chunk_budget <= 0: + return None + position = (prompt_tokens // chunk_budget - 1) * chunk_budget + return position if position > (instruction_boundary or 0) else None + + +def boundary_checkpoint_slots(request: Request) -> tuple[int, ...]: + """Allocate only the endpoints that this request can publish.""" + slots: tuple[int, ...] = (PROMPT_CHECKPOINT_SLOT, RESPONSE_CHECKPOINT_SLOT) + if request.recurrent_instruction_boundary is not None: + slots += (INSTRUCTION_CHECKPOINT_SLOT,) + if request.recurrent_prefill_tail_boundary is not None: + slots += (PREFILL_TAIL_CHECKPOINT_SLOT,) + return slots @dataclass(frozen=True) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 7e36eca00747..84812b69baed 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -11,11 +11,13 @@ from vllm.utils.math_utils import cdiv from vllm.v1.core.boundary_checkpoint import ( INSTRUCTION_CHECKPOINT_SLOT, + PREFILL_TAIL_CHECKPOINT_SLOT, PROMPT_CHECKPOINT_SLOT, RESPONSE_CHECKPOINT_SLOT, BoundaryCheckpoint, BoundaryCheckpointCache, BoundaryCheckpointKind, + boundary_checkpoint_slots, ) from vllm.v1.core.kv_cache_coordinator import ( HybridKVCacheCoordinator, @@ -595,9 +597,7 @@ def allocate_slots( boundary_replay_managers = [] if request.use_boundary_checkpoints: if request.request_id not in self._boundary_allocations: - checkpoint_slots = 2 + int( - request.recurrent_instruction_boundary is not None - ) + checkpoint_slots = len(boundary_checkpoint_slots(request)) boundary_blocks = checkpoint_slots * (self.num_kv_cache_groups + 1) if request.boundary_checkpoint is not None: boundary_blocks += sum( @@ -798,14 +798,18 @@ def allocate_slots( request.request_id not in self._boundary_allocations ): width = self.num_kv_cache_groups + 1 - checkpoint_slots = 2 + int( - request.recurrent_instruction_boundary is not None - ) - allocation = self.block_pool.get_new_blocks(checkpoint_slots * width) + slots = boundary_checkpoint_slots(request) + allocation = self.block_pool.get_new_blocks(len(slots) * width) self._boundary_allocations[request.request_id] = allocation + by_slot = { + slot: tuple( + block.block_id + for block in allocation[index * width : (index + 1) * width] + ) + for index, slot in enumerate(slots) + } request.boundary_checkpoint_blocks = tuple( - tuple(block.block_id for block in allocation[start : start + width]) - for start in range(0, checkpoint_slots * width, width) + by_slot.get(slot, (0,) * width) for slot in range(max(slots) + 1) ) # P/D: delay caching blocks if we have to recv from @@ -871,10 +875,16 @@ def publish_boundary_checkpoint( if num_tokens != request.num_prompt_tokens: return None slot = PROMPT_CHECKPOINT_SLOT - else: + elif kind == "instruction": if num_tokens != request.recurrent_instruction_boundary: return None slot = INSTRUCTION_CHECKPOINT_SLOT + elif kind == "prefill_tail": + if num_tokens != request.recurrent_prefill_tail_boundary: + return None + slot = PREFILL_TAIL_CHECKPOINT_SLOT + else: + return None if slot >= len(allocation): return None grouped_blocks = [] @@ -976,7 +986,10 @@ def reserve_external_boundary_checkpoint( raise ValueError("Imported boundary must cover an existing request prefix") if not 0 <= draft_prefix_len <= num_tokens: raise ValueError("Imported draft prefix exceeds the target prefix") - if kind not in ("instruction", "prompt", "response") or num_ranks < 1: + if ( + kind not in ("instruction", "prompt", "response", "prefill_tail") + or num_ranks < 1 + ): raise ValueError("Invalid external checkpoint kind or rank count") if page_positions != self.boundary_checkpoint_page_positions(num_tokens): raise ValueError("External checkpoint does not cover every live cache page") diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 85f5c062ec82..078033fc55cb 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1081,14 +1081,32 @@ def _request_boundary_reserve_blocks(vllm_config: VllmConfig, num_groups: int) - if not num_groups or not vllm_config.use_request_boundary_checkpoints: return 0 - from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS + from vllm.v1.core.boundary_checkpoint import ( + NUM_BOUNDARY_CHECKPOINT_SLOTS, + get_prefill_tail_checkpoint_position, + ) # KVCacheManager.allocate_slots reserves one block per group and one # auxiliary block per endpoint, separately from the live sequence state. # Restoring a checkpoint pins its source until the reader finishes, while # partial-tail CoW creates writable state. One additional bundle covers # these source pages and auxiliary state alongside the private endpoints. - return (NUM_BOUNDARY_CHECKPOINT_SLOTS + 1) * (num_groups + 1) + scheduler = vllm_config.scheduler_config + chunk_budget = ( + scheduler.max_num_scheduled_tokens + if scheduler.max_num_scheduled_tokens is not None + else scheduler.max_num_batched_tokens + ) + has_prefill_tail = ( + get_prefill_tail_checkpoint_position( + vllm_config.model_config.max_model_len, chunk_budget, None + ) + is not None + ) + # Short-context capacity must not reserve a tail endpoint that no request + # can reach. The capacity search calls this for each candidate context. + slots = NUM_BOUNDARY_CHECKPOINT_SLOTS - int(not has_prefill_tail) + return (slots + 1) * (num_groups + 1) def get_max_concurrency_for_kv_cache_config( diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 46ea472171c3..bdb08e48d295 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -54,6 +54,7 @@ class NewRequestData: boundary_checkpoint: BoundaryCheckpoint | None = None boundary_checkpoint_blocks: tuple[tuple[int, ...], ...] | None = None recurrent_instruction_boundary: int | None = None + recurrent_prefill_tail_boundary: int | None = None @classmethod def from_request( @@ -82,6 +83,7 @@ def from_request( boundary_checkpoint=request.boundary_checkpoint, boundary_checkpoint_blocks=request.boundary_checkpoint_blocks, recurrent_instruction_boundary=request.recurrent_instruction_boundary, + recurrent_prefill_tail_boundary=request.recurrent_prefill_tail_boundary, ) @property diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2ace136d5055..785212c3ae6b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -32,6 +32,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.utils import get_mm_features_in_window +from vllm.v1.core.boundary_checkpoint import get_prefill_tail_checkpoint_position from vllm.v1.core.encoder_cache_manager import ( EncoderCacheManager, EncoderDecoderCacheManager, @@ -513,14 +514,21 @@ def _mamba_block_aligned_split( + num_external_computed_tokens ) if request.use_boundary_checkpoints: - # Running-state migration still happens in the worker, but these - # requests publish only semantic endpoints. Stop exactly at the - # leading instruction boundary so the worker can snapshot the - # recurrent state associated with that token prefix. - instruction_boundary = request.recurrent_instruction_boundary - if instruction_boundary is not None and start < instruction_boundary: - return min(num_new_tokens, instruction_boundary - start) - return num_new_tokens + # The worker can snapshot recurrent state only at a computed + # endpoint, never retroactively inside an already executed chunk. + end = start + num_new_tokens + boundary_stops = ( + request.recurrent_instruction_boundary, + request.recurrent_prefill_tail_boundary, + ) + return min( + ( + stop - start + for stop in boundary_stops + if stop is not None and start < stop < end + ), + default=num_new_tokens, + ) # Split only during prefill: `request.num_tokens - 1` extends this to # resumed requests replaying their output tokens. prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1) @@ -2930,9 +2938,18 @@ def update_from_output( finish_reason = None captures = model_runner_output.boundary_checkpoint_tokens if captures is not None and not output_is_stale: - prompt_boundary, response_boundary, instruction_boundary = captures[ - req_index - ] + ( + prompt_boundary, + response_boundary, + instruction_boundary, + prefill_tail_boundary, + ) = captures[req_index] + if prefill_tail_boundary: + checkpoint = self.kv_cache_manager.publish_boundary_checkpoint( + request, prefill_tail_boundary, kind="prefill_tail" + ) + if checkpoint is not None and self.connector is not None: + self.connector.store_boundary_checkpoint(request, checkpoint) if instruction_boundary: checkpoint = self.kv_cache_manager.publish_boundary_checkpoint( request, instruction_boundary, kind="instruction" @@ -3342,6 +3359,14 @@ def add_request(self, request: Request) -> None: # Streaming-input session finished. self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) else: + if self.kv_cache_manager.boundary_checkpoints is not None: + request.recurrent_prefill_tail_boundary = ( + get_prefill_tail_checkpoint_position( + request.num_prompt_tokens, + self.max_num_scheduled_tokens, + request.recurrent_instruction_boundary, + ) + ) if request.resumable: request.streaming_queue = deque() self._enqueue_waiting_request(request) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index fab96031f3ec..783ce3b7d6b2 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -376,8 +376,8 @@ class ModelRunnerOutput: # ``None`` when ``return_sampling_mask`` is off. sampling_masks: SamplingMaskLists | None = None - # Per request: completed prompt, response, and leading-instruction - # checkpoint token counts, or 0. + # Per request: completed prompt, response, leading-instruction, and + # prefill-tail checkpoint token counts, or 0 for an uncaptured endpoint. boundary_checkpoint_tokens: list[list[int]] | None = None @staticmethod diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 99ad320761c8..4a129136e49f 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -150,6 +150,9 @@ def __init__( "The recurrent instruction boundary must be inside the prompt" ) self.recurrent_instruction_boundary = recurrent_instruction_boundary + # The scheduler fixes this once from its prefill budget. Preemption + # must not move the endpoint while workers own its private storage. + self.recurrent_prefill_tail_boundary: int | None = None self._output_token_ids: list[int] = [] if self.prompt_token_ids is None: self._all_token_ids: list[int] = [0] * self.num_prompt_tokens diff --git a/vllm/v1/worker/gpu/boundary_checkpoint.py b/vllm/v1/worker/gpu/boundary_checkpoint.py index 82c16b0628e8..2f6bcec3f87a 100644 --- a/vllm/v1/worker/gpu/boundary_checkpoint.py +++ b/vllm/v1/worker/gpu/boundary_checkpoint.py @@ -13,6 +13,7 @@ INSTRUCTION_CHECKPOINT_SLOT, MAX_BOUNDARY_STOP_TOKENS, NUM_BOUNDARY_CHECKPOINT_SLOTS, + PREFILL_TAIL_CHECKPOINT_SLOT, PROMPT_CHECKPOINT_SLOT, RESPONSE_CHECKPOINT_SLOT, ) @@ -29,6 +30,7 @@ _PROMPT_CHECKPOINT_SLOT = tl.constexpr(PROMPT_CHECKPOINT_SLOT) _RESPONSE_CHECKPOINT_SLOT = tl.constexpr(RESPONSE_CHECKPOINT_SLOT) _INSTRUCTION_CHECKPOINT_SLOT = tl.constexpr(INSTRUCTION_CHECKPOINT_SLOT) +_PREFILL_TAIL_CHECKPOINT_SLOT = tl.constexpr(PREFILL_TAIL_CHECKPOINT_SLOT) @triton.jit @@ -58,12 +60,14 @@ def prepare_boundary_capture( prompt_capture = 0 response_capture = 0 instruction_capture = 0 + prefill_tail_capture = 0 if enabled: prompt_len = tl.load(metadata_ptr + metadata_offset + 1) instruction_len = tl.load(metadata_ptr + metadata_offset + 2) min_len = tl.load(metadata_ptr + metadata_offset + 3) max_len = tl.load(metadata_ptr + metadata_offset + 4) num_stops = tl.load(metadata_ptr + metadata_offset + 5) + prefill_tail_len = tl.load(metadata_ptr + metadata_offset + 6) offsets = tl.arange(0, STOP_CAPACITY) stops = tl.load( stop_tokens_ptr + req_idx * STOP_CAPACITY + offsets, @@ -105,6 +109,19 @@ def prepare_boundary_capture( tl.store( seen_ptr + req_idx * NUM_CAPTURES + _INSTRUCTION_CHECKPOINT_SLOT, 1 ) + if ( + prefill_tail_len > 0 + and num_computed < prefill_tail_len + and computed_after == prefill_tail_len + and tl.load( + seen_ptr + req_idx * NUM_CAPTURES + _PREFILL_TAIL_CHECKPOINT_SLOT + ) + == 0 + ): + prefill_tail_capture = prefill_tail_len + tl.store( + seen_ptr + req_idx * NUM_CAPTURES + _PREFILL_TAIL_CHECKPOINT_SLOT, 1 + ) if ( stopped and tl.load(seen_ptr + req_idx * NUM_CAPTURES + _RESPONSE_CHECKPOINT_SLOT) @@ -131,6 +148,15 @@ def prepare_boundary_capture( tl.maximum(num_sampled - 1, 0), ) tl.store(capture_bias_ptr + output_offset + _INSTRUCTION_CHECKPOINT_SLOT, 0) + tl.store( + capture_tokens_ptr + output_offset + _PREFILL_TAIL_CHECKPOINT_SLOT, + prefill_tail_capture, + ) + tl.store(capture_bias_ptr + output_offset + _PREFILL_TAIL_CHECKPOINT_SLOT, 0) + tl.store( + capture_rows_ptr + output_offset + _PREFILL_TAIL_CHECKPOINT_SLOT, + query_start + prefill_tail_capture - num_computed - 1, + ) tl.store( capture_rows_ptr + output_offset + _PROMPT_CHECKPOINT_SLOT, query_start + prompt_capture - num_computed - 1, @@ -295,7 +321,7 @@ def __init__( for group in kv_cache_config.kv_cache_groups ] self.metadata = torch.zeros( - (self.max_reqs, 6), dtype=torch.int32, device=self.device + (self.max_reqs, 7), dtype=torch.int32, device=self.device ) self.stop_tokens = torch.full( (self.max_reqs, MAX_BOUNDARY_STOP_TOKENS), @@ -600,6 +626,7 @@ def add_request(self, slot: int, request: "NewRequestData") -> None: self.model_state.model_config.max_model_len, ), len(stops), + request.recurrent_prefill_tail_boundary or 0, ], dtype=torch.int32, device=self.device, @@ -631,6 +658,11 @@ def add_request(self, slot: int, request: "NewRequestData") -> None: and checkpoint.num_tokens >= request.recurrent_instruction_boundary ): self.seen[slot, INSTRUCTION_CHECKPOINT_SLOT] = 1 + if ( + request.recurrent_prefill_tail_boundary is not None + and checkpoint.num_tokens >= request.recurrent_prefill_tail_boundary + ): + self.seen[slot, PREFILL_TAIL_CHECKPOINT_SLOT] = 1 _restore_auxiliary_state_kernel[(self.auxiliary_metadata.shape[0],)]( self.auxiliary_metadata, self.pool, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index eae20c697116..32e6def84926 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -575,6 +575,7 @@ def _post_update_kernel( boundary_capture_bias_ptr=None, boundary_capture_rows_ptr=None, NUM_CAPTURES: tl.constexpr = 0, + BOUNDARY_METADATA_WIDTH: tl.constexpr = 0, ): req_id = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_id) @@ -616,7 +617,7 @@ def _post_update_kernel( boundary_capture_rows_ptr, STOP_CAPACITY=128, NUM_CAPTURES=NUM_CAPTURES, - METADATA_WIDTH=6, + METADATA_WIDTH=BOUNDARY_METADATA_WIDTH, ) tl.store(num_sampled_ptr + req_id, num_sampled) tl.store(num_rejected_ptr + req_id, num_rejected) @@ -696,6 +697,9 @@ def post_update( NUM_CAPTURES=( NUM_BOUNDARY_CHECKPOINT_SLOTS if boundary_capture is not None else 0 ), + BOUNDARY_METADATA_WIDTH=( + boundary_state.metadata.shape[1] if boundary_state is not None else 0 + ), num_warps=1, ) From 42f0491173f788d3fc1eef84d2f1c7142f57bb52 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 13:51:41 +0000 Subject: [PATCH 113/365] fix(moe): retain immutable B12X input-scale ownership Pass the loaded inference-scale contract through PackedWeights so exact uniform per-expert scales remain eligible for shared NVFP4 input quantization. Keep canonical vectors and the B12X equality and mutation guards intact. Add a host-only integration regression for uniform and one-ULP-different scales. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 54 ++++++++++++++++++++ vllm/model_executor/layers/fused_moe/b12x.py | 3 ++ 2 files changed, 57 insertions(+) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index e844dc362733..c78f8501c265 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -795,6 +795,60 @@ def prepare_for_b12x(**kwargs): torch.testing.assert_close(prepared[7], expected_a2_scale) +@pytest.mark.parametrize("uniform", [True, False]) +def test_b12x_nvfp4_preparation_preserves_static_scale_contract( + monkeypatch: pytest.MonkeyPatch, uniform: bool +) -> None: + """Loaded uniform scales admit shared quantization without losing the vector.""" + fused_moe = pytest.importorskip("b12x.moe.fused_moe") + from b12x.moe.fused_moe._impl import B12XFP4ExpertWeights + + # Keep the real scale-ownership check; only CUDA weight packing is excluded. + def prepare_weights(*, plan, weights): + return B12XFP4ExpertWeights( + plan=plan._impl, + a1_gscale=weights.input_scale, + a2_gscale=weights.intermediate_scale, + w1_fp4=weights.w13, + w1_blockscale=weights.w13_block_scales, + w1_alphas=weights.w13_global_scales, + w2_fp4=weights.w2, + w2_blockscale=weights.w2_block_scales, + w2_alphas=weights.w2_global_scales, + immutable_input_scales=weights.immutable_input_scales, + ) + + monkeypatch.setattr(fused_moe, "prepare_weights", prepare_weights) + monkeypatch.setattr(b12x, "_require_b12x_fused_moe", lambda: fused_moe) + monkeypatch.setattr(b12x, "_is_current_stream_capturing", lambda: False) + scales = torch.ones(4) + if not uniform: + scales[-1] = torch.nextafter(scales[-1], torch.tensor(float("inf"))) + config = make_dummy_moe_config(num_experts=4, hidden_dim=128, intermediate_size=128) + quant = FusedMoEQuantConfig.make( + "nvfp4", + w1_scale=torch.ones(4, 256, 8, dtype=torch.float8_e4m3fn), + w2_scale=torch.ones(4, 128, 8, dtype=torch.float8_e4m3fn), + g1_alphas=torch.ones(4), + g2_alphas=torch.ones(4), + a1_gscale=scales, + a2_gscale=torch.ones(4), + ) + experts = B12xExperts(config, quant) + prepared = experts._prepare_experts( + w1=torch.empty(4, 256, 64, dtype=torch.uint8), + w2=torch.empty(4, 128, 64, dtype=torch.uint8), + activation=MoEActivation.SILU, + params_dtype=torch.bfloat16, + ) + + assert prepared.a1_gscale is scales + assert prepared.can_share_input(input_scales_static=True) is uniform + assert not prepared.can_share_input(input_scales_static=False) + scales[-1] = 2.0 + assert not prepared.can_share_input(input_scales_static=True) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_b12x_nvfp4_preparation_pads_each_gated_half() -> None: device = torch.device("cuda") diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 28d57bba7fef..5fa3d8fa79d0 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -394,6 +394,9 @@ def _prepare_experts( w2_global_scales=w2_global_scale, input_scale=a1_gscale, intermediate_scale=a2_gscale, + # Loaded scales are fixed while prepared plans and their + # captured CUDA graphs execute. + immutable_input_scales=True, ), ) From 3e5c7d8b283e6975ae0925727d6e72dedc16baa0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 16:10:37 +0000 Subject: [PATCH 114/365] fix(b12x): declare loaded DeepSeek WO weight contracts Pass tile-major storage availability and exact block-scale replication from packed loader weights into WO preparation. B12X retains responsibility for decode tile selection. This requires B12X WO packed-weight declaration support. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/v1/attention/test_b12x_sparse_mla_api.py | 10 ++++++++-- vllm/models/deepseek_v4/nvidia/b12x.py | 2 ++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 534f971261d3..83b4690dae01 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1912,7 +1912,12 @@ def test_deepseek_v4_wo_declares_exact_rows_before_profiling(monkeypatch, eager_ layer = _deepseek_v4_wo_layer(torch.device("cpu")) layer._b12x_wo_projection_weights = SimpleNamespace( - groups=2, group_width=4096, rank=128, hidden=256 + groups=2, + group_width=4096, + rank=128, + hidden=256, + sfb_k_replicated=True, + wo_b=SimpleNamespace(values_tiled=None), ) monkeypatch.setattr(b12x_mla, "_require_b12x_wo_projection", lambda: wo_projection) workload = B12xWorkload( @@ -1939,6 +1944,7 @@ def test_deepseek_v4_wo_declares_exact_rows_before_profiling(monkeypatch, eager_ assert query.operation == "inv_rope" and not query.dynamic_tokens assert (query.heads_per_group, query.nope_dim, query.rope_dim) == (8, 448, 64) assert query.positions_dtype == "int64" and query.cos_sin_dtype == "bfloat16" + assert query.sfb_k_replicated and not query.wo_b_tiled (repeated,) = layer.get_b12x_preparation_units(layer, workload) assert all(a.plan is b.plan for a, b in zip(unit.requests, repeated.requests)) assert ( @@ -2113,7 +2119,7 @@ def test_mhc_prepares_model_operands_and_first_layer_broadcast( layer = _deepseek_v4_mhc_layer(first_layer) hidden = layer.attn_norm.weight.numel() operands = None - operations = ("post_pre", "post_pre_bf16", "post") + operations: tuple[str, ...] = ("post_pre", "post_pre_bf16", "post") if glm_operands: layer.input_layernorm = layer.attn_norm layer.post_attention_layernorm = layer.ffn_norm diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 98c210d79767..af83cdd2acee 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -856,6 +856,8 @@ def _b12x_wo_plan(self, rows: int): rope_dim=self.rope_head_dim, positions_dtype="int64", cos_sin_dtype=str(table.dtype).removeprefix("torch."), + sfb_k_replicated=weights.sfb_k_replicated, + wo_b_tiled=weights.wo_b.values_tiled is not None, ), ) self._b12x_wo_plans[rows] = plan From eb2a442a1fa9bd79a22cd9937125bcd9bf1faef1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 18:26:56 +0000 Subject: [PATCH 115/365] fix(b12x): retain packed first-projection storage in WO plans Forward loader-owned tile-major WO-A availability alongside WO-B and scale provenance. Cover generic and packed declarations before graph capture. Requires the matching B12X WO planning contract. Co-authored-by: Codex Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 22 +++++++++++++------ vllm/models/deepseek_v4/nvidia/b12x.py | 1 + 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 83b4690dae01..576d4be37f3b 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -1903,21 +1903,28 @@ def _deepseek_v4_wo_layer(device, groups=2, heads_per_group=8, rank=128, hidden= @pytest.mark.parametrize("eager_only", [False, True]) -def test_deepseek_v4_wo_declares_exact_rows_before_profiling(monkeypatch, eager_only): +@pytest.mark.parametrize("tiled", [False, True]) +def test_deepseek_v4_wo_declares_exact_rows_before_profiling( + monkeypatch, eager_only, tiled +): from dataclasses import replace from b12x.gemm import wo_projection from vllm.utils.b12x import B12xWorkload - layer = _deepseek_v4_wo_layer(torch.device("cpu")) + groups, rank, hidden = (4, 1024, 4096) if tiled else (2, 128, 256) + layer = _deepseek_v4_wo_layer( + torch.device("cpu"), groups=groups, rank=rank, hidden=hidden + ) layer._b12x_wo_projection_weights = SimpleNamespace( - groups=2, + groups=groups, group_width=4096, - rank=128, - hidden=256, + rank=rank, + hidden=hidden, sfb_k_replicated=True, - wo_b=SimpleNamespace(values_tiled=None), + wo_a=SimpleNamespace(values_tiled=object() if tiled else None), + wo_b=SimpleNamespace(values_tiled=object() if tiled else None), ) monkeypatch.setattr(b12x_mla, "_require_b12x_wo_projection", lambda: wo_projection) workload = B12xWorkload( @@ -1944,7 +1951,8 @@ def test_deepseek_v4_wo_declares_exact_rows_before_profiling(monkeypatch, eager_ assert query.operation == "inv_rope" and not query.dynamic_tokens assert (query.heads_per_group, query.nope_dim, query.rope_dim) == (8, 448, 64) assert query.positions_dtype == "int64" and query.cos_sin_dtype == "bfloat16" - assert query.sfb_k_replicated and not query.wo_b_tiled + assert query.sfb_k_replicated + assert query.wo_a_tiled is tiled and query.wo_b_tiled is tiled (repeated,) = layer.get_b12x_preparation_units(layer, workload) assert all(a.plan is b.plan for a, b in zip(unit.requests, repeated.requests)) assert ( diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index af83cdd2acee..6eb258527d3c 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -857,6 +857,7 @@ def _b12x_wo_plan(self, rows: int): positions_dtype="int64", cos_sin_dtype=str(table.dtype).removeprefix("torch."), sfb_k_replicated=weights.sfb_k_replicated, + wo_a_tiled=weights.wo_a.values_tiled is not None, wo_b_tiled=weights.wo_b.values_tiled is not None, ), ) From 89694ddb4fc3ad8c7270e8769dd938062ac730ed Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 21:10:59 +0000 Subject: [PATCH 116/365] perf(moe): measure routed candidates with shared FP8 experts Include the joined shared-expert MLP in small W4A8 preparation races when serving overlaps it with routed experts. Keep serving execution and numerical kernels unchanged, and separate cached choices by the shared branch execution contract. Unsupported shared paths retain isolated preparation. Validation: nine CPU metadata and stream-order tests pass; applicable pre-commit hooks including mypy pass. GPU serving qualification remains pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 146 ++++++++++++++++++ vllm/model_executor/layers/fused_moe/b12x.py | 101 +++++++++++- .../layers/fused_moe/runner/moe_runner.py | 9 ++ 3 files changed, 248 insertions(+), 8 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index e844dc362733..9e86d8451087 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -6,6 +6,7 @@ import gc import weakref +from contextlib import contextmanager from dataclasses import dataclass, replace from types import SimpleNamespace from typing import TYPE_CHECKING @@ -65,6 +66,151 @@ from b12x.preparation import PreparationSession +def _shared_fp8_tuning_fixture(): + mlp = torch.nn.Module() + mlp.act_fn = torch.nn.SiLU() + for name, shape in (("gate_up_proj", (16, 8)), ("down_proj", (8, 8))): + linear = torch.nn.Module() + linear.weight = torch.empty(shape, dtype=torch.float8_e4m3fn) + linear.scale = torch.ones((1, 1)) + linear.reduce_results = False + linear.deep_gemm_warmup_provider = SimpleNamespace( + get_deep_gemm_warmup_weights=lambda layer: (layer.weight, layer.scale) + ) + setattr(mlp, name, linear) + shared = torch.nn.Module() + shared._layer = mlp + shared._stream = object() + shared._disable_shared_experts_overlap = False + shared._mk_can_overlap_shared_experts = lambda: False + layer = torch.nn.Module() + layer.shared_experts_for_preparation = weakref.ref(shared) + return layer, shared + + +def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + other_layer, other_shared = _shared_fp8_tuning_fixture() + found, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert found is shared + assert key == b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] + other_shared._layer.gate_up_proj.weight = torch.empty( + (32, 8), dtype=torch.float8_e4m3fn + ) + assert key != b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] + assert list(layer.modules()) == [layer] + + +@pytest.mark.parametrize( + "exclusion", + ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], +) +def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + quant_mode, hidden_size = "w4a8_mx", 8 + if exclusion == "stream": + shared._stream = None + elif exclusion == "internal": + shared._mk_can_overlap_shared_experts = lambda: True + elif exclusion == "parallel": + shared._disable_shared_experts_overlap = True + elif exclusion == "collective": + shared._layer.down_proj.reduce_results = True + elif exclusion == "dtype": + shared._layer.down_proj.weight = torch.empty((8, 8)) + elif exclusion == "width": + hidden_size = 16 + else: + quant_mode = "nvfp4" + assert b12x._shared_expert_tuning_context(layer, quant_mode, hidden_size) == ( + None, + None, + ) + + +def test_shared_fp8_tuning_joins_auxiliary_work_before_timing_ends(monkeypatch): + from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExpertsOrder, + ) + + events: list[object] = [] + + class Stream: + def __init__(self, name): + self.name = name + + def wait_stream(self, other): + events.append((self.name, "wait", other.name)) + + primary, auxiliary = Stream("primary"), Stream("auxiliary") + + @contextmanager + def use_stream(stream): + events.append(("enter", stream.name)) + yield + events.append(("exit", stream.name)) + + class Output: + def record_stream(self, stream): + events.append(("record", stream.name)) + + def shared_forward(hidden): + events.append(("shared", tuple(hidden.shape))) + return Output() + + def bind(**kwargs): + def run(): + events.append("routed") + kwargs["output"].fill_(1) + + return SimpleNamespace(run=run) + + monkeypatch.setattr(torch.cuda, "current_stream", lambda: primary) + monkeypatch.setattr(torch.cuda, "stream", use_stream) + state = SimpleNamespace( + scratch=SimpleNamespace(scratch_specs=lambda: ()), bind=bind + ) + shared = SimpleNamespace( + _stream=auxiliary, + _layer=shared_forward, + _determine_shared_experts_order=lambda _: ( + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED + ), + ) + call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, shared) + hidden = torch.zeros((2, 8), dtype=torch.bfloat16) + ids = torch.zeros((2, 2), dtype=torch.int32) + weights = torch.ones((2, 2)) + prepared = call.make( + ( + hidden, + hidden.clone(), + torch.empty_like(hidden), + ids[None], + weights, + ids, + weights, + ) + ) + prepared.run() + assert events == [ + ("auxiliary", "wait", "primary"), + ("enter", "auxiliary"), + ("shared", (2, 8)), + ("exit", "auxiliary"), + "routed", + ("primary", "wait", "auxiliary"), + ("record", "primary"), + ] + assert torch.all(prepared.output == 1) + + def _prepare( layer, *, diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 28d57bba7fef..5bf8912db2ac 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -10,6 +10,7 @@ import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -79,6 +80,7 @@ class _PreparedMoECall: topk: int prepared: Any output_dtype: torch.dtype + shared_experts: Any | None = None def make(self, tensors): from b12x.preparation import PreparedCall @@ -123,8 +125,29 @@ def restore() -> None: output=output, input_scales_static=True, ) + run = binding.run + if self.shared_experts is not None: + from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExpertsOrder, + ) + + shared = self.shared_experts + if shared._determine_shared_experts_order(hidden) == ( + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED + ): + + def run() -> None: + primary = torch.cuda.current_stream() + auxiliary = shared._stream + auxiliary.wait_stream(primary) + with torch.cuda.stream(auxiliary): + shared_output = shared._layer(hidden) + binding.run() + primary.wait_stream(auxiliary) + shared_output.record_stream(primary) + return PreparedCall( - run=binding.run, + run=run, output=output, produce=produce, reset=reset, @@ -137,7 +160,12 @@ def restore() -> None: def _prepared_moe_call_factory( - *, tokens: int, topk: int, prepared: Any, output_dtype: torch.dtype + *, + tokens: int, + topk: int, + prepared: Any, + output_dtype: torch.dtype, + shared_experts: Any | None = None, ): shared = None @@ -190,11 +218,61 @@ def factory(state: Any): topk=topk, prepared=prepared, output_dtype=output_dtype, + shared_experts=shared_experts if tokens <= 8 else None, ).make(tensors) return factory +def _shared_expert_tuning_context( + layer: torch.nn.Module, quant_mode: str, hidden_size: int +): + """Describe the independently executable block-FP8 MLP used during tuning.""" + if quant_mode != "w4a8_mx": + return None, None + reference = getattr(layer, "shared_experts_for_preparation", None) + shared = None if reference is None else reference() + if ( + shared is None + or shared._stream is None + or shared._disable_shared_experts_overlap + or shared._mk_can_overlap_shared_experts() + ): + return None, None + from b12x.preparation import FrozenMapping + + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used + + mlp = shared._layer + if not callable(getattr(mlp, "act_fn", None)): + return None, None + projections = [] + for name in ("gate_up_proj", "down_proj"): + linear = getattr(mlp, name, None) + provider = getattr(linear, "deep_gemm_warmup_provider", None) + get_weights = getattr(provider, "get_deep_gemm_warmup_weights", None) + if not callable(get_weights) or getattr(linear, "reduce_results", False): + return None, None + weight, scale = get_weights(linear) + if weight.dtype != torch.float8_e4m3fn or weight.ndim != 2: + return None, None + projections.append((tuple(weight.shape), tuple(scale.shape), str(scale.dtype))) + if not (projections[0][0][1] == projections[1][0][0] == hidden_size): + return None, None + descriptor = FrozenMapping( + { + "version": 1, + "backend": "deep_gemm_block_fp8_mlp", + "projections": projections, + "activation": type(mlp.act_fn).__qualname__, + "ue8m0": is_deep_gemm_e8m0_used(), + "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, + "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), + } + ) + return shared, descriptor + + def _is_current_stream_capturing() -> bool: is_capturing = getattr(torch.cuda, "is_current_stream_capturing", None) return bool(is_capturing is not None and is_capturing()) @@ -634,9 +712,17 @@ def get_b12x_preparation_units( topk = int(self.moe_config.experts_per_token) route_on_input = bool(layer.apply_router_weight_on_input) fused_moe = _require_b12x_fused_moe() - plan_key = (counts, activation, route_on_input, id(prepared)) + shared, shared_context = _shared_expert_tuning_context( + layer, self._quant_mode, int(prepared.hidden_size) + ) + plan_key = (counts, activation, route_on_input, id(prepared), shared_context) plan = self._plan if getattr(self, "_plan_key", None) == plan_key else None if plan is None: + invocation: dict[str, Any] = { + "tuning_route_pattern": TUNING_WORKLOAD_VERSION + } + if shared_context is not None: + invocation["shared_expert_context"] = shared_context # The layer holds one plan for its serving shapes; a later call with # the same workload reuses it so the prepared state stays installed. plan = fused_moe.plan_execution( @@ -653,11 +739,7 @@ def get_b12x_preparation_units( # Tuning choices depend on the routing corpus. Invalidate only # MoE choices when its distribution changes, not compiled code # or unrelated component selections. - invocation=FrozenMapping( - { - "tuning_route_pattern": TUNING_WORKLOAD_VERSION, - } - ), + invocation=FrozenMapping(invocation), ) self._plan = plan self._plan_key = plan_key @@ -680,6 +762,7 @@ def get_b12x_preparation_units( topk=topk, prepared=prepared, output_dtype=self.output_dtype, + shared_experts=shared, ) for count in plan.token_counts } @@ -700,6 +783,7 @@ def get_b12x_preparation_units( topk=topk, prepared=prepared, output_dtype=self.output_dtype, + shared_experts=shared, ) request = plan.request( name=name, @@ -714,6 +798,7 @@ def get_b12x_preparation_units( route_on_input, counts, workload.output_dtype, + shared_context, ) return ( B12xPreparationUnit( diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 3dbb76215c2f..836c7ddb1195 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref from collections.abc import Callable, Iterable from contextlib import nullcontext from typing import TYPE_CHECKING, cast @@ -297,6 +298,14 @@ def __init__( enable_dbo=enable_dbo, mk_can_overlap_shared_experts=can_overlap, ) + if routed_input_transform is None and shared_expert_gate is None: + # Preparation can time the independent branch without registering + # its parameters a second time or extending the runner's lifetime. + object.__setattr__( + routed_experts, + "shared_experts_for_preparation", + weakref.ref(self._shared_experts), + ) # Needed for string -> MoERunner layer lookup in custom ops. self.layer_name = layer_name From 8c5db18614690843b2fdc95b822a856da0eabf20 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 21:36:47 +0000 Subject: [PATCH 117/365] perf(moe): include loaded gate work in shared-expert tuning Execute the eligible BF16 router projection before each routed candidate while the FP8 shared branch runs. Preserve the serving order without an artificial delay or a hardware-specific grid override. Include gate geometry and dispatch in the preparation cache identity. The selector remains represented by the existing synthetic routing corpus. Validation: 15 CPU contract and bounded-ownership tests pass; applicable pre-commit hooks including mypy pass. GPU serving qualification is pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 32 +++++++++++++++- vllm/model_executor/layers/fused_moe/b12x.py | 38 +++++++++++++++++-- .../layers/fused_moe/runner/moe_runner.py | 6 +++ 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index 9e86d8451087..aa984d93b528 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -95,7 +95,7 @@ def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch layer, shared = _shared_fp8_tuning_fixture() other_layer, other_shared = _shared_fp8_tuning_fixture() found, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) - assert found is shared + assert found.shared is shared and found.gate is None assert key == b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] other_shared._layer.gate_up_proj.weight = torch.empty( (32, 8), dtype=torch.float8_e4m3fn @@ -104,6 +104,28 @@ def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch assert list(layer.modules()) == [layer] +def test_shared_fp8_tuning_gate_identity_includes_loaded_geometry(monkeypatch): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + without_gate = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + gate = torch.nn.Module() + gate.weight = torch.empty((4, 8), dtype=torch.bfloat16) + gate.out_dtype = torch.float32 + gate.allow_ll_bf16_gemm = True + layer.routing_gate_for_preparation = weakref.ref(gate) + context, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert context.shared is shared and context.gate is gate + assert key != without_gate + gate.weight = torch.empty((8, 8), dtype=torch.bfloat16) + assert key != b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + gate.allow_ll_bf16_gemm = False + context, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert context.gate is None and key == without_gate + assert list(layer.modules()) == [layer] + + @pytest.mark.parametrize( "exclusion", ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], @@ -183,7 +205,12 @@ def run(): SharedExpertsOrder.MULTI_STREAM_OVERLAPPED ), ) - call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, shared) + + def gate(hidden): + events.append(("gate", tuple(hidden.shape))) + + context = b12x._SharedExpertTuning(shared, gate) + call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, context) hidden = torch.zeros((2, 8), dtype=torch.bfloat16) ids = torch.zeros((2, 2), dtype=torch.int32) weights = torch.ones((2, 2)) @@ -204,6 +231,7 @@ def run(): ("enter", "auxiliary"), ("shared", (2, 8)), ("exit", "auxiliary"), + ("gate", (2, 8)), "routed", ("primary", "wait", "auxiliary"), ("record", "primary"), diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 5bf8912db2ac..ad0fd43e3768 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -67,6 +67,12 @@ def _b12x_activation_name(activation: MoEActivation) -> str: return activation.value +@dataclass(frozen=True) +class _SharedExpertTuning: + shared: Any + gate: Any | None = None + + @dataclass(frozen=True) class _PreparedMoECall: """Priming tensors for one exact prepared MoE variant. @@ -131,7 +137,8 @@ def restore() -> None: SharedExpertsOrder, ) - shared = self.shared_experts + shared = self.shared_experts.shared + gate = self.shared_experts.gate if shared._determine_shared_experts_order(hidden) == ( SharedExpertsOrder.MULTI_STREAM_OVERLAPPED ): @@ -142,6 +149,8 @@ def run() -> None: auxiliary.wait_stream(primary) with torch.cuda.stream(auxiliary): shared_output = shared._layer(hidden) + if gate is not None: + gate(hidden) binding.run() primary.wait_stream(auxiliary) shared_output.record_stream(primary) @@ -259,18 +268,41 @@ def _shared_expert_tuning_context( projections.append((tuple(weight.shape), tuple(scale.shape), str(scale.dtype))) if not (projections[0][0][1] == projections[1][0][0] == hidden_size): return None, None + gate_reference = getattr(layer, "routing_gate_for_preparation", None) + gate = None if gate_reference is None else gate_reference() + gate_weight = getattr(gate, "weight", None) + if not ( + isinstance(gate_weight, torch.Tensor) + and gate_weight.dtype == torch.bfloat16 + and gate_weight.ndim == 2 + and gate_weight.shape[1] == hidden_size + and getattr(gate, "out_dtype", None) == torch.float32 + and getattr(gate, "allow_ll_bf16_gemm", False) + ): + gate = None + gate_context = ( + None + if gate is None + else { + "backend": "ll_bf16_router", + "shape": tuple(gate.weight.shape), + "dtype": str(gate.weight.dtype), + "output_dtype": str(gate.out_dtype), + } + ) descriptor = FrozenMapping( { - "version": 1, + "version": 2, "backend": "deep_gemm_block_fp8_mlp", "projections": projections, "activation": type(mlp.act_fn).__qualname__, "ue8m0": is_deep_gemm_e8m0_used(), "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), + "gate": gate_context, } ) - return shared, descriptor + return _SharedExpertTuning(shared, gate), descriptor def _is_current_stream_capturing() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 836c7ddb1195..4467872fee66 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -306,6 +306,12 @@ def __init__( "shared_experts_for_preparation", weakref.ref(self._shared_experts), ) + if gate is not None: + object.__setattr__( + routed_experts, + "routing_gate_for_preparation", + weakref.ref(gate), + ) # Needed for string -> MoERunner layer lookup in custom ops. self.layer_name = layer_name From bc1921bef26f4d6577a7348d31d2a78b8a4651b2 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 21:10:59 +0000 Subject: [PATCH 118/365] perf(moe): measure routed candidates with shared FP8 experts Include the joined shared-expert MLP in small W4A8 preparation races when serving overlaps it with routed experts. Keep serving execution and numerical kernels unchanged, and separate cached choices by the shared branch execution contract. Unsupported shared paths retain isolated preparation. Validation: nine CPU metadata and stream-order tests pass; applicable pre-commit hooks including mypy pass. GPU serving qualification remains pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 146 ++++++++++++++++++ vllm/model_executor/layers/fused_moe/b12x.py | 101 +++++++++++- .../layers/fused_moe/runner/moe_runner.py | 9 ++ 3 files changed, 248 insertions(+), 8 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index c78f8501c265..b3672dd50d9f 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -6,6 +6,7 @@ import gc import weakref +from contextlib import contextmanager from dataclasses import dataclass, replace from types import SimpleNamespace from typing import TYPE_CHECKING @@ -65,6 +66,151 @@ from b12x.preparation import PreparationSession +def _shared_fp8_tuning_fixture(): + mlp = torch.nn.Module() + mlp.act_fn = torch.nn.SiLU() + for name, shape in (("gate_up_proj", (16, 8)), ("down_proj", (8, 8))): + linear = torch.nn.Module() + linear.weight = torch.empty(shape, dtype=torch.float8_e4m3fn) + linear.scale = torch.ones((1, 1)) + linear.reduce_results = False + linear.deep_gemm_warmup_provider = SimpleNamespace( + get_deep_gemm_warmup_weights=lambda layer: (layer.weight, layer.scale) + ) + setattr(mlp, name, linear) + shared = torch.nn.Module() + shared._layer = mlp + shared._stream = object() + shared._disable_shared_experts_overlap = False + shared._mk_can_overlap_shared_experts = lambda: False + layer = torch.nn.Module() + layer.shared_experts_for_preparation = weakref.ref(shared) + return layer, shared + + +def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + other_layer, other_shared = _shared_fp8_tuning_fixture() + found, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert found is shared + assert key == b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] + other_shared._layer.gate_up_proj.weight = torch.empty( + (32, 8), dtype=torch.float8_e4m3fn + ) + assert key != b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] + assert list(layer.modules()) == [layer] + + +@pytest.mark.parametrize( + "exclusion", + ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], +) +def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + quant_mode, hidden_size = "w4a8_mx", 8 + if exclusion == "stream": + shared._stream = None + elif exclusion == "internal": + shared._mk_can_overlap_shared_experts = lambda: True + elif exclusion == "parallel": + shared._disable_shared_experts_overlap = True + elif exclusion == "collective": + shared._layer.down_proj.reduce_results = True + elif exclusion == "dtype": + shared._layer.down_proj.weight = torch.empty((8, 8)) + elif exclusion == "width": + hidden_size = 16 + else: + quant_mode = "nvfp4" + assert b12x._shared_expert_tuning_context(layer, quant_mode, hidden_size) == ( + None, + None, + ) + + +def test_shared_fp8_tuning_joins_auxiliary_work_before_timing_ends(monkeypatch): + from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExpertsOrder, + ) + + events: list[object] = [] + + class Stream: + def __init__(self, name): + self.name = name + + def wait_stream(self, other): + events.append((self.name, "wait", other.name)) + + primary, auxiliary = Stream("primary"), Stream("auxiliary") + + @contextmanager + def use_stream(stream): + events.append(("enter", stream.name)) + yield + events.append(("exit", stream.name)) + + class Output: + def record_stream(self, stream): + events.append(("record", stream.name)) + + def shared_forward(hidden): + events.append(("shared", tuple(hidden.shape))) + return Output() + + def bind(**kwargs): + def run(): + events.append("routed") + kwargs["output"].fill_(1) + + return SimpleNamespace(run=run) + + monkeypatch.setattr(torch.cuda, "current_stream", lambda: primary) + monkeypatch.setattr(torch.cuda, "stream", use_stream) + state = SimpleNamespace( + scratch=SimpleNamespace(scratch_specs=lambda: ()), bind=bind + ) + shared = SimpleNamespace( + _stream=auxiliary, + _layer=shared_forward, + _determine_shared_experts_order=lambda _: ( + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED + ), + ) + call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, shared) + hidden = torch.zeros((2, 8), dtype=torch.bfloat16) + ids = torch.zeros((2, 2), dtype=torch.int32) + weights = torch.ones((2, 2)) + prepared = call.make( + ( + hidden, + hidden.clone(), + torch.empty_like(hidden), + ids[None], + weights, + ids, + weights, + ) + ) + prepared.run() + assert events == [ + ("auxiliary", "wait", "primary"), + ("enter", "auxiliary"), + ("shared", (2, 8)), + ("exit", "auxiliary"), + "routed", + ("primary", "wait", "auxiliary"), + ("record", "primary"), + ] + assert torch.all(prepared.output == 1) + + def _prepare( layer, *, diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 5fa3d8fa79d0..4758575f6b4f 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -10,6 +10,7 @@ import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -79,6 +80,7 @@ class _PreparedMoECall: topk: int prepared: Any output_dtype: torch.dtype + shared_experts: Any | None = None def make(self, tensors): from b12x.preparation import PreparedCall @@ -123,8 +125,29 @@ def restore() -> None: output=output, input_scales_static=True, ) + run = binding.run + if self.shared_experts is not None: + from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExpertsOrder, + ) + + shared = self.shared_experts + if shared._determine_shared_experts_order(hidden) == ( + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED + ): + + def run() -> None: + primary = torch.cuda.current_stream() + auxiliary = shared._stream + auxiliary.wait_stream(primary) + with torch.cuda.stream(auxiliary): + shared_output = shared._layer(hidden) + binding.run() + primary.wait_stream(auxiliary) + shared_output.record_stream(primary) + return PreparedCall( - run=binding.run, + run=run, output=output, produce=produce, reset=reset, @@ -137,7 +160,12 @@ def restore() -> None: def _prepared_moe_call_factory( - *, tokens: int, topk: int, prepared: Any, output_dtype: torch.dtype + *, + tokens: int, + topk: int, + prepared: Any, + output_dtype: torch.dtype, + shared_experts: Any | None = None, ): shared = None @@ -190,11 +218,61 @@ def factory(state: Any): topk=topk, prepared=prepared, output_dtype=output_dtype, + shared_experts=shared_experts if tokens <= 8 else None, ).make(tensors) return factory +def _shared_expert_tuning_context( + layer: torch.nn.Module, quant_mode: str, hidden_size: int +): + """Describe the independently executable block-FP8 MLP used during tuning.""" + if quant_mode != "w4a8_mx": + return None, None + reference = getattr(layer, "shared_experts_for_preparation", None) + shared = None if reference is None else reference() + if ( + shared is None + or shared._stream is None + or shared._disable_shared_experts_overlap + or shared._mk_can_overlap_shared_experts() + ): + return None, None + from b12x.preparation import FrozenMapping + + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used + + mlp = shared._layer + if not callable(getattr(mlp, "act_fn", None)): + return None, None + projections = [] + for name in ("gate_up_proj", "down_proj"): + linear = getattr(mlp, name, None) + provider = getattr(linear, "deep_gemm_warmup_provider", None) + get_weights = getattr(provider, "get_deep_gemm_warmup_weights", None) + if not callable(get_weights) or getattr(linear, "reduce_results", False): + return None, None + weight, scale = get_weights(linear) + if weight.dtype != torch.float8_e4m3fn or weight.ndim != 2: + return None, None + projections.append((tuple(weight.shape), tuple(scale.shape), str(scale.dtype))) + if not (projections[0][0][1] == projections[1][0][0] == hidden_size): + return None, None + descriptor = FrozenMapping( + { + "version": 1, + "backend": "deep_gemm_block_fp8_mlp", + "projections": projections, + "activation": type(mlp.act_fn).__qualname__, + "ue8m0": is_deep_gemm_e8m0_used(), + "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, + "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), + } + ) + return shared, descriptor + + def _is_current_stream_capturing() -> bool: is_capturing = getattr(torch.cuda, "is_current_stream_capturing", None) return bool(is_capturing is not None and is_capturing()) @@ -637,9 +715,17 @@ def get_b12x_preparation_units( topk = int(self.moe_config.experts_per_token) route_on_input = bool(layer.apply_router_weight_on_input) fused_moe = _require_b12x_fused_moe() - plan_key = (counts, activation, route_on_input, id(prepared)) + shared, shared_context = _shared_expert_tuning_context( + layer, self._quant_mode, int(prepared.hidden_size) + ) + plan_key = (counts, activation, route_on_input, id(prepared), shared_context) plan = self._plan if getattr(self, "_plan_key", None) == plan_key else None if plan is None: + invocation: dict[str, Any] = { + "tuning_route_pattern": TUNING_WORKLOAD_VERSION + } + if shared_context is not None: + invocation["shared_expert_context"] = shared_context # The layer holds one plan for its serving shapes; a later call with # the same workload reuses it so the prepared state stays installed. plan = fused_moe.plan_execution( @@ -656,11 +742,7 @@ def get_b12x_preparation_units( # Tuning choices depend on the routing corpus. Invalidate only # MoE choices when its distribution changes, not compiled code # or unrelated component selections. - invocation=FrozenMapping( - { - "tuning_route_pattern": TUNING_WORKLOAD_VERSION, - } - ), + invocation=FrozenMapping(invocation), ) self._plan = plan self._plan_key = plan_key @@ -683,6 +765,7 @@ def get_b12x_preparation_units( topk=topk, prepared=prepared, output_dtype=self.output_dtype, + shared_experts=shared, ) for count in plan.token_counts } @@ -703,6 +786,7 @@ def get_b12x_preparation_units( topk=topk, prepared=prepared, output_dtype=self.output_dtype, + shared_experts=shared, ) request = plan.request( name=name, @@ -717,6 +801,7 @@ def get_b12x_preparation_units( route_on_input, counts, workload.output_dtype, + shared_context, ) return ( B12xPreparationUnit( diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 3dbb76215c2f..836c7ddb1195 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref from collections.abc import Callable, Iterable from contextlib import nullcontext from typing import TYPE_CHECKING, cast @@ -297,6 +298,14 @@ def __init__( enable_dbo=enable_dbo, mk_can_overlap_shared_experts=can_overlap, ) + if routed_input_transform is None and shared_expert_gate is None: + # Preparation can time the independent branch without registering + # its parameters a second time or extending the runner's lifetime. + object.__setattr__( + routed_experts, + "shared_experts_for_preparation", + weakref.ref(self._shared_experts), + ) # Needed for string -> MoERunner layer lookup in custom ops. self.layer_name = layer_name From 6902d06cdd5d9ee0f7b3cf29b36dd30a7555603b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 21:36:47 +0000 Subject: [PATCH 119/365] perf(moe): include loaded gate work in shared-expert tuning Execute the eligible BF16 router projection before each routed candidate while the FP8 shared branch runs. Preserve the serving order without an artificial delay or a hardware-specific grid override. Include gate geometry and dispatch in the preparation cache identity. The selector remains represented by the existing synthetic routing corpus. Validation: 15 CPU contract and bounded-ownership tests pass; applicable pre-commit hooks including mypy pass. GPU serving qualification is pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 32 +++++++++++++++- vllm/model_executor/layers/fused_moe/b12x.py | 38 +++++++++++++++++-- .../layers/fused_moe/runner/moe_runner.py | 6 +++ 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index b3672dd50d9f..56dc5cc24893 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -95,7 +95,7 @@ def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch layer, shared = _shared_fp8_tuning_fixture() other_layer, other_shared = _shared_fp8_tuning_fixture() found, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) - assert found is shared + assert found.shared is shared and found.gate is None assert key == b12x._shared_expert_tuning_context(other_layer, "w4a8_mx", 8)[1] other_shared._layer.gate_up_proj.weight = torch.empty( (32, 8), dtype=torch.float8_e4m3fn @@ -104,6 +104,28 @@ def test_shared_fp8_tuning_context_keys_geometry_not_module_identity(monkeypatch assert list(layer.modules()) == [layer] +def test_shared_fp8_tuning_gate_identity_includes_loaded_geometry(monkeypatch): + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + without_gate = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + gate = torch.nn.Module() + gate.weight = torch.empty((4, 8), dtype=torch.bfloat16) + gate.out_dtype = torch.float32 + gate.allow_ll_bf16_gemm = True + layer.routing_gate_for_preparation = weakref.ref(gate) + context, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert context.shared is shared and context.gate is gate + assert key != without_gate + gate.weight = torch.empty((8, 8), dtype=torch.bfloat16) + assert key != b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + gate.allow_ll_bf16_gemm = False + context, key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8) + assert context.gate is None and key == without_gate + assert list(layer.modules()) == [layer] + + @pytest.mark.parametrize( "exclusion", ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], @@ -183,7 +205,12 @@ def run(): SharedExpertsOrder.MULTI_STREAM_OVERLAPPED ), ) - call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, shared) + + def gate(hidden): + events.append(("gate", tuple(hidden.shape))) + + context = b12x._SharedExpertTuning(shared, gate) + call = b12x._PreparedMoECall(state, 2, 2, None, torch.bfloat16, context) hidden = torch.zeros((2, 8), dtype=torch.bfloat16) ids = torch.zeros((2, 2), dtype=torch.int32) weights = torch.ones((2, 2)) @@ -204,6 +231,7 @@ def run(): ("enter", "auxiliary"), ("shared", (2, 8)), ("exit", "auxiliary"), + ("gate", (2, 8)), "routed", ("primary", "wait", "auxiliary"), ("record", "primary"), diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 4758575f6b4f..d8956c42865b 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -67,6 +67,12 @@ def _b12x_activation_name(activation: MoEActivation) -> str: return activation.value +@dataclass(frozen=True) +class _SharedExpertTuning: + shared: Any + gate: Any | None = None + + @dataclass(frozen=True) class _PreparedMoECall: """Priming tensors for one exact prepared MoE variant. @@ -131,7 +137,8 @@ def restore() -> None: SharedExpertsOrder, ) - shared = self.shared_experts + shared = self.shared_experts.shared + gate = self.shared_experts.gate if shared._determine_shared_experts_order(hidden) == ( SharedExpertsOrder.MULTI_STREAM_OVERLAPPED ): @@ -142,6 +149,8 @@ def run() -> None: auxiliary.wait_stream(primary) with torch.cuda.stream(auxiliary): shared_output = shared._layer(hidden) + if gate is not None: + gate(hidden) binding.run() primary.wait_stream(auxiliary) shared_output.record_stream(primary) @@ -259,18 +268,41 @@ def _shared_expert_tuning_context( projections.append((tuple(weight.shape), tuple(scale.shape), str(scale.dtype))) if not (projections[0][0][1] == projections[1][0][0] == hidden_size): return None, None + gate_reference = getattr(layer, "routing_gate_for_preparation", None) + gate = None if gate_reference is None else gate_reference() + gate_weight = getattr(gate, "weight", None) + if not ( + isinstance(gate_weight, torch.Tensor) + and gate_weight.dtype == torch.bfloat16 + and gate_weight.ndim == 2 + and gate_weight.shape[1] == hidden_size + and getattr(gate, "out_dtype", None) == torch.float32 + and getattr(gate, "allow_ll_bf16_gemm", False) + ): + gate = None + gate_context = ( + None + if gate is None + else { + "backend": "ll_bf16_router", + "shape": tuple(gate.weight.shape), + "dtype": str(gate.weight.dtype), + "output_dtype": str(gate.out_dtype), + } + ) descriptor = FrozenMapping( { - "version": 1, + "version": 2, "backend": "deep_gemm_block_fp8_mlp", "projections": projections, "activation": type(mlp.act_fn).__qualname__, "ue8m0": is_deep_gemm_e8m0_used(), "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), + "gate": gate_context, } ) - return shared, descriptor + return _SharedExpertTuning(shared, gate), descriptor def _is_current_stream_capturing() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 836c7ddb1195..4467872fee66 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -306,6 +306,12 @@ def __init__( "shared_experts_for_preparation", weakref.ref(self._shared_experts), ) + if gate is not None: + object.__setattr__( + routed_experts, + "routing_gate_for_preparation", + weakref.ref(gate), + ) # Needed for string -> MoERunner layer lookup in custom ops. self.layer_name = layer_name From 20185db718fce41a2e7536e5b8918c1dae950759 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 23:13:00 +0000 Subject: [PATCH 120/365] test(moe): verify all preparation producers forward the B12X corpus Check candidate input sharing and every supplied route pattern against the B12X-owned corpus instead of fixing its realization count in the adapter test. B12X tests retain independent coverage of per-token uniqueness and sharing levels. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index aa984d93b528..ef39d30cd2ae 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -1050,6 +1050,7 @@ def test_b12x_moe_candidate_calls_share_bounded_trial_storage( """Repeated candidate calls reuse one activation/output tensor set (a weakref cache), while each call's scratch is a fresh, correctly shaped trial-only allocation rather than a caller-owned workspace region.""" + from b12x.moe.fused_moe.workloads import make_tuning_routes class FakeState: def __init__(self): @@ -1097,23 +1098,23 @@ def bind(self, **kwargs): assert not second_call.capture_safe first_call.restore() ids = first_state.bound["topk_ids"] - shared = 2 <= tokens <= 8 - unique = max(topk, (3 * tokens * topk + 2) // 5) if shared else tokens * topk - assert ids.unique().numel() == min(unique, num_experts) + expected_routes = make_tuning_routes(tokens, topk, num_experts, device="cpu") + torch.testing.assert_close(ids, expected_routes[0]) assert all(row.unique().numel() == topk for row in ids) assert torch.isfinite(first_state.bound["a"]).all() torch.testing.assert_close( first_state.bound["topk_weights"].sum(dim=1), torch.ones(tokens) ) - assert len(first_call.benchmark_producers) == (4 if shared else 1) - for first_producer, second_producer in zip( + assert len(first_call.benchmark_producers) == len(expected_routes) + for first_producer, second_producer, expected_ids in zip( first_call.benchmark_producers, second_call.benchmark_producers, + expected_routes, strict=True, ): first_producer() first_ids = ids.clone() - assert ids.unique().numel() == min(unique, num_experts) + torch.testing.assert_close(ids, expected_ids) assert all(row.unique().numel() == topk for row in ids) second_producer() torch.testing.assert_close(ids, first_ids) From 76ea484b2a56076d0d5435d877236a14e17e7304 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 23:13:00 +0000 Subject: [PATCH 121/365] test(moe): verify all preparation producers forward the B12X corpus Check candidate input sharing and every supplied route pattern against the B12X-owned corpus instead of fixing its realization count in the adapter test. B12X tests retain independent coverage of per-token uniqueness and sharing levels. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index 56dc5cc24893..aeda02b75ca4 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -1104,6 +1104,7 @@ def test_b12x_moe_candidate_calls_share_bounded_trial_storage( """Repeated candidate calls reuse one activation/output tensor set (a weakref cache), while each call's scratch is a fresh, correctly shaped trial-only allocation rather than a caller-owned workspace region.""" + from b12x.moe.fused_moe.workloads import make_tuning_routes class FakeState: def __init__(self): @@ -1151,23 +1152,23 @@ def bind(self, **kwargs): assert not second_call.capture_safe first_call.restore() ids = first_state.bound["topk_ids"] - shared = 2 <= tokens <= 8 - unique = max(topk, (3 * tokens * topk + 2) // 5) if shared else tokens * topk - assert ids.unique().numel() == min(unique, num_experts) + expected_routes = make_tuning_routes(tokens, topk, num_experts, device="cpu") + torch.testing.assert_close(ids, expected_routes[0]) assert all(row.unique().numel() == topk for row in ids) assert torch.isfinite(first_state.bound["a"]).all() torch.testing.assert_close( first_state.bound["topk_weights"].sum(dim=1), torch.ones(tokens) ) - assert len(first_call.benchmark_producers) == (4 if shared else 1) - for first_producer, second_producer in zip( + assert len(first_call.benchmark_producers) == len(expected_routes) + for first_producer, second_producer, expected_ids in zip( first_call.benchmark_producers, second_call.benchmark_producers, + expected_routes, strict=True, ): first_producer() first_ids = ids.clone() - assert ids.unique().numel() == min(unique, num_experts) + torch.testing.assert_close(ids, expected_ids) assert all(row.unique().numel() == topk for row in ids) second_producer() torch.testing.assert_close(ids, first_ids) From ff06e6ba6294213168d9c036e1fe43724f9e47ac Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 23:31:47 +0000 Subject: [PATCH 122/365] fix(moe): key overlap tuning by the shared activation recipe Include the activation module type and static parameters in the shared-expert measurement context. Retain isolated preparation for unsupported non-module activations. Serving math and stream ordering are unchanged. Validation: 19 CPU context and ownership tests pass with source-matched activation definitions. Whole-server qualification of the context-key refinement remains pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 31 +++++++++++++++++++- vllm/model_executor/layers/fused_moe/b12x.py | 11 +++++-- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index ef39d30cd2ae..f574f1f4de92 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -126,9 +126,35 @@ def test_shared_fp8_tuning_gate_identity_includes_loaded_geometry(monkeypatch): assert list(layer.modules()) == [layer] +@pytest.mark.parametrize("parameter", ("swiglu_limit", "alpha", "beta")) +def test_shared_fp8_tuning_context_tracks_activation_recipe(monkeypatch, parameter): + from vllm.config import DeviceConfig + from vllm.model_executor.layers.activation import SiluAndMulWithClamp + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))): + shared._layer.act_fn = SiluAndMulWithClamp(10.0, compile_native=False) + key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + setattr( + shared._layer.act_fn, parameter, getattr(shared._layer.act_fn, parameter) + 1 + ) + assert key != b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + + @pytest.mark.parametrize( "exclusion", - ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], + [ + "stream", + "internal", + "parallel", + "collective", + "dtype", + "width", + "quant", + "activation", + ], ) def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion): from vllm.utils import deep_gemm @@ -148,6 +174,9 @@ def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion shared._layer.down_proj.weight = torch.empty((8, 8)) elif exclusion == "width": hidden_size = 16 + elif exclusion == "activation": + del shared._layer.act_fn + shared._layer.act_fn = lambda value: value else: quant_mode = "nvfp4" assert b12x._shared_expert_tuning_context(layer, quant_mode, hidden_size) == ( diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index ad0fd43e3768..1b690cd17c77 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -253,7 +253,7 @@ def _shared_expert_tuning_context( from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used mlp = shared._layer - if not callable(getattr(mlp, "act_fn", None)): + if not isinstance(getattr(mlp, "act_fn", None), torch.nn.Module): return None, None projections = [] for name in ("gate_up_proj", "down_proj"): @@ -292,10 +292,15 @@ def _shared_expert_tuning_context( ) descriptor = FrozenMapping( { - "version": 2, + "version": 3, "backend": "deep_gemm_block_fp8_mlp", "projections": projections, - "activation": type(mlp.act_fn).__qualname__, + "activation": { + "type": ( + f"{type(mlp.act_fn).__module__}.{type(mlp.act_fn).__qualname__}" + ), + "parameters": mlp.act_fn.extra_repr(), + }, "ue8m0": is_deep_gemm_e8m0_used(), "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), From cfdecfaa4739ee3141c9745997cd8a83f31745a4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 19 Sep 2026 23:31:47 +0000 Subject: [PATCH 123/365] fix(moe): key overlap tuning by the shared activation recipe Include the activation module type and static parameters in the shared-expert measurement context. Retain isolated preparation for unsupported non-module activations. Serving math and stream ordering are unchanged. Validation: 19 CPU context and ownership tests pass with source-matched activation definitions. Whole-server qualification of the context-key refinement remains pending. Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/kernels/moe/test_b12x.py | 31 +++++++++++++++++++- vllm/model_executor/layers/fused_moe/b12x.py | 11 +++++-- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index aeda02b75ca4..80c85bf65af5 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -126,9 +126,35 @@ def test_shared_fp8_tuning_gate_identity_includes_loaded_geometry(monkeypatch): assert list(layer.modules()) == [layer] +@pytest.mark.parametrize("parameter", ("swiglu_limit", "alpha", "beta")) +def test_shared_fp8_tuning_context_tracks_activation_recipe(monkeypatch, parameter): + from vllm.config import DeviceConfig + from vllm.model_executor.layers.activation import SiluAndMulWithClamp + from vllm.utils import deep_gemm + + monkeypatch.setattr(deep_gemm, "is_deep_gemm_e8m0_used", lambda: True) + layer, shared = _shared_fp8_tuning_fixture() + with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))): + shared._layer.act_fn = SiluAndMulWithClamp(10.0, compile_native=False) + key = b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + setattr( + shared._layer.act_fn, parameter, getattr(shared._layer.act_fn, parameter) + 1 + ) + assert key != b12x._shared_expert_tuning_context(layer, "w4a8_mx", 8)[1] + + @pytest.mark.parametrize( "exclusion", - ["stream", "internal", "parallel", "collective", "dtype", "width", "quant"], + [ + "stream", + "internal", + "parallel", + "collective", + "dtype", + "width", + "quant", + "activation", + ], ) def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion): from vllm.utils import deep_gemm @@ -148,6 +174,9 @@ def test_shared_fp8_tuning_excludes_unsupported_execution(monkeypatch, exclusion shared._layer.down_proj.weight = torch.empty((8, 8)) elif exclusion == "width": hidden_size = 16 + elif exclusion == "activation": + del shared._layer.act_fn + shared._layer.act_fn = lambda value: value else: quant_mode = "nvfp4" assert b12x._shared_expert_tuning_context(layer, quant_mode, hidden_size) == ( diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index d8956c42865b..f43df5289406 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -253,7 +253,7 @@ def _shared_expert_tuning_context( from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used mlp = shared._layer - if not callable(getattr(mlp, "act_fn", None)): + if not isinstance(getattr(mlp, "act_fn", None), torch.nn.Module): return None, None projections = [] for name in ("gate_up_proj", "down_proj"): @@ -292,10 +292,15 @@ def _shared_expert_tuning_context( ) descriptor = FrozenMapping( { - "version": 2, + "version": 3, "backend": "deep_gemm_block_fp8_mlp", "projections": projections, - "activation": type(mlp.act_fn).__qualname__, + "activation": { + "type": ( + f"{type(mlp.act_fn).__module__}.{type(mlp.act_fn).__qualname__}" + ), + "parameters": mlp.act_fn.extra_repr(), + }, "ue8m0": is_deep_gemm_e8m0_used(), "tma_aligned_scales": envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES, "max_tokens": min(8, envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD), From 7d8bda56179dcc18d5a051307137638ed8f547b7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 20 Sep 2026 00:35:25 +0000 Subject: [PATCH 124/365] docs(moe): describe shared-expert tuning context contract Document supported inputs and the absent-context result without changing runtime behavior. Co-authored-by: Codex Signed-off-by: Martin Vit --- vllm/model_executor/layers/fused_moe/b12x.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 1b690cd17c77..9765dfba4802 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -236,7 +236,17 @@ def factory(state: Any): def _shared_expert_tuning_context( layer: torch.nn.Module, quant_mode: str, hidden_size: int ): - """Describe the independently executable block-FP8 MLP used during tuning.""" + """Describe the independently executable block-FP8 MLP used during tuning. + + Args: + layer: Routed MoE layer holding shared-expert and router references. + quant_mode: B12X quantization scheme of the routed experts. + hidden_size: Local input and output width shared by both expert paths. + + Returns: + The shared-expert tuning binding and immutable cache descriptor, or + ``(None, None)`` when the layer cannot provide this overlap context. + """ if quant_mode != "w4a8_mx": return None, None reference = getattr(layer, "shared_experts_for_preparation", None) From d730eb3f04bbd7055817db340c31aeb07f3592fa Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 20 Sep 2026 16:06:30 -0400 Subject: [PATCH 125/365] feat(qwen): support QSA decode context parallelism Assisted-by: OpenAI Codex Signed-off-by: derek --- tests/models/qwen4_exp/test_b12x_qsa.py | 330 +++++++++++++++ vllm/config/model.py | 12 +- vllm/models/qwen3_8_flash_next/config.py | 1 + vllm/models/qwen4_exp/nvidia/b12x_qsa.py | 512 +++++++++++++++++++++-- vllm/platforms/interface.py | 4 +- vllm/v1/attention/backend.py | 9 + 6 files changed, 819 insertions(+), 49 deletions(-) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 48540e675c36..09b15e031c94 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/tests/models/qwen4_exp/test_b12x_qsa.py @@ -25,6 +25,7 @@ ) from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState from vllm.platforms import current_platform +from vllm.platforms.interface import Platform from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.kv_cache_interface import FullAttentionSpec from vllm.v1.worker.utils import select_common_block_size @@ -90,6 +91,76 @@ def cache_requirements(**kwargs): ] +@pytest.mark.parametrize( + ( + "dcp_size", + "initial_block_size", + "expected_block_size", + "expected_page_size", + ), + [ + (1, 64, 1472, 847_872), + (2, 64, 1472, 847_872), + (4, 64, 768, 835_584), + (4, 16, 752, 818_176), + ], +) +def test_qsa_hybrid_alignment_uses_materialized_dcp_heads( + monkeypatch, + dcp_size: int, + initial_block_size: int, + expected_block_size: int, + expected_page_size: int, +) -> None: + monkeypatch.setattr( + qsa_cache_module, + "get_b12x_qsa", + lambda: SimpleNamespace( + cache_requirements=lambda **kwargs: SimpleNamespace( + compressed_page_nbytes=256 + ) + ), + ) + model_cls = SimpleNamespace( + get_mamba_state_shape_from_config=lambda config: ((817_152,),), + get_mamba_state_dtype_from_config=lambda config: (torch.uint8,), + ) + monkeypatch.setattr( + "vllm.model_executor.models.ModelRegistry.resolve_model_cls", + lambda *args, **kwargs: (model_cls, None), + ) + model_config = SimpleNamespace( + architecture="Qwen4ExpForConditionalGeneration", + dtype=torch.bfloat16, + use_mla=False, + get_head_size=lambda: 256, + get_num_kv_heads=lambda parallel_config: 1, + get_total_num_kv_heads=lambda: 2, + ) + cache_config = SimpleNamespace( + block_size=initial_block_size, + cache_dtype="fp8", + mamba_block_size=None, + user_specified_mamba_block_size=False, + mamba_cache_mode="align", + mamba_page_size_padded=None, + ) + config = SimpleNamespace( + model_config=model_config, + cache_config=cache_config, + parallel_config=SimpleNamespace( + tensor_parallel_size=4, + decode_context_parallel_size=dcp_size, + ), + ) + + Platform._align_hybrid_block_size(config, Qwen4ExpQSABackend) + + assert cache_config.block_size == expected_block_size + assert cache_config.mamba_block_size == expected_block_size + assert cache_config.mamba_page_size_padded == expected_page_size + + def test_qsa_backend_selects_the_manager_block_without_dense_page_limits() -> None: assert select_common_block_size(384, [Qwen4ExpQSABackend]) == 384 assert select_common_block_size(512, [Qwen4ExpQSABackend]) == 512 @@ -99,6 +170,252 @@ def test_qsa_backend_selects_the_manager_block_without_dense_page_limits() -> No assert Qwen4ExpQSABackend.get_preferred_block_size(70) == 72 +def test_qsa_dcp_packs_qkv_into_one_collective() -> None: + layer = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + layer.dcp_size = 4 + layer.dcp_num_heads = 24 + layer.dcp_num_kv_heads = 2 + layer.dcp_kv_replicas = 2 + layer.num_heads = 6 + layer.total_num_heads = 24 + layer.num_projected_kv_heads = 1 + layer.total_num_kv_heads = 2 + layer.head_dim = 3 + + rank_packs = [] + for rank in range(4): + rank_query = torch.full((2, 6, 3), rank, dtype=torch.float32) + kv_head = rank // 2 + rank_key = torch.full((2, 1, 3), 10 + kv_head, dtype=torch.float32) + rank_value = torch.full((2, 1, 3), 20 + kv_head, dtype=torch.float32) + rank_packs.append(torch.cat((rank_query, rank_key, rank_value), dim=1)) + + calls = [] + + def gather(packed): + calls.append(packed) + return torch.cat(rank_packs, dim=1) + + layer._gather_dcp_heads = gather + local_query = rank_packs[0][:, :6] + local_key = rank_packs[0][:, 6:7] + local_value = rank_packs[0][:, 7:8] + + query, key, value = layer._gather_dcp_qkv(local_query, local_key, local_value) + + assert len(calls) == 1 + torch.testing.assert_close(calls[0], rank_packs[0]) + assert query.shape == (2, 24, 3) + assert key.shape == value.shape == (2, 2, 3) + for rank in range(4): + assert torch.all(query[:, rank * 6 : (rank + 1) * 6] == rank) + assert torch.all(key[:, 0] == 10) + assert torch.all(key[:, 1] == 11) + assert torch.all(value[:, 0] == 20) + assert torch.all(value[:, 1] == 21) + + +def test_qsa_dcp_subgroup_deduplicates_replicated_kv_heads() -> None: + layer = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + layer.dcp_size = 2 + layer.dcp_num_heads = 12 + layer.dcp_num_kv_heads = 1 + layer.dcp_kv_replicas = 2 + layer.num_heads = 6 + layer.total_num_heads = 24 + layer.num_projected_kv_heads = 1 + layer.total_num_kv_heads = 2 + layer.head_dim = 3 + + rank_packs = [] + for rank in range(2): + query = torch.full((2, 6, 3), rank, dtype=torch.float32) + key = torch.full((2, 1, 3), 10, dtype=torch.float32) + value = torch.full((2, 1, 3), 20, dtype=torch.float32) + rank_packs.append(torch.cat((query, key, value), dim=1)) + + layer._gather_dcp_heads = lambda _packed: torch.cat(rank_packs, dim=1) + query, key, value = layer._gather_dcp_qkv( + rank_packs[0][:, :6], + rank_packs[0][:, 6:7], + rank_packs[0][:, 7:8], + ) + + assert query.shape == (2, 12, 3) + assert key.shape == value.shape == (2, 1, 3) + assert torch.all(key == 10) + assert torch.all(value == 20) + + +def test_qsa_dcp_fp32_combine_weights_before_reduce_scatter(monkeypatch) -> None: + layer = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + layer.dcp_size = 2 + layer.max_decode_rows = 8 + + remote_lse = torch.tensor( + [[1.0, 0.0, -1.0, 2.0], [0.5, -0.5, 1.5, -1.5]], + dtype=torch.float32, + ) + + class Group: + def all_gather(self, tensor, dim): + assert dim == 0 + return torch.cat((tensor, remote_lse), dim=0) + + def reduce_scatter(self, tensor, dim): + assert dim == 0 + assert tensor.dtype == torch.float32 + return tensor[: tensor.shape[0] // 2] + + monkeypatch.setattr(qsa_module, "get_dcp_group", lambda: Group()) + output = torch.arange(24, dtype=torch.bfloat16).view(2, 4, 3) + local_lse = torch.tensor( + [[0.0, 1.0, 2.0, -1.0], [1.5, 0.5, -0.5, 1.0]], + dtype=torch.float32, + ) + staged = SimpleNamespace( + sequence_lengths=torch.ones(2, dtype=torch.int32), + query_start_loc=torch.arange(3, dtype=torch.int32), + ) + + actual = layer._combine_dcp_output(output, local_lse, staged) + + global_lse = torch.logsumexp(torch.stack((local_lse, remote_lse)), dim=0) + expected = ( + output[:, :2].float() + * torch.exp(local_lse[:, :2] - global_lse[:, :2]).unsqueeze(-1) + ).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_qsa_dcp_falls_back_to_manager_above_decode_row_limit() -> None: + layer = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + expected = torch.empty(9, 6, 256) + calls = [] + layer.max_decode_rows = 8 + layer.dcp_manager = SimpleNamespace( + combine=lambda *args, **kwargs: calls.append((args, kwargs)) or expected + ) + output = torch.empty(9, 24, 256) + lse = torch.empty(9, 24) + staged = SimpleNamespace( + sequence_lengths=torch.ones(9, dtype=torch.int32), + query_start_loc=torch.arange(10, dtype=torch.int32), + ) + + actual = layer._combine_dcp_output(output, lse, staged) + + assert actual is expected + assert calls == [ + ( + (output, lse), + { + "seq_lens": staged.sequence_lengths, + "query_start_loc": staged.query_start_loc, + }, + ) + ] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_qsa_dcp_fused_pack_unpack_kernels_preserve_heads() -> None: + rows = 2 + query_heads = 4 + kv_heads = 1 + total_query_heads = 16 + total_kv_heads = 2 + head_dim = 256 + local_heads = query_heads + 2 * kv_heads + output_heads = total_query_heads + 2 * total_kv_heads + device = torch.device("cuda") + + rank_packs = [] + for rank in range(4): + query = torch.full( + (rows, query_heads, head_dim), rank, dtype=torch.bfloat16, device=device + ) + key = ( + torch.arange(head_dim, dtype=torch.bfloat16, device=device).expand( + rows, 1, -1 + ) + + (rank // 2) * 1000 + ) + value = key + 100 + rank_packs.append( + torch.cat((query, key, value), dim=1).transpose(0, 1).contiguous() + ) + + local = rank_packs[0].transpose(0, 1).contiguous() + query = local[:, :query_heads].flatten(1) + key = local[:, query_heads : query_heads + kv_heads].flatten(1) + value = local[:, query_heads + kv_heads :].flatten(1) + packed = torch.empty( + local_heads, rows, head_dim, dtype=torch.bfloat16, device=device + ) + qsa_module._pack_dcp_qkv_kernel[(local_heads, rows)]( + query, + key, + value, + packed, + query.stride(0), + key.stride(0), + value.stride(0), + rows, + NUM_QUERY_HEADS=query_heads, + NUM_KV_HEADS=kv_heads, + HEAD_DIM=head_dim, + BLOCK_DIM=head_dim, + num_warps=4, + ) + torch.testing.assert_close(packed, rank_packs[0], rtol=0, atol=0) + + gathered = torch.cat(rank_packs, dim=0).contiguous() + output = torch.empty( + rows * output_heads * head_dim, dtype=torch.bfloat16, device=device + ) + qsa_module._unpack_dcp_qkv_kernel[(output_heads, rows)]( + gathered, + output, + rows, + NUM_QUERY_HEADS=query_heads, + NUM_PROJECTED_KV_HEADS=kv_heads, + TOTAL_QUERY_HEADS=total_query_heads, + TOTAL_KV_HEADS=total_kv_heads, + KV_REPLICAS=2, + HEAD_DIM=head_dim, + BLOCK_DIM=head_dim, + num_warps=4, + ) + query_elements = rows * total_query_heads * head_dim + kv_elements = rows * total_kv_heads * head_dim + actual_query = output[:query_elements].view(rows, total_query_heads, head_dim) + actual_key = output[query_elements : query_elements + kv_elements].view( + rows, total_kv_heads, head_dim + ) + actual_value = output[query_elements + kv_elements :].view_as(actual_key) + expected_query = torch.cat( + [rank_pack[:query_heads].transpose(0, 1) for rank_pack in rank_packs], + dim=1, + ) + expected_key = torch.stack( + [rank_packs[0][query_heads], rank_packs[2][query_heads]], dim=1 + ) + expected_value = torch.stack( + [ + rank_packs[0][query_heads + kv_heads], + rank_packs[2][query_heads + kv_heads], + ], + dim=1, + ) + torch.testing.assert_close(actual_query, expected_query, rtol=0, atol=0) + torch.testing.assert_close(actual_key, expected_key, rtol=0, atol=0) + torch.testing.assert_close(actual_value, expected_value, rtol=0, atol=0) + + def test_qsa_selector_tail_is_zero_copy_in_block_outer_layer_pages() -> None: num_pages = 3 num_layers = 2 @@ -513,6 +830,17 @@ def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None: ) +def test_qsa_compressed_dcp_interleave_keeps_dcp1_local() -> None: + assert qsa_module._compressed_dcp_interleave(1, 1, 4) == 1 + assert qsa_module._compressed_dcp_interleave(4, 4, 4) == 1 + + +def test_qsa_dcp_group_kv_geometry_tracks_replica_subgroups() -> None: + assert qsa_module._qsa_dcp_group_kv_geometry(2, 2, 2) == (2, 1) + assert qsa_module._qsa_dcp_group_kv_geometry(4, 2, 2) == (1, 2) + assert qsa_module._qsa_dcp_group_kv_geometry(4, 4, 2) == (2, 2) + + def test_qsa_run_consumes_projection_views_and_writes_live_output(monkeypatch) -> None: rows = 2 query = torch.randn(rows, 6, 256, dtype=torch.bfloat16) @@ -546,6 +874,7 @@ def run(binding, **inputs): binding.output.fill_(7) owner = SimpleNamespace( + dcp_manager=None, _qsa_binding_for_workload=lambda **_: context, _prepare_qsa_metadata=lambda *_: staged, _shared_qsa_rope_positions=lambda *_: positions, @@ -1535,6 +1864,7 @@ def run(binding, *, query, request_ids, query_positions, reuse): layer = SimpleNamespace( layer_name="qsa", max_seqs=2, + dcp_manager=None, _mtp_source_rows=source_rows, _qsa_binding_for_workload=lambda **_: context, _stage_runtime_metadata=lambda *_args, **_kwargs: staged, diff --git a/vllm/config/model.py b/vllm/config/model.py index 33ecae91f3fb..707404a79805 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1461,7 +1461,15 @@ def verify_with_parallel_config( decode_context_parallel_size = parallel_config.decode_context_parallel_size if decode_context_parallel_size > 1 and not self.use_mla: total_num_kv_heads = self.get_total_num_kv_heads() - if tensor_parallel_size <= total_num_kv_heads: + supports_full_tp_dcp = bool( + getattr( + self.hf_text_config, + "supports_full_tp_dcp_with_kv_gather", + False, + ) + and decode_context_parallel_size == tensor_parallel_size + ) + if tensor_parallel_size <= total_num_kv_heads and not supports_full_tp_dcp: raise ValueError( "Decode context parallelism for GQA/MQA requires " f"`--tensor-parallel-size` ({tensor_parallel_size}) to be " @@ -1471,7 +1479,7 @@ def verify_with_parallel_config( ) max_dcp_size = tensor_parallel_size // total_num_kv_heads - if decode_context_parallel_size > max_dcp_size: + if decode_context_parallel_size > max_dcp_size and not supports_full_tp_dcp: raise ValueError( "`--decode-context-parallel-size` " f"({decode_context_parallel_size}) exceeds the maximum " diff --git a/vllm/models/qwen3_8_flash_next/config.py b/vllm/models/qwen3_8_flash_next/config.py index ff22c33b234c..d00d4c59463c 100644 --- a/vllm/models/qwen3_8_flash_next/config.py +++ b/vllm/models/qwen3_8_flash_next/config.py @@ -22,6 +22,7 @@ def __init__(self, **kwargs: Any) -> None: class Qwen3_8FlashNextTextConfig(Qwen4ExpTextConfig): model_type = "qwen3_8_flash_next_text" + supports_full_tp_dcp_with_kv_gather = True def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) diff --git a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py index 81114b664603..2e0fd84831b2 100644 --- a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py @@ -15,7 +15,7 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.config.compilation import CUDAGraphMode -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import get_dcp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention.attention import set_default_quant_scales from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -61,6 +61,8 @@ B12xPagedMetadata, B12xPagedMetadataBuilder, ) +from vllm.v1.attention.backends.mla.b12x_indexer import _merge_dcp_topk +from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.kv_cache_interface import ( AttentionSpec, FullAttentionSpec, @@ -109,6 +111,33 @@ def _qsa_prefill_context_capacities( return tuple(capacities) +def _dcp_local_length(length: int, size: int, rank: int, interleave: int) -> int: + round_width = size * interleave + full_rounds, remainder = divmod(length, round_width) + return full_rounds * interleave + min( + max(remainder - rank * interleave, 0), interleave + ) + + +def _compressed_dcp_interleave(size: int, interleave: int, ratio: int) -> int: + if size == 1: + return 1 + return interleave // ratio + + +def _qsa_dcp_group_kv_geometry( + tp_size: int, + dcp_size: int, + total_kv_heads: int, +) -> tuple[int, int]: + projected_kv_heads = max(1, total_kv_heads // tp_size) + replicas_per_kv_head = max(1, tp_size // total_kv_heads) + if max(dcp_size, replicas_per_kv_head) % min(dcp_size, replicas_per_kv_head): + raise ValueError("QSA DCP groups must align with replicated KV heads") + replicas_in_group = min(dcp_size, replicas_per_kv_head) + return dcp_size * projected_kv_heads // replicas_in_group, replicas_in_group + + def _register_qsa_compilation_context( compilation_config: Any, layer_name: str, @@ -314,6 +343,26 @@ def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: ), ) + @classmethod + def customize_hybrid_kv_cache_spec( + cls, + spec: AttentionSpec, + vllm_config: VllmConfig, + ) -> AttentionSpec: + parallel = vllm_config.parallel_config + dcp_size = parallel.decode_context_parallel_size + if dcp_size > 1: + num_kv_heads, _ = _qsa_dcp_group_kv_geometry( + parallel.tensor_parallel_size, + dcp_size, + vllm_config.model_config.get_total_num_kv_heads(), + ) + spec = replace( + spec, + num_kv_heads=num_kv_heads, + ) + return cls.customize_spec(spec) + @classmethod def get_impl_cls(cls) -> type[Qwen4ExpQSAImpl]: return Qwen4ExpQSAImpl @@ -397,7 +446,9 @@ class Qwen4ExpQSAImpl(AttentionImpl[Qwen4ExpQSAMetadata]): is_sparse: ClassVar[bool] = True supports_dense_mha_prefill: ClassVar[bool] = False - supports_dcp: bool = False + can_return_lse_for_decode: bool = True + supports_dcp: bool = True + supports_mtp_with_cp_non_trivial_interleave_size: bool = True def __init__( self, @@ -429,8 +480,6 @@ def __init__( raise ValueError("QSA requires head_dim=256 and valid grouped-query heads") if not math.isclose(scale, head_size**-0.5, rel_tol=1e-5, abs_tol=1e-7): raise ValueError("QSA requires canonical head_dim**-0.5 scaling") - if self.total_cp_world_size > 1: - raise NotImplementedError("QSA does not support context parallelism") self.num_heads = int(num_heads) self.head_size = int(head_size) self.output_head_size = self.head_size @@ -676,10 +725,113 @@ def _stage_qsa_rope_positions_kernel( ) +@triton.jit +def _pack_dcp_qkv_kernel( + query, + key, + value, + packed, + query_row_stride, + key_row_stride, + value_row_stride, + rows, + NUM_QUERY_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_DIM: tl.constexpr, +) -> None: + head = tl.program_id(0) + row = tl.program_id(1) + dims = tl.arange(0, BLOCK_DIM) + dim_mask = dims < HEAD_DIM + is_query = head < NUM_QUERY_HEADS + is_key = (head >= NUM_QUERY_HEADS) & (head < NUM_QUERY_HEADS + NUM_KV_HEADS) + kv_head = head - NUM_QUERY_HEADS + value_head = kv_head - NUM_KV_HEADS + data = tl.load( + query + row * query_row_stride + head * HEAD_DIM + dims, + mask=is_query & dim_mask, + other=0.0, + ) + data += tl.load( + key + row * key_row_stride + kv_head * HEAD_DIM + dims, + mask=is_key & dim_mask, + other=0.0, + ) + data += tl.load( + value + row * value_row_stride + value_head * HEAD_DIM + dims, + mask=~(is_query | is_key) & dim_mask, + other=0.0, + ) + tl.store( + packed + (head * rows + row) * HEAD_DIM + dims, + data, + mask=dim_mask, + ) + + +@triton.jit +def _unpack_dcp_qkv_kernel( + gathered, + output, + rows, + NUM_QUERY_HEADS: tl.constexpr, + NUM_PROJECTED_KV_HEADS: tl.constexpr, + TOTAL_QUERY_HEADS: tl.constexpr, + TOTAL_KV_HEADS: tl.constexpr, + KV_REPLICAS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_DIM: tl.constexpr, +) -> None: + output_head = tl.program_id(0) + row = tl.program_id(1) + dims = tl.arange(0, BLOCK_DIM) + dim_mask = dims < HEAD_DIM + local_heads = NUM_QUERY_HEADS + 2 * NUM_PROJECTED_KV_HEADS + is_query = output_head < TOTAL_QUERY_HEADS + is_key = (output_head >= TOTAL_QUERY_HEADS) & ( + output_head < TOTAL_QUERY_HEADS + TOTAL_KV_HEADS + ) + + query_rank = output_head // NUM_QUERY_HEADS + query_head = output_head % NUM_QUERY_HEADS + kv_head = output_head - TOTAL_QUERY_HEADS + value_head = kv_head - TOTAL_KV_HEADS + selected_kv_head = tl.where(is_key, kv_head, value_head) * KV_REPLICAS + kv_rank = selected_kv_head // NUM_PROJECTED_KV_HEADS + local_kv_head = selected_kv_head % NUM_PROJECTED_KV_HEADS + source_head = tl.where( + is_query, + query_rank * local_heads + query_head, + kv_rank * local_heads + + NUM_QUERY_HEADS + + local_kv_head + + tl.where(is_key, 0, NUM_PROJECTED_KV_HEADS), + ) + data = tl.load( + gathered + (source_head * rows + row) * HEAD_DIM + dims, + mask=dim_mask, + ) + + query_elements = rows * TOTAL_QUERY_HEADS * HEAD_DIM + key_elements = rows * TOTAL_KV_HEADS * HEAD_DIM + query_offset = (row * TOTAL_QUERY_HEADS + output_head) * HEAD_DIM + key_offset = query_elements + (row * TOTAL_KV_HEADS + kv_head) * HEAD_DIM + value_offset = ( + query_elements + key_elements + (row * TOTAL_KV_HEADS + value_head) * HEAD_DIM + ) + output_offset = tl.where( + is_query, + query_offset, + tl.where(is_key, key_offset, value_offset), + ) + tl.store(output + output_offset + dims, data, mask=dim_mask) + + class Qwen4ExpQSAAttention(nn.Module, AttentionLayerBase): """Merged main attention, selector state, and b12x QSA transaction.""" - supports_dcp = False + supports_dcp = True def __init__( self, @@ -708,11 +860,10 @@ def __init__( if getattr(quant_config, "kv_cache_scheme", None) is not None: raise NotImplementedError("QSA does not support KV-cache quantization") parallel = vllm_config.parallel_config - if ( - parallel.prefill_context_parallel_size > 1 - or parallel.decode_context_parallel_size > 1 - ): - raise NotImplementedError("QSA does not support context parallelism") + if parallel.prefill_context_parallel_size > 1: + raise NotImplementedError( + "QSA does not support prefill context parallelism" + ) if not getattr(config, "is_causal", True): raise NotImplementedError("QSA requires causal decoder attention") if getattr(config, "dual_chunk_attention_config", None) is not None: @@ -733,6 +884,16 @@ def __init__( aux_stream() self.hidden_size = int(config.hidden_size) tp_size = get_tensor_model_parallel_world_size() + self.dcp_size = int(parallel.decode_context_parallel_size) + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0 + self.cp_kv_cache_interleave_size = int(parallel.cp_kv_cache_interleave_size) + if self.dcp_size > 1: + if self.cp_kv_cache_interleave_size % _QSA_COMPRESS_RATIO: + raise ValueError( + "QSA DCP requires cp_kv_cache_interleave_size divisible by " + f"{_QSA_COMPRESS_RATIO}" + ) + self.overlap_input_projections = False self.total_num_heads = int(config.num_attention_heads) if self.total_num_heads % tp_size: raise ValueError("QSA query heads must divide across TP ranks") @@ -743,12 +904,21 @@ def __init__( raise ValueError("QSA KV heads must divide across TP ranks") elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must divide replicated QSA KV heads") - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.num_projected_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.dcp_num_heads = self.num_heads * self.dcp_size + self.dcp_num_kv_heads, self.dcp_kv_replicas = _qsa_dcp_group_kv_geometry( + tp_size, + self.dcp_size, + self.total_num_kv_heads, + ) + self.num_kv_heads = ( + self.dcp_num_kv_heads if self.dcp_size > 1 else self.num_projected_kv_heads + ) self.head_dim = int(config.head_dim) if self.head_dim != 256: raise NotImplementedError("The SM12x QSA integration requires head_dim=256") self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim + self.kv_size = self.num_projected_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.attn_output_gate = True @@ -841,6 +1011,22 @@ def __init__( self.max_decode_rows = self.max_seqs * (1 + self.max_speculative_tokens) device = torch.device("cuda", torch.accelerator.current_device_index()) self.device = device + self.dcp_manager = ( + MLADCPManager( + vllm_config=vllm_config, + device=device, + num_heads=self.num_heads, + query_head_dim=self.head_dim, + output_head_dim=self.head_dim, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + is_lse_base_on_e=True, + use_pcp=False, + ) + if self.dcp_size > 1 + else None + ) self._selector_done = ( torch.cuda.Event() if self.overlap_input_projections else None ) @@ -913,7 +1099,7 @@ def __init__( "_qsa_output", torch.empty( self.max_tokens, - self.num_heads, + self.dcp_num_heads if self.dcp_size > 1 else self.num_heads, self.head_dim, dtype=torch.bfloat16, device=device, @@ -1003,7 +1189,7 @@ def _qsa_caps(self, **geometry): max_batch=self.max_seqs, max_raw_state_slots=self.max_seqs, max_speculative_tokens=self.max_speculative_tokens, - q_heads=self.num_heads, + q_heads=(self.dcp_num_heads if self.dcp_size > 1 else self.num_heads), kv_heads=self.num_kv_heads, head_dim=self.head_dim, index_heads=self.index_heads, @@ -1020,6 +1206,9 @@ def _qsa_caps(self, **geometry): rms_norm_eps=float(self.indexer.q_layernorm.variance_epsilon), dtype=torch.bfloat16, kv_dtype=self.kv_cache_kernel_dtype, + dcp_size=self.dcp_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, **geometry, ) @@ -1210,10 +1399,24 @@ def prepare(state): raise PreparationResourceUnavailableError( "QSA primer requires a live cache page" ) - main_pages = triton.cdiv(rows, caps.main_page_size) - compressed_pages = triton.cdiv( - rows // caps.compress_ratio, caps.compressed_page_size + local_rows = _dcp_local_length( + rows, + caps.dcp_size, + caps.dcp_rank, + caps.cp_kv_cache_interleave_size, + ) + local_groups = _dcp_local_length( + rows // caps.compress_ratio, + caps.dcp_size, + caps.dcp_rank, + _compressed_dcp_interleave( + caps.dcp_size, + caps.cp_kv_cache_interleave_size, + caps.compress_ratio, + ), ) + main_pages = triton.cdiv(local_rows, caps.main_page_size) + compressed_pages = triton.cdiv(local_groups, caps.compressed_page_size) output = torch.empty( (rows, caps.q_heads, caps.head_dim), dtype=torch.bfloat16, device=device ) @@ -1391,11 +1594,164 @@ def _project_qkv_gate( query, gate = q_gate.chunk(2, dim=-1) query = self.q_norm(query).flatten(-2) key = self.k_norm( - key.unflatten(-1, (self.num_kv_heads, self.head_dim)) + key.unflatten(-1, (self.num_projected_kv_heads, self.head_dim)) ).flatten(-2) query, key = self.rotary_emb(positions, query, key) return query, key, value, gate.flatten(-2) + def _combine_dcp_output( + self, + partial_output: torch.Tensor, + partial_lse: torch.Tensor, + staged: _StagedQSAMetadata, + ) -> torch.Tensor: + if partial_output.shape[0] <= self.max_decode_rows: + from vllm.v1.attention.ops.dcp import mask_dcp_empty_shards_ + + mask_dcp_empty_shards_( + partial_lse, + staged.sequence_lengths, + staged.query_start_loc, + ) + group = get_dcp_group() + lses = group.all_gather(partial_lse.contiguous(), dim=0).view( + self.dcp_size, *partial_lse.shape + ) + lses = torch.where( + torch.isfinite(lses), lses, lses.new_full((), float("-inf")) + ) + global_lse = torch.logsumexp(lses, dim=0) + weights = torch.exp(partial_lse - global_lse) + weights = torch.where(torch.isfinite(weights), weights, 0.0) + corrected = partial_output.float() * weights.unsqueeze(-1) + combined = group.reduce_scatter( + corrected.transpose(0, 1).contiguous(), dim=0 + ) + return combined.transpose(0, 1).to(partial_output.dtype) + assert self.dcp_manager is not None + return self.dcp_manager.combine( + partial_output, + partial_lse, + seq_lens=staged.sequence_lengths, + query_start_loc=staged.query_start_loc, + ) + + def _gather_dcp_heads(self, tensor: torch.Tensor) -> torch.Tensor: + gathered = get_dcp_group().all_gather( + tensor.transpose(0, 1).contiguous(), dim=0 + ) + return gathered.transpose(0, 1).contiguous() + + def _gather_dcp_qkv_fused_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + query = query.flatten(1) + key = key.flatten(1) + value = value.flatten(1) + rows = int(query.shape[0]) + local_heads = self.num_heads + 2 * self.num_projected_kv_heads + packed = torch.empty( + (local_heads, rows, self.head_dim), + dtype=query.dtype, + device=query.device, + ) + block_dim = triton.next_power_of_2(self.head_dim) + _pack_dcp_qkv_kernel[(local_heads, rows)]( + query, + key, + value, + packed, + query.stride(0), + key.stride(0), + value.stride(0), + rows, + NUM_QUERY_HEADS=self.num_heads, + NUM_KV_HEADS=self.num_projected_kv_heads, + HEAD_DIM=self.head_dim, + BLOCK_DIM=block_dim, + num_warps=4, + ) + gathered = get_dcp_group().all_gather(packed, dim=0) + output_heads = self.dcp_num_heads + 2 * self.dcp_num_kv_heads + output = torch.empty( + rows * output_heads * self.head_dim, + dtype=query.dtype, + device=query.device, + ) + _unpack_dcp_qkv_kernel[(output_heads, rows)]( + gathered, + output, + rows, + NUM_QUERY_HEADS=self.num_heads, + NUM_PROJECTED_KV_HEADS=self.num_projected_kv_heads, + TOTAL_QUERY_HEADS=self.dcp_num_heads, + TOTAL_KV_HEADS=self.dcp_num_kv_heads, + KV_REPLICAS=self.dcp_kv_replicas, + HEAD_DIM=self.head_dim, + BLOCK_DIM=block_dim, + num_warps=4, + ) + query_elements = rows * self.dcp_num_heads * self.head_dim + kv_elements = rows * self.dcp_num_kv_heads * self.head_dim + query = output[:query_elements].view(rows, self.dcp_num_heads, self.head_dim) + key = output[query_elements : query_elements + kv_elements].view( + rows, self.dcp_num_kv_heads, self.head_dim + ) + value = output[query_elements + kv_elements :].view_as(key) + return query, key, value + + def _gather_dcp_qkv( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather Q/K/V and remove replicated KV heads.""" + + if query.is_cuda and query.shape[0] <= self.max_decode_rows: + return self._gather_dcp_qkv_fused_cuda(query, key, value) + if not query.is_cuda: + local_heads = self.num_heads + 2 * self.num_projected_kv_heads + packed = torch.cat((query, key, value), dim=1) + gathered = self._gather_dcp_heads(packed).view( + query.shape[0], self.dcp_size, local_heads, self.head_dim + ) + query = ( + gathered[:, :, : self.num_heads] + .contiguous() + .view(query.shape[0], self.dcp_num_heads, self.head_dim) + ) + kv_start = self.num_heads + key = ( + gathered[:, :, kv_start : kv_start + self.num_projected_kv_heads] + .contiguous() + .view(query.shape[0], -1, self.head_dim) + ) + value = ( + gathered[:, :, kv_start + self.num_projected_kv_heads :] + .contiguous() + .view(query.shape[0], -1, self.head_dim) + ) + else: + query = self._gather_dcp_heads(query) + key = self._gather_dcp_heads(key) + value = self._gather_dcp_heads(value) + if key.ndim == 2: + key = key.view( + query.shape[0], + self.dcp_size * self.num_projected_kv_heads, + self.head_dim, + ) + value = value.view_as(key) + return ( + query, + key[:, :: self.dcp_kv_replicas].contiguous(), + value[:, :: self.dcp_kv_replicas].contiguous(), + ) + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: super().bind_kv_cache(kv_cache) self.max_seq_len = int(self._qsa_model_config.max_model_len) @@ -1430,9 +1786,18 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: // self.compress_ratio * self.compress_ratio ) - table_width = math.ceil(qsa_max_seq_len / main_page_size) + dcp_round = self.dcp_size * self.cp_kv_cache_interleave_size + full_rounds, remainder = divmod(qsa_max_seq_len, dcp_round) + local_max_seq_len = full_rounds * self.cp_kv_cache_interleave_size + min( + max( + remainder - self.dcp_rank * self.cp_kv_cache_interleave_size, + 0, + ), + self.cp_kv_cache_interleave_size, + ) + table_width = math.ceil(local_max_seq_len / main_page_size) compressed_table_width = math.ceil( - (qsa_max_seq_len // self.compress_ratio) / compressed_page_size + (local_max_seq_len // self.compress_ratio) / compressed_page_size ) if compressed_table_width != table_width: raise RuntimeError( @@ -1941,6 +2306,10 @@ def _run_reused_qsa(self, query, key, value, output) -> None: compressed_block_table=context.compressed_block_table, ) self._b12x_diagnostic_request_ids = staged.request_ids + if self.dcp_manager is not None: + query, key, value = self._gather_dcp_qkv( + query[:rows], key[:rows], value[:rows] + ) impl = cast(Qwen4ExpQSAImpl, self.impl) impl.do_kv_cache_update( self, key[:rows], value[:rows], self.kv_cache, metadata.slot_mapping[:rows] @@ -1948,13 +2317,26 @@ def _run_reused_qsa(self, query, key, value, output) -> None: qsa = get_b12x_qsa() if qsa is None: raise RuntimeError("b12x QSA is unavailable for MTP selection reuse") - qsa.run( - self._bind_qsa_context(context, staged, output[:rows]), - query=query[:rows], - request_ids=staged.request_ids, - query_positions=staged.logical_positions, - reuse=qsa.DraftSelectionReuse(source_rows=self._mtp_source_rows), - ) + reuse = qsa.DraftSelectionReuse(source_rows=self._mtp_source_rows) + if self.dcp_manager is None: + qsa.run( + self._bind_qsa_context(context, staged, output[:rows]), + query=query[:rows], + request_ids=staged.request_ids, + query_positions=staged.logical_positions, + reuse=reuse, + ) + else: + binding = self._bind_qsa_context(context, staged, self._qsa_output[:rows]) + partial_output, partial_lse = qsa.attend_reuse( + binding, + query=query, + request_ids=staged.request_ids, + query_positions=staged.logical_positions, + reuse=reuse, + ) + combined = self._combine_dcp_output(partial_output, partial_lse, staged) + output[:rows].copy_(combined) if rows < output.shape[0]: output[rows:].zero_() @@ -2148,6 +2530,10 @@ def _run_b12x_qsa( rope_positions = self._shared_qsa_rope_positions( positions, staged.request_ids, rows ) + if self.dcp_manager is not None: + query, key, value = self._gather_dcp_qkv( + query[:rows], key[:rows], value[:rows] + ) impl = cast(Qwen4ExpQSAImpl, self.impl) impl.do_kv_cache_update( self, @@ -2159,20 +2545,56 @@ def _run_b12x_qsa( qsa = get_b12x_qsa() if qsa is None: raise RuntimeError("b12x QSA disappeared after cache binding") - binding = self._bind_qsa_context(context, staged, output[:rows]) - qsa.run( - binding, - query=query[:rows], - index_query=index_query[:rows], - raw_index_key=raw_index_key[:rows], - request_ids=staged.request_ids, - query_positions=staged.logical_positions, - rope_positions=rope_positions, - sequence_lengths=staged.sequence_lengths, - query_start_loc=staged.query_start_loc, - num_accepted_tokens=staged.num_accepted_tokens, - is_prefilling=staged.is_prefilling, - ) + if self.dcp_manager is None: + binding = self._bind_qsa_context(context, staged, output[:rows]) + qsa.run( + binding, + query=query[:rows], + index_query=index_query[:rows], + raw_index_key=raw_index_key[:rows], + request_ids=staged.request_ids, + query_positions=staged.logical_positions, + rope_positions=rope_positions, + sequence_lengths=staged.sequence_lengths, + query_start_loc=staged.query_start_loc, + num_accepted_tokens=staged.num_accepted_tokens, + is_prefilling=staged.is_prefilling, + ) + else: + binding = self._bind_qsa_context(context, staged, self._qsa_output[:rows]) + selection = qsa.select( + binding, + query=query, + index_query=index_query[:rows], + raw_index_key=raw_index_key[:rows], + request_ids=staged.request_ids, + query_positions=staged.logical_positions, + rope_positions=rope_positions, + sequence_lengths=staged.sequence_lengths, + query_start_loc=staged.query_start_loc, + num_accepted_tokens=staged.num_accepted_tokens, + is_prefilling=staged.is_prefilling, + ) + _merge_dcp_topk( + selection.group_ids, + selection.scores, + self.dcp_rank, + self.dcp_size, + _compressed_dcp_interleave( + self.dcp_size, + self.cp_kv_cache_interleave_size, + self.compress_ratio, + ), + ) + partial_output, partial_lse = qsa.attend( + binding, + query=query, + request_ids=staged.request_ids, + query_positions=staged.logical_positions, + selection=selection, + ) + combined = self._combine_dcp_output(partial_output, partial_lse, staged) + output[:rows].copy_(combined) if rows < output.shape[0]: output[rows:].zero_() @@ -2247,8 +2669,8 @@ def forward( query, key, value, gate = self._project_qkv_gate(qkv, positions) rows = int(hidden_states.shape[0]) query = query.view(rows, self.num_heads, self.head_dim) - key = key.view(rows, self.num_kv_heads, self.head_dim) - value = value.view(rows, self.num_kv_heads, self.head_dim) + key = key.view(rows, self.num_projected_kv_heads, self.head_dim) + value = value.view(rows, self.num_projected_kv_heads, self.head_dim) output = torch.empty_like(query) torch.ops.vllm.qwen4_exp_b12x_qsa_run_projected( positions, @@ -2267,8 +2689,8 @@ def forward( index_query, raw_index_key = self.indexer.project(hidden_states) num_tokens = int(hidden_states.shape[0]) query = query.view(num_tokens, self.num_heads, self.head_dim) - key = key.view(num_tokens, self.num_kv_heads, self.head_dim) - value = value.view(num_tokens, self.num_kv_heads, self.head_dim) + key = key.view(num_tokens, self.num_projected_kv_heads, self.head_dim) + value = value.view(num_tokens, self.num_projected_kv_heads, self.head_dim) output = torch.empty_like(query) layer_name = _encode_layer_name(self.layer_name) if current_platform.opaque_attention_op(): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 3283694e4e49..f1e70ee9a719 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -975,8 +975,8 @@ def _align_hybrid_block_size( dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ) - attn_page_size_1_token = backend_cls.customize_spec( - attn_spec + attn_page_size_1_token = backend_cls.customize_hybrid_kv_cache_spec( + attn_spec, vllm_config ).page_size_bytes # Compute mamba page size diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 98b0dcc35a3f..b43e5f507162 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -150,6 +150,15 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """ return spec + @classmethod + def customize_hybrid_kv_cache_spec( + cls, + spec: "AttentionSpec", + vllm_config: "VllmConfig", + ) -> "AttentionSpec": + """Adjust the probe used to size attention against recurrent state.""" + return cls.customize_spec(spec) + @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: supported_sizes = cls.get_supported_kernel_block_sizes() From ca3609cd9ab1eb62f8681a776669f9439fcf4368 Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 20 Sep 2026 16:07:17 -0400 Subject: [PATCH 126/365] fix(spec_decode): include graph shape in config hash Assisted-by: OpenAI Codex Signed-off-by: derek --- tests/config/test_config_utils.py | 13 +++++++++++++ vllm/config/speculative.py | 2 +- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index ff31ddbe828c..2f20bba9f695 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -8,6 +8,7 @@ from vllm.config.cache import CacheConfig from vllm.config.scheduler import SchedulerConfig +from vllm.config.speculative import SpeculativeConfig from vllm.config.utils import get_hash_factors, hash_factors, normalize_value # Helpers @@ -235,6 +236,18 @@ def test_scheduler_config_hash_includes_max_num_seqs(): assert larger_batch_hash != base_hash +def test_speculative_config_hash_includes_graph_shape(): + def config(method: str, num_speculative_tokens: int) -> SpeculativeConfig: + return SpeculativeConfig( + method=method, + num_speculative_tokens=num_speculative_tokens, + prompt_lookup_max=num_speculative_tokens, + ) + + assert config("ngram", 1).compute_hash() != config("ngram", 3).compute_hash() + assert config("ngram", 3).compute_hash() != config("ngram_gpu", 3).compute_hash() + + def test_cache_config_hash_ignores_prefix_cache_retention_interval(): base_hash = CacheConfig().compute_hash() assert CacheConfig(prefix_cache_retention_interval=64).compute_hash() == base_hash diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 42df7dbd911d..591d029d1acd 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -627,7 +627,7 @@ def compute_hash(self) -> str: excluding anything before input ids/embeddings and after the final hidden states. """ - factors: list[Any] = [] + factors: list[Any] = [self.method, self.num_speculative_tokens] # Eagle3 and extract_hidden_states affect the computation graph because # they return intermediate hidden states in addition to the final hidden state. uses_aux_hidden_states = self.method in ( From 8f740b1f9f403779309cd3adfc0c87d62b123642 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 20 Sep 2026 23:18:34 +0200 Subject: [PATCH 127/365] docs(release): require integration changelog fragments [skip ci] --- .lil/README.md | 33 +++++++++++++++++++++++++++++++++ AGENTS.md | 7 +++++++ 2 files changed, 40 insertions(+) create mode 100644 .lil/README.md diff --git a/.lil/README.md b/.lil/README.md new file mode 100644 index 000000000000..c86c03acb3fb --- /dev/null +++ b/.lil/README.md @@ -0,0 +1,33 @@ +# Integration container changelog fragments + +A runtime-affecting change that will be pushed or imported into an +`integration/*` branch must include one versioned JSON fragment in +`.lil/changes/`. The community container publisher compares these fragments +between exact component revisions and creates the container release notes. + +The complete schema, examples, correction policy, and publication flow are in +the [integration container changelog runbook](https://github.com/local-inference-lab/blackwell-llm-docker/blob/main/docs/integration-release-changelog.md). + +Use a component-prefixed identifier that matches the filename. For example, +`.lil/changes/vllm-816.json` contains: + +```json +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-816", + "category": "fix", + "summary": "Merge QSA selection across DCP ranks", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "DCP1, DCP2, and DCP4 use the same distributed selection contract." + ], + "pull_requests": [816], + "requires": ["b12x-402"] +} +``` + +The fragment belongs in the implementation PR or in the commit series that +imports the reviewed change. Preserve contributor attribution. Once a container +release has published a fragment, do not edit or delete it; add a corrective +fragment instead. diff --git a/AGENTS.md b/AGENTS.md index a1060975b0b4..ec1a0cea9c96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -156,3 +156,10 @@ vulnerability process. - **Editing these instructions**: [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) — Rules for modifying AGENTS.md or any domain-specific guide it references. + +## LIL integration container changelog + +Before pushing or importing a runtime-affecting change into an +`integration/*` branch, read [`.lil/README.md`](.lil/README.md) and add the +required `.lil/changes/.json` fragment to the same commit series. Preserve +published fragments unchanged; corrections use a new fragment. From ac0e6c47b7d8bab6678b0ac3d6e4786ed0040f7e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 20 Sep 2026 21:30:55 +0000 Subject: [PATCH 128/365] docs(release): describe container-visible behavior --- .lil/changes/vllm-815.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-815.json diff --git a/.lil/changes/vllm-815.json b/.lil/changes/vllm-815.json new file mode 100644 index 000000000000..48643bd11398 --- /dev/null +++ b/.lil/changes/vllm-815.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-815", + "category": "fix", + "summary": "Separate CUDA graph caches by speculative method and draft depth", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "Speculative configurations with different execution graph shapes no longer reuse an incompatible compiled cache entry." + ], + "pull_requests": [ + 815 + ], + "requires": [] +} From dadcac59e3795e017636403cd4fdcd93e9d60470 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 20 Sep 2026 21:30:56 +0000 Subject: [PATCH 129/365] docs(release): describe container-visible behavior --- .lil/changes/vllm-816.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-816.json diff --git a/.lil/changes/vllm-816.json b/.lil/changes/vllm-816.json new file mode 100644 index 000000000000..d53c08cb2e81 --- /dev/null +++ b/.lil/changes/vllm-816.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-816", + "category": "feature", + "summary": "Support Qwen QSA decode context parallelism", + "models": [ + "Qwen3.8-Flash-Next" + ], + "compatibility": "No user action required.", + "details": [ + "DCP ranks own interleaved KV and selector state, merge candidates globally, and combine decode attention with FP32 weighting.", + "QSA external KV-cache transfer remains unsupported." + ], + "pull_requests": [ + 816 + ], + "requires": [ + "b12x-401", + "b12x-402" + ] +} From 7c943ab634075c744425a33ab385cfe363d8b686 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 10:33:36 +0000 Subject: [PATCH 130/365] feat(glm5next): recover KDA state after speculative verification Keep one FP32 recurrent checkpoint per GLM KDA draft window and reconstruct the accepted state from correction and key/gate records. The optional --use-replayssm mode retains B12X prefill and leaves default serving unchanged. Validate cache geometry and metadata with CPU tests, recurrence and accepted-state parity with a CUDA test, and DFlash2 TP4 serving with full CUDA graphs, 32k prefill, and C1/C8 decode. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit aa9b31b86b915b1b820dc3657946abee7589ee61) --- tests/models/glm5next/test_kda_recoverssm.py | 163 ++++++++++++++++++ tests/models/glm5next/test_kda_recurrent.py | 83 +++++++++ vllm/config/vllm.py | 6 +- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 99 ++++++++--- vllm/models/glm5next/nvidia/kda.py | 122 ++++++++++++- vllm/models/glm5next/nvidia/model.py | 48 +++++- vllm/v1/attention/backends/gdn_attn.py | 17 +- 7 files changed, 499 insertions(+), 39 deletions(-) create mode 100644 tests/models/glm5next/test_kda_recoverssm.py diff --git a/tests/models/glm5next/test_kda_recoverssm.py b/tests/models/glm5next/test_kda_recoverssm.py new file mode 100644 index 000000000000..95a97c3aaaef --- /dev/null +++ b/tests/models/glm5next/test_kda_recoverssm.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, +) +from vllm.config.compilation import CUDAGraphMode +from vllm.models.glm5next.nvidia.kda import ( + Glm5NextRecoverKDAMetadataBuilder, +) +from vllm.models.glm5next.nvidia.model import Glm5NextForCausalLM +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import MambaSpec + + +@pytest.fixture(autouse=True) +def disable_pinned_host_staging(monkeypatch): + monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False) + monkeypatch.setattr("vllm.v1.attention.backends.utils.PIN_MEMORY", False) + + +def test_glm_recoverssm_reserves_one_recurrent_state_per_draft_window(): + config = SimpleNamespace( + model_config=SimpleNamespace( + dtype=torch.bfloat16, + hf_config=SimpleNamespace( + linear_num_heads=64, + linear_head_dim=128, + linear_conv_kernel_dim=4, + ), + ), + cache_config=SimpleNamespace( + mamba_cache_dtype="auto", + mamba_ssm_cache_dtype="auto", + use_kda_recoverssm=False, + mamba_cache_mode="align", + ), + parallel_config=SimpleNamespace(tensor_parallel_size=2), + speculative_config=SimpleNamespace(num_speculative_tokens=7), + ) + baseline = MambaSpec( + block_size=256, + shapes=Glm5NextForCausalLM.get_mamba_state_shape_from_config(config), + dtypes=Glm5NextForCausalLM.get_mamba_state_dtype_from_config(config), + mamba_cache_mode="align", + num_speculative_blocks=7, + ) + config.cache_config.use_kda_recoverssm = True + recovered = MambaSpec( + block_size=256, + shapes=Glm5NextForCausalLM.get_mamba_state_shape_from_config(config), + dtypes=Glm5NextForCausalLM.get_mamba_state_dtype_from_config(config), + mamba_cache_mode="align", + num_speculative_blocks=0, + ) + assert len(baseline.shapes) == 2 + assert len(recovered.shapes) == 4 + assert recovered.dtypes[1] == torch.float32 + assert recovered.dtypes[2] == torch.float32 + assert baseline.page_size_bytes == 2_342_912 + assert recovered.page_size_bytes == 2_605_056 + assert ( + baseline.max_memory_usage_bytes(config) + - recovered.max_memory_usage_bytes(config) + == 15_876_096 + ) + + +def _builder() -> Glm5NextRecoverKDAMetadataBuilder: + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(model_type="glm5_next") + ), + cache_config=SimpleNamespace( + mamba_cache_mode="none", + use_kda_recoverssm=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + speculative_config=SimpleNamespace( + num_speculative_tokens=2, + parallel_drafting=False, + use_eagle_block_drop=Mock(return_value=False), + ), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, + max_cudagraph_capture_size=None, + static_forward_context={}, + ), + scheduler_config=SimpleNamespace(max_num_seqs=4, max_num_batched_tokens=256), + additional_config={}, + num_speculative_tokens=2, + use_v2_model_runner=False, + ) + spec = MambaSpec( + block_size=16, + shapes=((16, 64),), + dtypes=(torch.float16,), + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="none", + num_speculative_blocks=0, + ) + builder = Glm5NextRecoverKDAMetadataBuilder( + spec, + ["layer.0"], + config, + torch.device("cpu"), + ) + builder._recoverssm_context = Mock() + return builder + + +def test_glm_recoverssm_mixed_prefill_and_decode_keeps_one_state_slot(): + common = create_common_attn_metadata( + BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]), + 16, + torch.device("cpu"), + ).replace(is_prefilling=torch.tensor([True, True, False])) + metadata = _builder().build( + 0, + common, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2]), + num_accepted_tokens=torch.tensor([3, 2, 2]), + ) + assert metadata.num_spec_decodes == 1 + assert metadata.num_prefills == 2 + assert metadata.spec_state_indices_tensor is not None + assert metadata.spec_state_indices_tensor.shape == (1, 1) + torch.testing.assert_close( + metadata.num_accepted_tokens, torch.ones(1, dtype=torch.int32) + ) + assert metadata.recoverssm_commit is not None + torch.testing.assert_close( + metadata.recoverssm_commit.request_indices, torch.tensor([2], dtype=torch.int32) + ) + sampled = torch.tensor([1, 1, 3], dtype=torch.int32) + assert metadata.commit_recoverssm_state(sampled) is None + metadata.recoverssm_context.commit.assert_called_once() + + +def test_glm_recoverssm_keeps_draftless_decode_out_of_prefill(): + common = create_common_attn_metadata( + BatchSpec(seq_lens=[20], query_lens=[1]), + 16, + torch.device("cpu"), + ).replace(is_prefilling=torch.tensor([False])) + metadata = _builder().build( + 0, + common, + num_decode_draft_tokens_cpu=torch.tensor([-1]), + num_accepted_tokens=torch.tensor([1]), + ) + assert metadata.num_spec_decodes == 1 + assert metadata.num_prefills == 0 + assert metadata.spec_state_indices_tensor is not None + assert metadata.spec_state_indices_tensor.shape == (1, 1) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 68cf493b443a..d76f401a5d83 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -8,10 +8,17 @@ match a pure-PyTorch recurrence, and reject layouts it cannot address. """ +from types import SimpleNamespace + import pytest import torch +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda +from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + kda_recoverssm_verify, +) from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -184,3 +191,79 @@ def test_fused_recurrent_kda_rejects_unaddressable_layouts(): broken["cu_seqlens"] = None if q.shape[0] > 1 else inputs["cu_seqlens"] with pytest.raises(AssertionError, match=r"torch.Size"): run_kernel(broken, state) + + +@torch.inference_mode() +def test_recoverssm_matches_glm_recurrence_and_accepted_state(): + torch.manual_seed(7) + num_seqs, query_len = 2, 3 + device = torch.device("cuda") + inputs, baseline_state = make_inputs(num_seqs, query_len, device) + initial_state = baseline_state.clone() + inputs["num_accepted_tokens"] = torch.ones( + num_seqs, dtype=torch.int32, device=device + ) + baseline_output = run_kernel(inputs, baseline_state) + + correction = torch.empty( + initial_state.shape[0], H, query_len, D, dtype=torch.float32, device=device + ) + key_gate = torch.empty( + initial_state.shape[0], + H, + query_len, + 2 * D, + dtype=torch.bfloat16, + device=device, + ) + recovery_indices = inputs["ssm_state_indices"][:, 0].contiguous() + recovered_output = kda_recoverssm_verify( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + raw_g=inputs["g"], + raw_beta=inputs["beta"], + A_log=inputs["a_log"], + dt_bias=inputs["g_bias"], + lower_bound=LOWER_BOUND, + checkpoint_state=initial_state, + correction_cache=correction, + kg_cache=key_gate, + query_start_loc=inputs["cu_seqlens"], + state_indices=recovery_indices, + spec_query_len=query_len, + ) + torch.testing.assert_close(recovered_output, baseline_output, rtol=3e-2, atol=3e-2) + + history = 3 + conv_shape = ( + (initial_state.shape[0], 12, history + query_len - 1) + if is_conv_state_dim_first() + else (initial_state.shape[0], history + query_len - 1, 12) + ) + layer = SimpleNamespace( + kv_cache=( + torch.zeros(conv_shape, dtype=torch.bfloat16, device=device), + initial_state, + correction, + key_gate, + ), + A_log=inputs["a_log"], + dt_bias=inputs["g_bias"], + local_num_heads=H, + head_dim=D, + gate_lower_bound=LOWER_BOUND, + ) + context = KDARecoverSSMCommitContext.create( + [layer], spec_query_len=query_len, max_num_reqs=num_seqs + ) + accepted = torch.tensor([1, 3], dtype=torch.int32, device=device) + context.commit(accepted, recovery_indices, inputs["cu_seqlens"]) + for seq, count in enumerate(accepted.tolist()): + baseline_index = inputs["ssm_state_indices"][seq, count - 1] + torch.testing.assert_close( + initial_state[recovery_indices[seq]], + baseline_state[baseline_index], + rtol=3e-3, + atol=3e-3, + ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3e5d4ff160f5..cd4b2d8bfc2a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3395,8 +3395,12 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": if self.model_config is not None and self.model_config.architecture not in ( "KimiLinearForCausalLM", "KimiK3ForConditionalGeneration", + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", ): - raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA") + raise ValueError( + "RecoverSSM is only supported for Kimi-K3 and GLM-5.3 KDA" + ) if self.mamba_config.enable_stochastic_rounding: raise ValueError( "RecoverSSM supports bfloat16/float32 " diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 47ad93288542..edea2b39410d 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -360,23 +360,39 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): def get_state_dtype( self, - ) -> tuple[torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, ...]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") - return MambaStateDtypeCalculator.kda_state_dtype( - self.model_config.dtype, self.cache_config.mamba_cache_dtype + dtypes = MambaStateDtypeCalculator.kda_state_dtype( + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, ) + if self.cache_config.use_kda_recoverssm: + return MambaStateDtypeCalculator.append_kda_recoverssm_record( + dtypes, self.model_config.dtype + ) + return dtypes def get_state_shape( self, - ) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.kda_state_shape( + ) -> tuple[tuple[int, ...], ...]: + shapes = MambaStateShapeCalculator.kda_state_shape( self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size, num_spec=self.num_spec, ) + if self.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + shapes, + self.num_heads, + self.head_dim, + tp_world_size=self.tp_size, + spec_query_len=1 + self.num_spec, + ) + return shapes def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: spec = super().get_kv_cache_spec(vllm_config) @@ -1285,7 +1301,8 @@ def rearrange_mixed_qkv( def _can_use_b12x_kda_decode(self, m: GDNAttentionMetadata) -> bool: if ( - self._b12x_kda_plan is None + self.cache_config.use_kda_recoverssm + or self._b12x_kda_plan is None or m.num_prefills != 0 or (m.num_decodes == 0 and m.num_spec_decodes == 0) ): @@ -1541,9 +1558,8 @@ def _forward( g2_actual = g2[:num_actual_tokens] use_b12x_kda = self._can_use_b12x_kda_decode(m) - constant_caches = self.kv_cache - - conv_state, recurrent_state = constant_caches + conv_state, recurrent_state = self.kv_cache[:2] + recoverssm_records = self.kv_cache[2:] # conv_state must be (..., dim, width-1) for the conv kernels. # DS layout stores it that way directly; SD layout needs a transpose. if not is_conv_state_dim_first(): @@ -1586,7 +1602,11 @@ def _forward( assert spec_state_indices_tensor is not None assert spec_query_start_loc is not None spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] - spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_max_query_len = ( + self.num_spec + 1 + if self.cache_config.use_kda_recoverssm + else spec_state_indices_tensor.size(-1) + ) # Sibling beta and, for full-rank gates, output-gate views remain # live, so write the convolution output separately. @@ -1635,21 +1655,50 @@ def _forward( if m.num_prefills == 0 and m.num_decodes == 0 else None ) - core_attn_out_spec, _ = fused_recurrent_kda( - q=q_spec, - k=k_spec, - v=v_spec, - raw_g=g1_spec, - raw_beta=beta_spec, - A_log=self.A_log, - dt_bias=self.dt_bias, - lower_bound=self.gate_lower_bound, - initial_state=recurrent_state, - cu_seqlens=spec_cu_seqlens, - ssm_state_indices=spec_state_indices_tensor, - num_accepted_tokens=num_accepted_tokens, - out=spec_out, - ) + if self.cache_config.use_kda_recoverssm: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + kda_recoverssm_verify, + ) + + if len(recoverssm_records) != 2: + raise ValueError( + "KDA RecoverSSM requires correction and key/gate buffers" + ) + core_attn_out_spec = kda_recoverssm_verify( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + checkpoint_state=recurrent_state, + correction_cache=recoverssm_records[0], + kg_cache=recoverssm_records[1], + query_start_loc=spec_cu_seqlens, + state_indices=spec_state_indices_tensor[ + : m.num_spec_decodes, 0 + ], + spec_query_len=self.num_spec + 1, + out=spec_out, + ) + else: + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) # ---------- non-spec path (prefill or plain decode) ---------- core_attn_out_non_spec = None diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 938136ab242b..d54d69f86a49 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -14,6 +14,8 @@ import os from copy import copy +from dataclasses import dataclass +from typing import TYPE_CHECKING import torch from einops import rearrange @@ -22,7 +24,13 @@ from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( KimiGatedDeltaNetAttention, ) -from vllm.v1.attention.backend import AttentionBackend +from vllm.models.kimi_k3.nvidia.kda_metadata import ( + KDARecoverSSMAlignMetadata, + KDARecoverSSMCommitMetadata, + KimiK3KDAMetadata, +) +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, GDNAttentionMetadataBuilder, @@ -33,6 +41,11 @@ use_preallocated_workspace, ) +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + class Glm5NextKDAMetadataBuilder(GDNAttentionMetadataBuilder): supports_varlen_decode_cudagraph = True @@ -64,6 +77,110 @@ def supports_device_cpu_query_lens_mismatch(cls) -> bool: return True +@dataclass +class Glm5NextRecoverKDAMetadata(KimiK3KDAMetadata): + """GLM prefill metadata with KDA speculative-state recovery.""" + + +class Glm5NextRecoverKDAMetadataBuilder(Glm5NextKDAMetadataBuilder): + supports_update_block_table = False + + def __init__( + self, + kv_cache_spec: MambaSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self._neutral_accepted = torch.ones( + vllm_config.scheduler_config.max_num_seqs, + dtype=torch.int32, + device=device, + ) + self._recoverssm_context: KDARecoverSSMCommitContext | None = None + + def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": + if self._recoverssm_context is None: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + + layers = self.vllm_config.compilation_config.static_forward_context + self._recoverssm_context = KDARecoverSSMCommitContext.create( + [layers[name] for name in self.layer_names], + spec_query_len=self.num_spec + 1, + max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, + ) + return self._recoverssm_context + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> Glm5NextRecoverKDAMetadata: + common = common_attn_metadata + if num_decode_draft_tokens_cpu is not None: + assert common.is_prefilling is not None + active_decodes = (~common.is_prefilling) & ( + common.query_start_loc_cpu.diff() > 0 + ) + missing = active_decodes & (num_decode_draft_tokens_cpu < 0) + if bool(torch.any(missing)): + num_decode_draft_tokens_cpu = num_decode_draft_tokens_cpu.clone() + num_decode_draft_tokens_cpu[missing] = 0 + metadata = super().build( + common_prefix_len, + common, + num_accepted_tokens=self._neutral_accepted[: common.num_reqs], + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + fast_build=fast_build, + ) + commit = None + if metadata.num_spec_decodes > 0: + assert metadata.spec_sequence_masks_cpu is not None + assert metadata.spec_state_indices_tensor is not None + assert metadata.spec_query_start_loc is not None + request_indices = async_tensor_h2d( + metadata.spec_sequence_masks_cpu.nonzero(as_tuple=True)[0], + dtype=torch.int32, + device=common.query_start_loc.device, + ) + align = None + if self.kv_cache_spec.mamba_cache_mode == "align": + align = KDARecoverSSMAlignMetadata( + block_table=common.block_table_tensor, + num_computed_tokens=common.compute_num_computed_tokens(), + block_size=self.kv_cache_spec.block_size, + ) + commit = KDARecoverSSMCommitMetadata( + state_indices=metadata.spec_state_indices_tensor, + query_start_loc=metadata.spec_query_start_loc, + request_indices=request_indices, + align=align, + ) + return Glm5NextRecoverKDAMetadata( + **vars(metadata), + recoverssm_commit=commit, + recoverssm_context=( + self._get_recoverssm_context() if commit is not None else None + ), + ) + + +class Glm5NextRecoverKDAAttentionBackend(Glm5NextKDAAttentionBackend): + @staticmethod + def get_name() -> str: + return "GLM5NEXT_RECOVER_KDA" + + @staticmethod + def get_builder_cls() -> type[Glm5NextRecoverKDAMetadataBuilder]: + return Glm5NextRecoverKDAMetadataBuilder + + _GATE_SIDE_STREAM = os.getenv("VLLM_GLM53_KDA_GATE_SIDE_STREAM", "1") != "0" _side_streams: dict[int, torch.cuda.Stream] = {} @@ -112,6 +229,7 @@ class Glm5NextLinearAttention(KimiGatedDeltaNetAttention): b12x_kda_null_state_index = 0 def __init__(self, config, vllm_config: VllmConfig, prefix: str = "") -> None: + self.enable_b12x_kda_decode = not vllm_config.cache_config.use_kda_recoverssm quant_config = vllm_config.quant_config if ( quant_config is not None @@ -127,6 +245,8 @@ def __init__(self, config, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__(config, vllm_config, prefix) def get_attn_backend(self) -> type[AttentionBackend]: + if self.cache_config.use_kda_recoverssm: + return Glm5NextRecoverKDAAttentionBackend if ( self.speculative_config is not None and self.speculative_config.enable_adaptive_verification diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index a1f5adecf675..2382d143faee 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -66,6 +66,7 @@ IsHybrid, MixtureOfExperts, SupportsEagle3, + SupportsReplaySSM, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -1306,7 +1307,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Glm5NextForCausalLM( - nn.Module, HasInnerState, MixtureOfExperts, IsHybrid, SupportsEagle3 + nn.Module, + HasInnerState, + MixtureOfExperts, + IsHybrid, + SupportsEagle3, + SupportsReplaySSM, ): packed_modules_mapping = GLM5NEXT_PACKED_MODULES_MAPPING supports_pp: ClassVar[Literal[False]] = False @@ -1364,15 +1370,30 @@ def forward( def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) -> tuple[torch.dtype, ...]: + dtypes = MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, ) + if vllm_config.cache_config.use_kda_recoverssm: + dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( + dtypes, vllm_config.model_config.dtype + ) + return dtypes @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig" - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> ( + tuple[tuple[int, int], tuple[int, int, int]] + | tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int, int], + ] + ): parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config tp_size = parallel_config.tensor_parallel_size @@ -1381,13 +1402,22 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) - return MambaStateShapeCalculator.kda_state_shape( + shapes = MambaStateShapeCalculator.kda_state_shape( tp_size, hf_config.linear_num_heads, hf_config.linear_head_dim, conv_kernel_size=hf_config.linear_conv_kernel_dim, num_spec=num_spec, ) + if vllm_config.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + shapes, + hf_config.linear_num_heads, + hf_config.linear_head_dim, + tp_world_size=tp_size, + spec_query_len=1 + num_spec, + ) + return shapes @classmethod def get_mamba_state_copy_func( @@ -1422,7 +1452,11 @@ def process_weights_after_loading(self) -> None: dummy_inputs=Glm4vDummyInputsBuilder, ) class Glm5NextForConditionalGeneration( - Glm4vForConditionalGeneration, HasInnerState, IsHybrid, SupportsEagle3 + Glm4vForConditionalGeneration, + HasInnerState, + IsHybrid, + SupportsEagle3, + SupportsReplaySSM, ): packed_modules_mapping = GLM5NEXT_PACKED_MODULES_MAPPING # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 133983779240..51bbd2809086 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -198,6 +198,9 @@ def __init__( else: self.num_spec = 0 self.use_spec_decode: bool = self.num_spec > 0 + self.state_index_columns = ( + 1 if vllm_config.cache_config.use_kda_recoverssm else self.num_spec + 1 + ) self._b12x_mixed = ( B12xGdnMixedMetadata( max_tokens=vllm_config.scheduler_config.max_num_batched_tokens, @@ -206,6 +209,7 @@ def __init__( device=device, ) if self.gdn_prefill_backend == "b12x" + and not vllm_config.cache_config.use_kda_recoverssm else None ) self._init_reorder_batch_threshold(1, self.use_spec_decode) @@ -223,7 +227,7 @@ def __init__( ) self.spec_state_indices_tensor: torch.Tensor = torch.empty( - (self.decode_cudagraph_max_bs, self.num_spec + 1), + (self.decode_cudagraph_max_bs, self.state_index_columns), dtype=torch.int32, device=device, ) @@ -264,7 +268,10 @@ def __init__( ) self._decode_state_indices_source: torch.Tensor | None = None self._decode_state_indices_view: torch.Tensor | None = None - self._reuse_spec_decode_inputs = envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH + self._reuse_spec_decode_inputs = ( + envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH + and not vllm_config.cache_config.use_kda_recoverssm + ) # Constant sources for the uniform spec-decode fast path. They are # copied into the builder-owned graph buffers above, never handed to # the layers directly: a full cudagraph captured from a uniform batch @@ -557,7 +564,7 @@ def build( # type: ignore[override] ) # Filter by spec_sequence_masks to exclude padded sequences spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks_cpu, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.state_index_columns ] non_spec_state_indices_tensor = None # Padded sequences are always at the back, so the first @@ -578,7 +585,7 @@ def build( # type: ignore[override] spec_token_indx = index[num_non_spec_tokens:] spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks_cpu, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.state_index_columns ] non_spec_state_indices_tensor = block_table_tensor[ non_spec_sequence_masks_cpu, 0 @@ -960,7 +967,7 @@ def update_block_table( else: non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu spec_state_indices = state_indices[ - spec_sequence_masks_cpu, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.state_index_columns ] non_spec_state_indices = state_indices[non_spec_sequence_masks_cpu, 0] prefill_state_indices = metadata.prefill_state_indices From a281a59c6c78f9c595110ce8cde41b1ad60df08e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 14:11:37 +0000 Subject: [PATCH 131/365] feat(glm5next): use B12X verification and accepted-state recovery Route GLM speculative KDA verification and grouped FP32 checkpoint recovery through prepared B12X CuTe executables. Pure decode and mixed prefill/speculative batches bind caller-owned record buffers; mixed batches retain one merged output normalization. Cache-boundary metadata and convolution compaction remain shared with Kimi. Require B12X record/commit support when recovery is enabled; ordinary serving and Kimi recovery remain unchanged. Ten vLLM metadata/recurrence tests and all changed-file pre-commit hooks pass. TP2/DCP2 Spark MTP3 serving passes full-and-piecewise graph startup, text, prefix reuse, vision and bounded decode/prefill checks. Builds on the Kimi-K3 RecoverSSM integration by Jiangyun Zhu (70afdedc108). Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit ade05be4ad76519060dd7648433074058568ea5f) --- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 36 +++++++++++-- vllm/models/glm5next/nvidia/kda.py | 7 ++- vllm/models/glm5next/nvidia/ops/recoverssm.py | 54 +++++++++++++++++++ vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 8 ++- 4 files changed, 97 insertions(+), 8 deletions(-) create mode 100644 vllm/models/glm5next/nvidia/ops/recoverssm.py diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index edea2b39410d..c4a854fe241c 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -618,6 +618,10 @@ def _initialize_b12x_kda_decode(self, vllm_config: VllmConfig) -> None: or not api.is_supported(device) ): return + if self.cache_config.use_kda_recoverssm and not hasattr(api, "bind_kda_commit"): + raise RuntimeError( + "GLM KDA state recovery requires B12X record/commit support" + ) max_seqs = int(vllm_config.scheduler_config.max_num_seqs) state_index_columns = max(1, self.num_spec + 1) if state_index_columns > 8: @@ -716,6 +720,7 @@ def _b12x_kda_decode_declaration(self, max_state_slots: int): gate_activation="sigmoid", qk_l2norm=True, null_state_index=self.b12x_kda_null_state_index, + recover_speculative_state=self.cache_config.use_kda_recoverssm, ) return api.plan( caps, @@ -1030,6 +1035,7 @@ def reset(): num_seqs=num_seqs, num_tokens=num_tokens, output=output, + **self._b12x_recovery_records(), ) return PreparedCall( run=lambda: state.run( @@ -1301,8 +1307,7 @@ def rearrange_mixed_qkv( def _can_use_b12x_kda_decode(self, m: GDNAttentionMetadata) -> bool: if ( - self.cache_config.use_kda_recoverssm - or self._b12x_kda_plan is None + self._b12x_kda_plan is None or m.num_prefills != 0 or (m.num_decodes == 0 and m.num_spec_decodes == 0) ): @@ -1334,6 +1339,7 @@ def _run_b12x_kda_decode_post_conv( query_start_loc: torch.Tensor, num_accepted_tokens: torch.Tensor | None, num_requests: int, + apply_output_norm: bool = True, ) -> None: """Execute B12X KDA after the convolution projection. @@ -1431,14 +1437,21 @@ def _run_b12x_kda_decode_post_conv( num_seqs=num_seqs, num_tokens=num_tokens_tensor, output=output, + **self._b12x_recovery_records(), ) api.run_kda( binding, lower_bound=self.gate_lower_bound, eps=self.o_norm.eps, scale=self.head_dim**-0.5, + apply_output_norm=apply_output_norm, ) + def _b12x_recovery_records(self) -> dict[str, torch.Tensor]: + if not self.cache_config.use_kda_recoverssm: + return {} + return {"correction_cache": self.kv_cache[2], "kg_cache": self.kv_cache[3]} + def forward( self, hidden_states: torch.Tensor, @@ -1655,7 +1668,24 @@ def _forward( if m.num_prefills == 0 and m.num_decodes == 0 else None ) - if self.cache_config.use_kda_recoverssm: + if self.cache_config.use_kda_recoverssm and self.enable_b12x_kda_decode: + if spec_out is None: + spec_out = torch.empty_like(v_spec) + self._run_b12x_kda_decode_post_conv( + metadata=m, + mixed_qkv=mixed_qkv_spec, + raw_g=g1_spec[0], + raw_beta=beta_spec[0], + z=g1_spec[0], + output=spec_out[0], + state_indices=spec_state_indices_tensor, + query_start_loc=spec_cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + num_requests=m.num_spec_decodes, + apply_output_norm=False, + ) + core_attn_out_spec = spec_out + elif self.cache_config.use_kda_recoverssm: from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( kda_recoverssm_verify, ) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index d54d69f86a49..d01e5de69d43 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -102,12 +102,12 @@ def __init__( def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": if self._recoverssm_context is None: - from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( - KDARecoverSSMCommitContext, + from vllm.models.glm5next.nvidia.ops.recoverssm import ( + B12XKDARecoverSSMCommitContext, ) layers = self.vllm_config.compilation_config.static_forward_context - self._recoverssm_context = KDARecoverSSMCommitContext.create( + self._recoverssm_context = B12XKDARecoverSSMCommitContext.create( [layers[name] for name in self.layer_names], spec_query_len=self.num_spec + 1, max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, @@ -229,7 +229,6 @@ class Glm5NextLinearAttention(KimiGatedDeltaNetAttention): b12x_kda_null_state_index = 0 def __init__(self, config, vllm_config: VllmConfig, prefix: str = "") -> None: - self.enable_b12x_kda_decode = not vllm_config.cache_config.use_kda_recoverssm quant_config = vllm_config.quant_config if ( quant_config is not None diff --git a/vllm/models/glm5next/nvidia/ops/recoverssm.py b/vllm/models/glm5next/nvidia/ops/recoverssm.py new file mode 100644 index 000000000000..58f170230481 --- /dev/null +++ b/vllm/models/glm5next/nvidia/ops/recoverssm.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""B12X recurrent-state recovery with vLLM's cache-boundary metadata.""" + +from typing import Any, cast + +import torch + +from vllm.models.kimi_k3.nvidia.ops.recoverssm import KDARecoverSSMCommitContext + + +class B12XKDARecoverSSMCommitContext(KDARecoverSSMCommitContext): + _b12x_layer: Any + _b12x_scratch: torch.Tensor + + @classmethod + def create(cls, layers, *, spec_query_len, max_num_reqs): + context = cast( + "B12XKDARecoverSSMCommitContext", + super().create( + layers, spec_query_len=spec_query_len, max_num_reqs=max_num_reqs + ), + ) + layer = layers[0] + if layer._b12x_kda_plan is None: + raise RuntimeError("B12X KDA recovery requires a prepared decode plan") + context._b12x_layer = layer + (spec,) = layer._b12x_kda_plan.scratch_specs() + context._b12x_scratch = torch.empty( + spec.shape, dtype=spec.dtype, device=context.checkpoints[0].device + ) + return context + + def _commit_recurrent_state(self, state_indices, batch, align_mode): + layer = self._b12x_layer + api = layer._b12x_kda_api + binding = api.bind_kda_commit( + layer._b12x_kda_plan, + scratch=self._b12x_scratch, + state_base_addrs=self.state_base_addrs, + state_block_strides=self.state_block_strides, + correction_cache_base_addrs=self.correction_cache_base_addrs, + correction_cache_block_strides=self.correction_cache_block_strides, + kg_cache_base_addrs=self.kg_cache_base_addrs, + kg_cache_block_strides=self.kg_cache_block_strides, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=state_indices, + commit_lens=self.commit_lens[:batch], + final_state_indices=self.final_state_indices[:batch], + boundary_state_indices=self.boundary_state_indices[:batch], + boundary_recovery_lens=self.boundary_recovery_lens[:batch], + ) + api.run_kda_commit(binding, lower_bound=self.lower_bound) diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index c603ec7ed1fc..2c0a79279fc3 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -1009,7 +1009,13 @@ def commit( num_warps=4, ) + self._commit_recurrent_state(state_indices, batch, block_table is not None) + + def _commit_recurrent_state( + self, state_indices: torch.Tensor, batch: int, align_mode: bool + ) -> None: state_ref = self.checkpoints[0] + num_layers = len(self.checkpoints) _, num_heads, value_dim, key_dim = state_ref.shape block_k = triton.next_power_of_2(key_dim) block_v = min(triton.next_power_of_2(value_dim), 32) @@ -1058,7 +1064,7 @@ def commit( BV=block_v, NUM_HEADS=num_heads, USE_LOWER_BOUND=self.lower_bound is not None, - ALIGN_MODE=block_table is not None, + ALIGN_MODE=align_mode, num_warps=4, num_stages=2, ) From 4eefb884ea4ca46922bbb60c802ea4ed5e919d85 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 15:00:31 +0000 Subject: [PATCH 132/365] feat(glm5next): recover speculative state before atomic checkpoint export Select native B12X state recovery automatically for eligible GLM MTP and DFlash on SM12x, preserving the explicit opt-out and other model defaults. Export request-boundary states after acceptance and convolution compaction so atomic external caches can restore committed state. Reject connectors without an atomic boundary adapter. Keep FP32 recurrent precision and cache layout identity. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/glm5next/test_kda_recoverssm.py | 79 +++++++++++++++++++ tests/test_config.py | 2 +- tests/v1/worker/test_gpu_model_runner_v2.py | 42 ++++++++++ tests/v1/worker/test_mamba_utils.py | 16 +++- vllm/config/cache.py | 11 ++- vllm/config/vllm.py | 82 +++++++++++++++++++- vllm/engine/arg_utils.py | 2 +- vllm/v1/worker/gpu/boundary_checkpoint.py | 9 ++- vllm/v1/worker/gpu/model_runner.py | 13 +++- vllm/v1/worker/gpu_model_runner.py | 6 +- vllm/v1/worker/mamba_utils.py | 14 +++- 11 files changed, 260 insertions(+), 16 deletions(-) diff --git a/tests/models/glm5next/test_kda_recoverssm.py b/tests/models/glm5next/test_kda_recoverssm.py index 95a97c3aaaef..eece79d7d235 100644 --- a/tests/models/glm5next/test_kda_recoverssm.py +++ b/tests/models/glm5next/test_kda_recoverssm.py @@ -11,7 +11,9 @@ BatchSpec, create_common_attn_metadata, ) +from vllm.config import CacheConfig, VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.config.mamba import MambaBackendEnum from vllm.models.glm5next.nvidia.kda import ( Glm5NextRecoverKDAMetadataBuilder, ) @@ -26,6 +28,83 @@ def disable_pinned_host_staging(monkeypatch): monkeypatch.setattr("vllm.v1.attention.backends.utils.PIN_MEMORY", False) +def _recovery_config(monkeypatch): + monkeypatch.setattr("vllm.platforms.current_platform.is_cuda", lambda: True) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_device_capability", + lambda: SimpleNamespace(major=12), + ) + monkeypatch.setattr( + "vllm.utils.b12x.get_b12x_gdn_decode", + lambda: SimpleNamespace(bind_kda_commit=Mock(), is_supported=lambda: True), + ) + return SimpleNamespace( + cache_config=CacheConfig(mamba_cache_mode="align"), + model_config=SimpleNamespace( + architecture="Glm5NextForConditionalGeneration", + supports_replayssm=True, + dtype=torch.bfloat16, + hf_text_config=SimpleNamespace(linear_head_dim=128), + ), + speculative_config=SimpleNamespace(method="mtp"), + num_speculative_tokens=3, + use_v2_model_runner=True, + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + mamba_config=SimpleNamespace( + backend=MambaBackendEnum.TRITON, enable_stochastic_rounding=False + ), + kv_transfer_config=None, + use_request_boundary_checkpoints=True, + ) + + +@pytest.mark.parametrize("method,depth", [("mtp", 3), ("dflash", 7)]) +@pytest.mark.parametrize("external", [False, True]) +def test_glm_recovery_defaults_to_native_with_atomic_cache( + monkeypatch, method, depth, external +): + config = _recovery_config(monkeypatch) + config.speculative_config.method = method + config.num_speculative_tokens = depth + if external: + config.kv_transfer_config = SimpleNamespace(is_kv_transfer_instance=True) + VllmConfig.validate_mamba_cached_kernel(config) + assert config.cache_config.use_replayssm is True + assert config.cache_config.use_kda_recoverssm is True + + +@pytest.mark.parametrize("unsupported", ["depth", "connector", "state", "v1", "b12x"]) +def test_glm_recovery_leaves_unsupported_configs_unchanged(monkeypatch, unsupported): + config = _recovery_config(monkeypatch) + if unsupported == "depth": + config.num_speculative_tokens = 10 + elif unsupported == "connector": + config.kv_transfer_config = SimpleNamespace(is_kv_transfer_instance=True) + config.use_request_boundary_checkpoints = False + elif unsupported == "state": + config.cache_config.mamba_ssm_cache_dtype = "bfloat16" + elif unsupported == "v1": + config.use_v2_model_runner = False + else: + monkeypatch.setattr("vllm.utils.b12x.get_b12x_gdn_decode", lambda: None) + VllmConfig.validate_mamba_cached_kernel(config) + assert not config.cache_config.use_kda_recoverssm + config.cache_config.use_replayssm = True + with pytest.raises(ValueError, match="GLM KDA recovery requires"): + VllmConfig.validate_mamba_cached_kernel(config) + + +def test_glm_recovery_respects_opt_out_and_other_model_defaults(monkeypatch): + config = _recovery_config(monkeypatch) + config.cache_config.use_replayssm = False + VllmConfig.validate_mamba_cached_kernel(config) + assert not config.cache_config.use_kda_recoverssm + config.cache_config.use_replayssm = None + config.model_config.architecture = "KimiLinearForCausalLM" + VllmConfig.validate_mamba_cached_kernel(config) + assert not config.cache_config.use_kda_recoverssm + + def test_glm_recoverssm_reserves_one_recurrent_state_per_draft_window(): config = SimpleNamespace( model_config=SimpleNamespace( diff --git a/tests/test_config.py b/tests/test_config.py index 7df3da2b24da..e378e0ceb3f8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -227,7 +227,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): config.cache_config.mamba_cache_mode = "none" config.model_config.architecture = "NemotronHForCausalLM" - with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"): + with pytest.raises(ValueError, match="only supported for Kimi-K3 and GLM-5.3 KDA"): VllmConfig.validate_mamba_cached_kernel(config) config.model_config.architecture = "KimiLinearForCausalLM" diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 35d9f2f5d491..d67c5977e56c 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -22,6 +22,48 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner +@pytest.mark.parametrize("recovery", [False, True]) +def test_boundary_capture_uses_recovered_accepted_state(monkeypatch, recovery): + """A stop inside a draft window must be applied before checkpoint export.""" + events = [] + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.is_last_pp_rank = False + runner.cache_config = SimpleNamespace(use_kda_recoverssm=recovery) + runner.req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([19])), + last_sampled_tokens=None, + all_token_ids=SimpleNamespace(gpu=None), + total_len=SimpleNamespace(gpu=None), + ) + sampled = torch.tensor([4]) + + def update(*args): + sampled.fill_(2) # The second accepted token terminates the response. + events.append("trim") + + def commit(*args): + assert sampled.item() == 2 + events.append("commit") + + def capture(*args, accepted_state_committed=False): + assert accepted_state_committed is recovery + events.append("capture") + + monkeypatch.setattr(model_runner_module, "post_update", update) + runner.model_state = SimpleNamespace(postprocess_state=commit) + runner.boundary_checkpoint_state = SimpleNamespace(capture_mamba=capture) + runner.postprocess_sampled( + torch.tensor([0]), + torch.tensor([[1, 2, 3, 4]]), + sampled, + torch.tensor([0]), + boundary_capture=torch.empty(3, 1, 3), + ) + assert events == ( + ["trim", "commit", "capture"] if recovery else ["trim", "capture", "commit"] + ) + + def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): runner = GPUModelRunner.__new__(GPUModelRunner) runner.max_model_len = 262144 diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 587a350b8789..948beca220f3 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -1072,7 +1072,10 @@ def t(values): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_boundary_checkpoint_copies_only_selected_committed_states(monkeypatch): +@pytest.mark.parametrize("accepted_state_committed", [False, True]) +def test_boundary_checkpoint_copies_only_selected_committed_states( + monkeypatch, accepted_state_committed +): """Prompt/response copies respect request slots and speculative rollback.""" monkeypatch.setattr( "vllm.v1.worker.mamba_utils.is_conv_state_dim_first", lambda: False @@ -1100,7 +1103,14 @@ def t(values): bias = t([[0, 2, 0], [0, 1, 0]]) def copy(): - ctx.checkpoint_request_boundaries(idx, states, capture, bias, destinations) + ctx.checkpoint_request_boundaries( + idx, + states, + capture, + bias, + destinations, + accepted_state_committed=accepted_state_committed, + ) copy() graph = torch.cuda.CUDAGraph() @@ -1109,6 +1119,8 @@ def copy(): graph.replay() for layer in range(len(names)): for dst, src, shift in ((20, 1, 0), (21, 1, 2), (23, 15, 1)): + if accepted_state_committed: + shift = 0 torch.testing.assert_close( conv[layer][dst, : cfg.conv_width - shift], conv_ref[layer][src, shift:], diff --git a/vllm/config/cache.py b/vllm/config/cache.py index a874be06494e..5fdf58b66a71 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -213,15 +213,20 @@ class CacheConfig: """ReplaySSM logical history length B for Mamba2. Triton uses B physical rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B. Default 16.""" - use_replayssm: bool = False - """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip + use_replayssm: bool | None = None + """Enable checkpoint recovery. None selects B12X recovery automatically for + supported GLM-5.3 speculative decoding, including atomic request-boundary + external caching. False (--no-use-replayssm) retains full speculative states. + Other models remain opt-in. + + For Mamba2, use the ReplaySSM decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton or FlashInfer mamba backend; standard (non-speculative) decode only. In align mode flushes are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) - """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" + """Whether Kimi-K3 or GLM-5.3 KDA uses speculative state recovery.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index cd4b2d8bfc2a..a7695fac3e6a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -68,6 +68,60 @@ logger = init_logger(__name__) + +def _glm_kda_recovery_unavailable(config: "VllmConfig") -> str | None: + """Describe why GLM cannot use the native checkpoint/record contract.""" + from vllm.platforms import current_platform + + cache = config.cache_config + model = config.model_config + spec = config.speculative_config + if model is None or model.architecture not in ( + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + ): + return "requires a GLM-5.3 target" + if ( + spec is None + or spec.method not in ("mtp", "dflash") + or not (1 <= config.num_speculative_tokens <= 7) + ): + return "requires MTP or DFlash with 1 to 7 draft tokens" + if ( + model.dtype != torch.bfloat16 + or model.hf_text_config.linear_head_dim != 128 + or cache.mamba_ssm_cache_dtype not in ("auto", "float32") + ): + return "requires BF16 activations, 128-wide heads and FP32 recurrent state" + if config.mamba_config.enable_stochastic_rounding: + return "does not support stochastic recurrent-state rounding" + if cache.mamba_cache_mode not in ("none", "align"): + return "requires none or align Mamba cache mode" + if not config.use_v2_model_runner: + return "requires Model Runner V2" + if config.parallel_config.pipeline_parallel_size != 1: + return "requires pipeline_parallel_size=1" + if config.mamba_config.backend != MambaBackendEnum.TRITON: + return "requires the triton Mamba metadata backend" + if ( + config.kv_transfer_config is not None + and config.kv_transfer_config.is_kv_transfer_instance + and not config.use_request_boundary_checkpoints + ): + return "requires an atomic request-boundary external-cache connector" + if not current_platform.is_cuda(): + return "requires CUDA" + capability = current_platform.get_device_capability() + if capability is None or capability.major != 12: + return "requires an SM12x GPU" + from vllm.utils.b12x import get_b12x_gdn_decode + + api = get_b12x_gdn_decode() + if api is None or not hasattr(api, "bind_kda_commit") or not api.is_supported(): + return "requires B12X speculative KDA verification and commit support" + return None + + # TODO(rocm): These models are either unsupported by MRV2 or slower with # MRV2 on AMD GPUs. ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset( @@ -3381,11 +3435,29 @@ def validate_mamba_block_size(self) -> "VllmConfig": @model_validator(mode="after") def validate_mamba_cached_kernel(self) -> "VllmConfig": + if self.cache_config.use_replayssm is None: + self.cache_config.use_replayssm = ( + _glm_kda_recovery_unavailable(self) is None + ) + if self.cache_config.use_replayssm: + logger.info( + "GLM speculative KDA uses B12X checkpoint recovery with FP32 " + "recurrent state. Use --no-use-replayssm to retain full " + "speculative states." + ) if not self.cache_config.use_replayssm: self.cache_config.use_kda_recoverssm = False return self self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 + if self.model_config is not None and self.model_config.architecture in ( + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + ): + reason = _glm_kda_recovery_unavailable(self) + if reason is not None: + raise ValueError(f"GLM KDA recovery {reason}") + if self.model_config is not None and not self.model_config.supports_replayssm: raise ValueError( "--use-replayssm is not supported for architecture " @@ -3447,10 +3519,18 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": if ( self.kv_transfer_config is not None and self.kv_transfer_config.is_kv_transfer_instance + and not ( + self.cache_config.use_kda_recoverssm + and self.model_config is not None + and self.model_config.architecture + in ("Glm5NextForCausalLM", "Glm5NextForConditionalGeneration") + and self.use_request_boundary_checkpoints + ) ): raise ValueError( "--use-replayssm is incompatible with KV connectors " - "(P/D disaggregation, KV cache offload)" + "except GLM KDA recovery with an atomic request-boundary " + "checkpoint connector" ) return self diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 96f4675f72c0..54c19d9ef2f0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -817,7 +817,7 @@ class EngineArgs: CacheConfig.enable_mamba_fine_grained_prefix_cache ) replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len - use_replayssm: bool = CacheConfig.use_replayssm + use_replayssm: bool | None = CacheConfig.use_replayssm mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON mamba_ssu_algorithm: MambaSSUAlgorithm | None = None diff --git a/vllm/v1/worker/gpu/boundary_checkpoint.py b/vllm/v1/worker/gpu/boundary_checkpoint.py index 82c16b0628e8..87d4dc19ec82 100644 --- a/vllm/v1/worker/gpu/boundary_checkpoint.py +++ b/vllm/v1/worker/gpu/boundary_checkpoint.py @@ -648,7 +648,13 @@ def add_request(self, slot: int, request: "NewRequestData") -> None: slot, checkpoint.num_tokens - 1 ) - def capture_mamba(self, idx_mapping: torch.Tensor, capture: torch.Tensor) -> None: + def capture_mamba( + self, + idx_mapping: torch.Tensor, + capture: torch.Tensor, + *, + accepted_state_committed: bool = False, + ) -> None: state = self.model_state if state._mamba_ctx is not None: state._mamba_ctx.checkpoint_request_boundaries( @@ -657,6 +663,7 @@ def capture_mamba(self, idx_mapping: torch.Tensor, capture: torch.Tensor) -> Non capture[0], capture[1], self.mamba_blocks, + accepted_state_committed=accepted_state_committed, ) def capture_auxiliary( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43da8b0e4646..f7621e174314 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -726,7 +726,7 @@ def initialize_kv_cache( initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config, - use_replayssm=self.cache_config.use_replayssm, + use_replayssm=bool(self.cache_config.use_replayssm), ) piecewise_capture_available = bool( envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model) @@ -1788,7 +1788,11 @@ def postprocess_sampled( boundary_capture, ) - if boundary_capture is not None: + # Full-checkpoint kernels retain every speculative state, so capture + # selects an accepted column before postprocessing compacts it. Recovery + # kernels only materialize the accepted state during postprocessing. + recover_state = self.cache_config.use_kda_recoverssm + if boundary_capture is not None and not recover_state: assert self.boundary_checkpoint_state is not None self.boundary_checkpoint_state.capture_mamba(idx_mapping, boundary_capture) @@ -1798,6 +1802,11 @@ def postprocess_sampled( self.model_state.postprocess_state( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu ) + if boundary_capture is not None and recover_state: + assert self.boundary_checkpoint_state is not None + self.boundary_checkpoint_state.capture_mamba( + idx_mapping, boundary_capture, accepted_state_committed=True + ) def _merge_ec_connector_no_forward( self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1fc4022327a5..3033912e1efe 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -782,7 +782,7 @@ def __init__( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, + use_replayssm=bool(self.cache_config.use_replayssm), ) # Separate cuda stream for overlapping transfer of sampled token ids from @@ -7473,7 +7473,7 @@ def may_reinitialize_input_batch( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, + use_replayssm=bool(self.cache_config.use_replayssm), slot_mapping_modes=slot_mapping_modes, ) @@ -7589,7 +7589,7 @@ def initialize_kv_cache( initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config, - use_replayssm=self.vllm_config.cache_config.use_replayssm, + use_replayssm=bool(self.vllm_config.cache_config.use_replayssm), ) # The kernel block size for all KV cache groups. For example, if # kv_cache_manager uses block_size 256 for a given group, but the attention diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 6a438680f7f4..6b05709a033c 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -395,6 +395,7 @@ def checkpoint_mamba_states_kernel( NUM_CAPTURES: tl.constexpr, CONV_STATE_DIM_FIRST: tl.constexpr, TEMPORAL_TILES: tl.constexpr, + ACCEPTED_STATE_COMMITTED: tl.constexpr, ): batch_idx = tl.program_id(0) // NUM_CAPTURES kind = tl.program_id(0) % NUM_CAPTURES @@ -406,7 +407,9 @@ def checkpoint_mamba_states_kernel( if req_idx < 0: return src_col = tl.load(state_idx_ptr + req_idx) - token_bias = tl.load(capture_bias_ptr + batch_idx * NUM_CAPTURES + kind) + token_bias = 0 + if not ACCEPTED_STATE_COMMITTED: + token_bias = tl.load(capture_bias_ptr + batch_idx * NUM_CAPTURES + kind) group_idx = tl.load(state_group_indices_ptr + state_idx) destination = tl.load( destination_blocks_ptr @@ -1425,8 +1428,14 @@ def checkpoint_request_boundaries( capture_tokens: torch.Tensor, capture_bias: torch.Tensor, destination_blocks: torch.Tensor, + *, + accepted_state_committed: bool = False, ) -> None: - """Copy only requested endpoints, selecting the committed spec state.""" + """Copy endpoints from speculative columns or an already recovered state. + + Recovery commits and compacts both convolution and recurrent state. + Applying the original acceptance offset again would copy a stale page. + """ assert self.is_initialized num_reqs = idx_mapping.numel() assert ( @@ -1469,6 +1478,7 @@ def checkpoint_request_boundaries( NUM_CAPTURES=NUM_BOUNDARY_CHECKPOINT_SLOTS, CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), TEMPORAL_TILES=_TEMPORAL_TILES, + ACCEPTED_STATE_COMMITTED=accepted_state_committed, ) def run_fused_postprocess_align( From 4c954ffe6b97c32a64def88f12c27f66be36cb87 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 15:03:43 +0000 Subject: [PATCH 133/365] test(cache): derive boundary capture slots from the runtime contract Keep both full-checkpoint and accepted-state recovery coverage independent of optional endpoint count. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/worker/test_mamba_utils.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 948beca220f3..dff19a289439 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -16,6 +16,7 @@ get_temporal_copy_spec, ) from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -1094,13 +1095,19 @@ def t(values): tables = t([[1, 2, 3, 4, 5, 6], [14, 15, 16, 17, 18, 19]]) ctx.initialize_from_forward_context(config, context, _COPY_FUNCS, [tables]) - destinations = torch.zeros((8, 3, 1), device=device, dtype=torch.int32) - destinations[3, :, 0] = t([20, 21, 0]) - destinations[1, :, 0] = t([22, 23, 0]) + destinations = torch.zeros( + (8, NUM_BOUNDARY_CHECKPOINT_SLOTS, 1), device=device, dtype=torch.int32 + ) + destinations[3, :2, 0] = t([20, 21]) + destinations[1, :2, 0] = t([22, 23]) idx = t([3, 1]) states = t([0, 1, 0, 0, 0, 0, 0, 0]) - capture = t([[7, 9, 0], [0, 23, 0]]) - bias = t([[0, 2, 0], [0, 1, 0]]) + capture = torch.zeros( + (2, NUM_BOUNDARY_CHECKPOINT_SLOTS), device=device, dtype=torch.int32 + ) + bias = torch.zeros_like(capture) + capture[:, :2] = t([[7, 9], [0, 23]]) + bias[:, :2] = t([[0, 2], [0, 1]]) def copy(): ctx.checkpoint_request_boundaries( From f20f8d91e06886ec94389f190e73d98ebf9c09d3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 15:55:21 +0000 Subject: [PATCH 134/365] Document GLM recovery and LMCache compatibility in the release changelog Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-821.json | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 .lil/changes/vllm-821.json diff --git a/.lil/changes/vllm-821.json b/.lil/changes/vllm-821.json new file mode 100644 index 000000000000..8fde8855c56d --- /dev/null +++ b/.lil/changes/vllm-821.json @@ -0,0 +1,20 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-821", + "category": "feature", + "summary": "Reduce GLM speculative recurrent-state reservations while retaining atomic LMCache restores.", + "models": ["GLM-5.3-Flash"], + "compatibility": "No user action required.", + "details": [ + "Eligible SM12x Model Runner V2 configurations select B12X state recovery automatically for MTP or DFlash with one to seven draft tokens.", + "LMCache checkpoints capture the accepted recurrent state after recovery and convolution compaction; changed prompt tails and generated-response continuation remain reusable.", + "In the qualified TP2 Spark MTP3 LMCache configuration, fixed per-request reservation falls by approximately 188 MiB per GPU; maximum single-request context is unchanged.", + "The existing atomic LMCache adapter remains text-only; multimodal generation and native GPU prefix caching are unchanged.", + "Use --no-use-replayssm for the full-state diagnostic path." + ], + "pull_requests": [821], + "requires": ["b12x-408"], + "evidence": [ + "https://github.com/local-inference-lab/rtx6kpro/blob/master/benchmarks/glm53-kda-recovery-lmcache.md" + ] +} From 7356c7602d8f7230f826dfaa1e3a0c978ca34014 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 16:09:23 +0000 Subject: [PATCH 135/365] fix(cache): export recovered state for logits-only responses A prompt restored at an exact checkpoint can terminate after sampling without a target forward. Preserve checkpoint export for that response, while skipping recovery and convolution commit because no forward records exist. Cover both recovery layouts with and without target computation. Signed-off-by: Martin Vit Assisted-by: OpenAI Codex --- tests/v1/worker/test_gpu_model_runner_v2.py | 19 ++++++++++++++----- vllm/v1/worker/gpu/model_runner.py | 12 ++++++------ 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index d67c5977e56c..be09bb6e3ae3 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -23,7 +23,10 @@ @pytest.mark.parametrize("recovery", [False, True]) -def test_boundary_capture_uses_recovered_accepted_state(monkeypatch, recovery): +@pytest.mark.parametrize("logits_only", [False, True]) +def test_boundary_capture_uses_recovered_accepted_state( + monkeypatch, recovery, logits_only +): """A stop inside a draft window must be applied before checkpoint export.""" events = [] runner = GPUModelRunner.__new__(GPUModelRunner) @@ -38,10 +41,12 @@ def test_boundary_capture_uses_recovered_accepted_state(monkeypatch, recovery): sampled = torch.tensor([4]) def update(*args): - sampled.fill_(2) # The second accepted token terminates the response. + # A restored prompt can terminate after sampling its first token. + sampled.fill_(1 if logits_only else 2) events.append("trim") def commit(*args): + assert not logits_only assert sampled.item() == 2 events.append("commit") @@ -58,10 +63,14 @@ def capture(*args, accepted_state_committed=False): sampled, torch.tensor([0]), boundary_capture=torch.empty(3, 1, 3), + logits_only=logits_only, ) - assert events == ( - ["trim", "commit", "capture"] if recovery else ["trim", "capture", "commit"] - ) + if logits_only: + assert events == ["trim", "capture"] + else: + assert events == ( + ["trim", "commit", "capture"] if recovery else ["trim", "capture", "commit"] + ) def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f7621e174314..ed74fab4db7f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1796,12 +1796,12 @@ def postprocess_sampled( assert self.boundary_checkpoint_state is not None self.boundary_checkpoint_state.capture_mamba(idx_mapping, boundary_capture) - if logits_only: - return - - self.model_state.postprocess_state( - idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu - ) + # A logits-only restore already has committed state and no recovery + # records to apply, but a terminating sample still needs export. + if not logits_only: + self.model_state.postprocess_state( + idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu + ) if boundary_capture is not None and recover_state: assert self.boundary_checkpoint_state is not None self.boundary_checkpoint_state.capture_mamba( From 1579dac2e5a8f2df074530b45d1119c42cfbd760 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 16:28:29 +0000 Subject: [PATCH 136/365] test(cache): follow the logits-only checkpoint reuse contract Logits-only sampling passes no boundary capture request: it reuses the checkpoint already restored and must not compact state or export a duplicate. Preserve the runner's early return and exercise both recovery layouts with the actual caller contract. Signed-off-by: Martin Vit Assisted-by: OpenAI Codex --- tests/v1/worker/test_gpu_model_runner_v2.py | 6 +++--- vllm/v1/worker/gpu/model_runner.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index be09bb6e3ae3..4b1828dc44b0 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -27,7 +27,7 @@ def test_boundary_capture_uses_recovered_accepted_state( monkeypatch, recovery, logits_only ): - """A stop inside a draft window must be applied before checkpoint export.""" + """Trim before export; logits-only restores reuse their existing checkpoint.""" events = [] runner = GPUModelRunner.__new__(GPUModelRunner) runner.is_last_pp_rank = False @@ -62,11 +62,11 @@ def capture(*args, accepted_state_committed=False): torch.tensor([[1, 2, 3, 4]]), sampled, torch.tensor([0]), - boundary_capture=torch.empty(3, 1, 3), + boundary_capture=None if logits_only else torch.empty(3, 1, 3), logits_only=logits_only, ) if logits_only: - assert events == ["trim", "capture"] + assert events == ["trim"] else: assert events == ( ["trim", "commit", "capture"] if recovery else ["trim", "capture", "commit"] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ed74fab4db7f..f7621e174314 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1796,12 +1796,12 @@ def postprocess_sampled( assert self.boundary_checkpoint_state is not None self.boundary_checkpoint_state.capture_mamba(idx_mapping, boundary_capture) - # A logits-only restore already has committed state and no recovery - # records to apply, but a terminating sample still needs export. - if not logits_only: - self.model_state.postprocess_state( - idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu - ) + if logits_only: + return + + self.model_state.postprocess_state( + idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu + ) if boundary_capture is not None and recover_state: assert self.boundary_checkpoint_state is not None self.boundary_checkpoint_state.capture_mamba( From ce1e25b8ed3d12e886ec926c5fbab9912f3e37ac Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 16:46:36 +0000 Subject: [PATCH 137/365] Document the measured throughput cost of GLM state recovery Keep the integration release notes explicit about the TP4 MTP3 memory-throughput tradeoff and link the matched LMCache evidence. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-821.json | 1 + 1 file changed, 1 insertion(+) diff --git a/.lil/changes/vllm-821.json b/.lil/changes/vllm-821.json index 8fde8855c56d..8f4db064f76e 100644 --- a/.lil/changes/vllm-821.json +++ b/.lil/changes/vllm-821.json @@ -9,6 +9,7 @@ "Eligible SM12x Model Runner V2 configurations select B12X state recovery automatically for MTP or DFlash with one to seven draft tokens.", "LMCache checkpoints capture the accepted recurrent state after recovery and convolution compaction; changed prompt tails and generated-response continuation remain reusable.", "In the qualified TP2 Spark MTP3 LMCache configuration, fixed per-request reservation falls by approximately 188 MiB per GPU; maximum single-request context is unchanged.", + "Matched TP4 LMCache measurements show MTP3 C1 output -3.1%, C8 output -1.8% and prefill -0.3%; DFlash2 prefill changes -0.1% with similar decode throughput. This is a memory optimization, not a speedup claim.", "The existing atomic LMCache adapter remains text-only; multimodal generation and native GPU prefix caching are unchanged.", "Use --no-use-replayssm for the full-state diagnostic path." ], From 4d5251b9199104e12bef9aa776e1474415211523 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 16:55:27 +0000 Subject: [PATCH 138/365] Require the reviewed B12X binding fix in GLM recovery releases Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-821.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-821.json b/.lil/changes/vllm-821.json index 8f4db064f76e..e9ec54cbf941 100644 --- a/.lil/changes/vllm-821.json +++ b/.lil/changes/vllm-821.json @@ -14,7 +14,7 @@ "Use --no-use-replayssm for the full-state diagnostic path." ], "pull_requests": [821], - "requires": ["b12x-408"], + "requires": ["b12x-408", "b12x-409"], "evidence": [ "https://github.com/local-inference-lab/rtx6kpro/blob/master/benchmarks/glm53-kda-recovery-lmcache.md" ] From 66bcca4e46264e3ca54fb179d5c5539fb3d6d619 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 17:22:29 +0000 Subject: [PATCH 139/365] Keep full speculative metadata for backends without KDA recovery Narrow state-index rows only for metadata builders that implement accepted-state recovery. Preserve every speculative state column for the generic and ROCm Kimi kernels. Add a reproducer for both full-state builders and document optional output normalization. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/glm5next/test_kda_recoverssm.py | 29 +++++++++++++++++-- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 2 ++ vllm/models/glm5next/nvidia/kda.py | 1 + vllm/models/kimi_k3/nvidia/kda_metadata.py | 1 + vllm/v1/attention/backends/gdn_attn.py | 13 +++++---- 5 files changed, 38 insertions(+), 8 deletions(-) diff --git a/tests/models/glm5next/test_kda_recoverssm.py b/tests/models/glm5next/test_kda_recoverssm.py index eece79d7d235..8231a61d462d 100644 --- a/tests/models/glm5next/test_kda_recoverssm.py +++ b/tests/models/glm5next/test_kda_recoverssm.py @@ -18,6 +18,8 @@ Glm5NextRecoverKDAMetadataBuilder, ) from vllm.models.glm5next.nvidia.model import Glm5NextForCausalLM +from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDAMetadataBuilder +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MambaSpec @@ -152,7 +154,7 @@ def test_glm_recoverssm_reserves_one_recurrent_state_per_draft_window(): ) -def _builder() -> Glm5NextRecoverKDAMetadataBuilder: +def _builder(builder_cls=Glm5NextRecoverKDAMetadataBuilder): config = SimpleNamespace( model_config=SimpleNamespace( hf_text_config=SimpleNamespace(model_type="glm5_next") @@ -186,7 +188,7 @@ def _builder() -> Glm5NextRecoverKDAMetadataBuilder: mamba_cache_mode="none", num_speculative_blocks=0, ) - builder = Glm5NextRecoverKDAMetadataBuilder( + builder = builder_cls( spec, ["layer.0"], config, @@ -196,6 +198,29 @@ def _builder() -> Glm5NextRecoverKDAMetadataBuilder: return builder +@pytest.mark.parametrize( + "builder_cls", [GDNAttentionMetadataBuilder, KimiK3ROCmKDAMetadataBuilder] +) +def test_full_state_backends_keep_every_speculative_state_column(builder_cls): + """A global recovery flag cannot narrow metadata for full-state kernels.""" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]), + 16, + torch.device("cpu"), + ) + builder = _builder(builder_cls) + metadata = builder.build( + 0, + common, + num_decode_draft_tokens_cpu=torch.tensor([2, 2]), + num_accepted_tokens=torch.tensor([3, 2]), + ) + assert builder.spec_state_indices_tensor.shape[1] == 3 + torch.testing.assert_close( + metadata.spec_state_indices_tensor, common.block_table_tensor[:, :3] + ) + + def test_glm_recoverssm_mixed_prefill_and_decode_keeps_one_state_slot(): common = create_common_attn_metadata( BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]), diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index c4a854fe241c..82ac15a8342d 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1355,6 +1355,8 @@ def _run_b12x_kda_decode_post_conv( num_accepted_tokens: Accepted speculative-token counts, or ``None`` for one-token decode requests. num_requests: Number of packed requests. + apply_output_norm: Whether to apply output normalization in the + KDA operation. Defaults to True. Raises: PreparationResourceUnavailableError: If KDA was not prepared for diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index d01e5de69d43..6b78aa40596e 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -84,6 +84,7 @@ class Glm5NextRecoverKDAMetadata(KimiK3KDAMetadata): class Glm5NextRecoverKDAMetadataBuilder(Glm5NextKDAMetadataBuilder): supports_update_block_table = False + supports_kda_state_recovery = True def __init__( self, diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 2c259b2953df..a2695e66eee6 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -313,6 +313,7 @@ def commit_recoverssm_state( class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): supports_update_block_table = False + supports_kda_state_recovery = True mamba_aligned_state_indices: torch.Tensor | None = None def __init__( diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 51bbd2809086..c975ea3af5ee 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -149,6 +149,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] kv_cache_spec: MambaSpec _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH supports_update_block_table: bool = True + supports_kda_state_recovery: bool = False # Runner-owned stable storage, with NULL_BLOCK_ID in padded request rows. mamba_aligned_state_indices: torch.Tensor | None = None @@ -198,9 +199,11 @@ def __init__( else: self.num_spec = 0 self.use_spec_decode: bool = self.num_spec > 0 - self.state_index_columns = ( - 1 if vllm_config.cache_config.use_kda_recoverssm else self.num_spec + 1 + use_kda_state_recovery = ( + self.supports_kda_state_recovery + and vllm_config.cache_config.use_kda_recoverssm ) + self.state_index_columns = 1 if use_kda_state_recovery else self.num_spec + 1 self._b12x_mixed = ( B12xGdnMixedMetadata( max_tokens=vllm_config.scheduler_config.max_num_batched_tokens, @@ -208,8 +211,7 @@ def __init__( state_columns=self.num_spec + 1, device=device, ) - if self.gdn_prefill_backend == "b12x" - and not vllm_config.cache_config.use_kda_recoverssm + if self.gdn_prefill_backend == "b12x" and not use_kda_state_recovery else None ) self._init_reorder_batch_threshold(1, self.use_spec_decode) @@ -269,8 +271,7 @@ def __init__( self._decode_state_indices_source: torch.Tensor | None = None self._decode_state_indices_view: torch.Tensor | None = None self._reuse_spec_decode_inputs = ( - envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH - and not vllm_config.cache_config.use_kda_recoverssm + envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH and not use_kda_state_recovery ) # Constant sources for the uniform spec-decode fast path. They are # copied into the builder-owned graph buffers above, never handed to From f229335c7c2f5cf8251ffd54bf23a6646e3060a9 Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Fri, 18 Sep 2026 12:45:22 +0100 Subject: [PATCH 140/365] fix(qwen): propagate MTP positional overrides Normalize internally constructed MTP drafts before applying target positional overrides, preserving their MTP discriminator for both explicit and automatically detected drafts. Co-authored-by: OMP Agent Signed-off-by: Naadir Jeewa (cherry picked from commit e0ad58ef5873fbd0cf4ed3970c0c8008baf65e03) Signed-off-by: Martin Vit --- .lil/changes/vllm-777.json | 14 +++++ .../test_speculative_draft_hf_overrides.py | 37 +++++++++++ tests/models/qwen4_exp/test_b12x_qsa.py | 22 +++++++ vllm/config/speculative.py | 61 ++++++++++++++++--- 4 files changed, 125 insertions(+), 9 deletions(-) create mode 100644 .lil/changes/vllm-777.json diff --git a/.lil/changes/vllm-777.json b/.lil/changes/vllm-777.json new file mode 100644 index 000000000000..8d4ae1126e50 --- /dev/null +++ b/.lil/changes/vllm-777.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-777", + "category": "fix", + "summary": "Preserve Qwen positional overrides in MTP draft configuration", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "YaRN deployments should verify their configured target and draft context limits before serving.", + "details": [ + "Dictionary-based target Hugging Face overrides reach the internally constructed MTP draft after its model type and architecture are normalized.", + "QSA prefill plan selection covers the configured one-million-token context ceiling." + ], + "pull_requests": [777], + "requires": [] +} diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index bc5300e53fce..b81f1f4e292f 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -22,6 +22,8 @@ from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig +pytestmark = pytest.mark.skip_global_cleanup + def _make_hf_config(**kwargs) -> PretrainedConfig: defaults = dict( @@ -43,6 +45,41 @@ def test_dict_overrides_are_not_forwarded_to_draft(): assert composed is SpeculativeConfig.hf_config_override +@pytest.mark.cpu_test +@pytest.mark.parametrize("method", ["mtp", "draft_model"]) +def test_mtp_draft_receives_target_dict_overrides(method: str): + """An in-model MTP draft shares the target's positional geometry.""" + target_hf_overrides = { + "architectures": ["Qwen3_8FlashNextForCausalLM"], + "model_type": "qwen3_8_flash_next", + "text_config": { + "max_position_embeddings": 1048576, + "rope_parameters": {"rope_type": "yarn", "factor": 4.0}, + }, + } + override = SpeculativeConfig.get_draft_hf_overrides(method, target_hf_overrides) + assert callable(override) + text_config = _make_hf_config( + max_position_embeddings=262144, + rope_parameters={"rope_type": "default"}, + hc_count=4, + mtp_num_hidden_layers=2, + num_attention_heads=4, + ) + source = _make_hf_config( + architectures=["Qwen3_8FlashNextForCausalLM"], + model_type="qwen3_8_flash_next", + text_config=text_config, + ) + + out = override(source) + + assert out.model_type == "qwen4_exp_mtp" + assert out.architectures == ["Qwen4ExpMTP"] + assert out.text_config.max_position_embeddings == 1048576 + assert out.text_config.rope_parameters == {"rope_type": "yarn", "factor": 4.0} + + @pytest.mark.cpu_test def test_none_overrides_fall_back_to_arch_mapping(): composed = SpeculativeConfig.compose_draft_hf_overrides(None) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 09b15e031c94..71dcb7bcd4f5 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/tests/models/qwen4_exp/test_b12x_qsa.py @@ -807,6 +807,7 @@ def test_qsa_selector_fork_requires_one_full_graph( def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None: + """Prefill plan capacities cover each requested context ceiling.""" assert qsa_module._qsa_prefill_context_capacities(262144, 4096) == ( 4096, 8192, @@ -841,6 +842,27 @@ def test_qsa_dcp_group_kv_geometry_tracks_replica_subgroups() -> None: assert qsa_module._qsa_dcp_group_kv_geometry(4, 4, 2) == (2, 2) +def test_qsa_prefill_binding_accepts_pass_one_at_one_million_tokens() -> None: + """Pass one uses a prefill plan at the target context ceiling.""" + owner = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) + owner.max_decode_rows = 16 + owner.max_seq_len = 1048576 + prefill_plan = object() + owner._qsa_prefill_bindings = ( + qsa_module._QSAContextPlan( + max_seq_len=1048576, + caps=None, + plan=prefill_plan, + ), + ) + owner._qsa_decode_context = None + + assert ( + owner._qsa_binding_for_workload(rows=17, max_seq_len=1048576).plan + is prefill_plan + ) + + def test_qsa_run_consumes_projection_views_and_writes_live_output(monkeypatch) -> None: rows = 2 query = torch.randn(rows, 6, 256, dtype=torch.bfloat16) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 591d029d1acd..cb79f8999373 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1088,9 +1088,36 @@ def _apply_composed_hf_override( target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig], hf_config: PretrainedConfig, ) -> PretrainedConfig: + """Apply the draft normalization before a target config transform.""" hf_config = SpeculativeConfig.hf_config_override(hf_config) return target_hf_overrides(hf_config) + @staticmethod + def _apply_mtp_dict_hf_override( + target_hf_overrides: dict[str, Any], + hf_config: PretrainedConfig, + *, + require_mtp: bool = False, + ) -> PretrainedConfig: + """Apply target patches without replacing the MTP discriminator.""" + hf_config = SpeculativeConfig.hf_config_override(hf_config) + if require_mtp and hf_config.model_type not in get_args(MTPModelTypes): + return hf_config + for key, value in target_hf_overrides.items(): + if key in {"architectures", "model_type"}: + continue + target = getattr(hf_config, key, None) + if isinstance(value, dict) and target is not None: + if isinstance(target, dict): + target.update(value) + continue + if hasattr(target, "__dict__"): + for nested_key, nested_value in value.items(): + setattr(target, nested_key, nested_value) + continue + setattr(hf_config, key, value) + return hf_config + @staticmethod def compose_draft_hf_overrides( target_hf_overrides: HfOverrides | None, @@ -1117,6 +1144,27 @@ def compose_draft_hf_overrides( SpeculativeConfig._apply_composed_hf_override, target_hf_overrides ) + @staticmethod + def get_draft_hf_overrides( + method: str, + target_hf_overrides: HfOverrides | None, + ) -> HfOverrides: + """Return overrides appropriate for the selected draft method.""" + if method == "medusa": + return {"model_type": "medusa"} + if method == "mtp" and isinstance(target_hf_overrides, dict): + return functools.partial( + SpeculativeConfig._apply_mtp_dict_hf_override, + target_hf_overrides, + ) + if method == "draft_model" and isinstance(target_hf_overrides, dict): + return functools.partial( + SpeculativeConfig._apply_mtp_dict_hf_override, + target_hf_overrides, + require_mtp=True, + ) + return SpeculativeConfig.compose_draft_hf_overrides(target_hf_overrides) + @staticmethod def _is_custom_proposer_path(model: str | None) -> bool: """True if ``model`` is a dotted import path (e.g. ``pkg.MyProposer``).""" @@ -1130,6 +1178,7 @@ def _is_custom_proposer_path(model: str | None) -> bool: return len(parts) >= 2 and all(part.isidentifier() for part in parts) def __post_init__(self): + """Validate speculative settings and construct the draft configuration.""" # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter # will be used to set the draft model, eagle head, or additional weight @@ -1318,15 +1367,9 @@ def __post_init__(self): # detect them. When the method is explicitly "medusa", inject # model_type so MedusaConfig.from_pretrained is used instead. draft_hf_overrides: HfOverrides - if self.method == "medusa": - draft_hf_overrides = {"model_type": "medusa"} - else: - # Compose any callable hf_overrides set on the target so the - # draft config receives the same transform (e.g. the test - # shrink). Dict overrides stay target-only. - draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides( - self.target_model_config.hf_overrides - ) + draft_hf_overrides = SpeculativeConfig.get_draft_hf_overrides( + self.method, self.target_model_config.hf_overrides + ) self.draft_model_config = ModelConfig( model=self.model, runner="draft", From 6e8a9d246958b5c90bd8dc76e057e28ecea189a9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 19:15:04 +0000 Subject: [PATCH 141/365] docs(release): describe Qwen MTP override propagation Limit the vLLM PR777 release fragment to its runtime change. The QSA context-capacity check is test coverage for existing behavior. Signed-off-by: Martin Vit --- .lil/changes/vllm-777.json | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.lil/changes/vllm-777.json b/.lil/changes/vllm-777.json index 8d4ae1126e50..b28ed461166d 100644 --- a/.lil/changes/vllm-777.json +++ b/.lil/changes/vllm-777.json @@ -6,8 +6,7 @@ "models": ["Qwen3.8-Flash-Next"], "compatibility": "YaRN deployments should verify their configured target and draft context limits before serving.", "details": [ - "Dictionary-based target Hugging Face overrides reach the internally constructed MTP draft after its model type and architecture are normalized.", - "QSA prefill plan selection covers the configured one-million-token context ceiling." + "Dictionary-based target Hugging Face overrides reach the internally constructed MTP draft after its model type and architecture are normalized." ], "pull_requests": [777], "requires": [] From 56a456a56d9c94782e6832e4812630ac768e4db7 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Wed, 2 Sep 2026 16:31:45 -0700 Subject: [PATCH 142/365] [Bugfix][KV Offload] SimpleCPUOffload: skip non-prefix-cacheable groups Scratch KV cache groups such as the QSA key ring are not hash-addressable and cannot be offloaded. Exclude them from lazy target sizing and from scheduler store/load transfer selection using KVCacheConfig.prefix_cacheable_group_ids. Add a QSA-shaped scheduler regression covering eager store/load and lazy target sizing. Co-authored-by: Claude Fable 5.1 Signed-off-by: Zhewen Li (cherry picked from commit 3cc857e33d3c693ba4beef2eb8a453a99d969877) (cherry picked from commit a598d29dabc857f0feb74d40fcc88a9951392ac1) --- tests/v1/simple_kv_offload/test_scheduler.py | 120 +++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 10 ++ 2 files changed, 130 insertions(+) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 901e076e0950..4678bf7498d9 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -41,6 +41,7 @@ register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -2843,3 +2844,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True + + +QSA_RING_BLOCK_SIZE = 4 + + +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: + register_all_kvcache_specs(vllm_config=None) + full_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring_spec = CircularBufferSpec( + block_size=QSA_RING_BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + ) + groups = [ + KVCacheGroupSpec(["layer_0"], full_spec), + KVCacheGroupSpec(["qsa_ring"], ring_spec), + ] + tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, + ) + for group in groups + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _make_qsa_scheduler( + lazy: bool = False, +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + sched = SimpleCPUOffloadScheduler( + vllm_config=_make_vllm_config(), + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + lazy_offload=lazy, + ) + gpu_pool = BlockPool( + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE + ) + sched.bind_gpu_block_pool(gpu_pool) + return sched, gpu_pool + + +def test_qsa_ring_group_is_never_stored_or_loaded() -> None: + sched, gpu_pool = _make_qsa_scheduler() + assert sched.prefix_cacheable_group_ids == {0} + + num_blocks = 2 + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + ring_block = gpu_pool.get_new_blocks(1)[0] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} + assert ring_block.block_id not in meta.store_gpu_blocks + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-qsa-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + ring_block2 = gpu_pool.get_new_blocks(1)[0] + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} + assert ring_block2.block_id not in meta2.load_gpu_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + without_ring = KVCacheConfig( + num_blocks=16, + kv_cache_tensors=with_ring.kv_cache_tensors[:1], + kv_cache_groups=with_ring.kv_cache_groups[:1], + ) + max_batched = 64 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + with_ring, max_batched + ) + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + without_ring, max_batched + ) + assert target_with == target_without diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index bc7e6c3a13b4..825911fc6a46 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -165,6 +165,9 @@ def __init__( kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks + self.prefix_cacheable_group_ids = frozenset( + kv_cache_config.prefix_cacheable_group_ids + ) self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU # Find the full attention kv group for prefix cache matching. self.fa_gidx = -1 @@ -299,6 +302,8 @@ def _estimate_lazy_target_blocks( target = 0 for g in kv_cache_config.kv_cache_groups: spec = g.kv_cache_spec + if not spec.prefix_cacheable: + continue # Only full attention is sharded across DCP ranks; replicated specs # (mamba, sliding window, chunked-local) keep their own block size. block_size = spec.block_size * dcp_world_size_for_kv_cache_spec( @@ -443,6 +448,9 @@ def update_state_after_alloc( # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): + if g not in self.prefix_cacheable_group_ids: + cpu_hit_blocks.append([]) + continue g_block_size = self.group_block_sizes[g] n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -831,6 +839,8 @@ def _select_eager_blocks_to_store( for g, group_gpu_ids in enumerate(block_ids_by_group): if len(gpu_block_ids) >= num_free: break + if g not in self.prefix_cacheable_group_ids: + continue group_manager = self.cpu_coordinator.single_type_managers[g] if not group_manager.has_positionally_stable_blocks: continue From f6485322f97357a21f6e92bd5f2be5f9b8555e3e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 19:24:54 +0000 Subject: [PATCH 143/365] docs(release): attribute QSA ring offload import Identify the upstream SimpleCPUOffloadConnector source and its scope in the integration container changelog. Signed-off-by: Martin Vit --- .lil/changes/vllm-upstream-55033.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-upstream-55033.json diff --git a/.lil/changes/vllm-upstream-55033.json b/.lil/changes/vllm-upstream-55033.json new file mode 100644 index 000000000000..9000cc06c9d8 --- /dev/null +++ b/.lil/changes/vllm-upstream-55033.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-upstream-55033", + "category": "fix", + "summary": "Exclude Qwen QSA scratch cache from Simple CPU offload", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "Applies only when SimpleCPUOffloadConnector is enabled. The regular OffloadingConnector remains unqualified for this model.", + "details": [ + "The optional Simple CPU KV connector excludes non-prefix-cacheable QSA ring blocks from lazy target sizing, external load selection, and eager stores.", + "The change preserves the original author and code from vllm-project/vllm#55033." + ], + "pull_requests": [], + "authors": ["Zhewen Li"], + "evidence": ["https://github.com/vllm-project/vllm/pull/55033"], + "requires": [] +} From 4d7f9023830a8dc9db2ed81196b843c3de3c9afa Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 20:01:29 +0000 Subject: [PATCH 144/365] fix(glm): keep long profile prefills out of KDA decode Signed-off-by: Martin Vit (cherry picked from commit 65f73469a324f54f43b305590081dc3fd83daf0e) --- .../v1/attention/test_gdn_metadata_builder.py | 39 +++++++++++++++++++ .../attention/backends/b12x_gdn_metadata.py | 6 ++- vllm/v1/attention/backends/gdn_attn.py | 10 ++++- 3 files changed, 53 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 2967f2f05e6d..20fea68b4c85 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -16,6 +16,7 @@ ) from vllm.config import SpeculativeConfig from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.b12x_gdn_metadata import B12xGdnMixedMetadata from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionMetadata, GDNAttentionMetadataBuilder, @@ -174,6 +175,26 @@ class GDNBuildTestCase: expected_num_prefill_tokens=0, expected_num_spec_decodes=0, ), + "long_prefill_with_zero_draft_profile_marker": GDNBuildTestCase( + seq_lens=[4112], + query_lens=[4096], + num_decode_draft_tokens=[0], + num_speculative_tokens=3, + expected_num_decodes=0, + expected_num_prefills=1, + expected_num_prefill_tokens=4096, + expected_num_spec_decodes=0, + ), + "single_token_zero_draft_decode": GDNBuildTestCase( + seq_lens=[33], + query_lens=[1], + num_decode_draft_tokens=[0], + num_speculative_tokens=3, + expected_num_decodes=1, + expected_num_prefills=0, + expected_num_prefill_tokens=0, + expected_num_spec_decodes=0, + ), # Multi-token prefill alongside spec decode — no decode to reclassify "spec_decode_with_real_prefill": GDNBuildTestCase( seq_lens=[100, 20], @@ -442,6 +463,24 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase): assert meta.num_spec_decodes == test_case.expected_num_spec_decodes +def test_b12x_mixed_metadata_keeps_zero_draft_profile_in_prefill() -> None: + builder = _create_gdn_builder(num_speculative_tokens=3) + builder._b12x_mixed = B12xGdnMixedMetadata( + max_tokens=4096, max_seqs=1, state_columns=4, device=DEVICE + ) + metadata = _build( + builder, + BatchSpec(seq_lens=[4112], query_lens=[4096]), + num_decode_draft_tokens=[0], + ) + + assert metadata.num_prefills == 1 + assert metadata.num_spec_decodes == 0 + assert metadata.b12x_mixed is not None + assert metadata.b12x_mixed._num_non_spec == 1 + assert metadata.b12x_mixed._num_spec == 0 + + def test_fresh_single_token_prompt_uses_prefill_state_initialization() -> None: builder = _create_gdn_builder() batch = BatchSpec(seq_lens=[1], query_lens=[1]) diff --git a/vllm/v1/attention/backends/b12x_gdn_metadata.py b/vllm/v1/attention/backends/b12x_gdn_metadata.py index 7f2db234446b..6642167c5e45 100644 --- a/vllm/v1/attention/backends/b12x_gdn_metadata.py +++ b/vllm/v1/attention/backends/b12x_gdn_metadata.py @@ -100,7 +100,11 @@ def stage( for row, (start, end) in enumerate(zip(starts, starts[1:])): if end == start: continue - if drafts is not None and drafts[row] >= 0 and end - start > 1: + if ( + drafts is not None + and drafts[row] >= 0 + and end - start == drafts[row] + 1 + ): if end - start > self.state_columns: raise ValueError("GDN verification exceeds planned state columns") spec_rows.append(row) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 133983779240..491449002e53 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -475,7 +475,15 @@ def build( # type: ignore[override] spec_sequence_masks = None num_spec_decodes = 0 else: - spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + # A speculative row contains exactly one target token followed by + # its draft tokens. Profiling may provide a zero-draft marker for + # a long prefill, which must not enter the bounded decode kernel. + query_lens_cpu = query_start_loc_cpu.diff() + spec_sequence_masks_cpu = ( + (num_decode_draft_tokens_cpu >= 0) + & (num_decode_draft_tokens_cpu <= self.num_spec) + & (query_lens_cpu == num_decode_draft_tokens_cpu + 1) + ) num_spec_decodes = spec_sequence_masks_cpu.sum().item() # A zero-draft varlen batch still recovers the previous step's # accepted recurrent state before consuming its bonus tokens. From e36edd0ee43e839fc5cebae9dbd7ce05a3210b73 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 20:02:36 +0000 Subject: [PATCH 145/365] docs(release): record GLM DCP profile classification fix Signed-off-by: Martin Vit --- .lil/changes/vllm-822.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-822.json diff --git a/.lil/changes/vllm-822.json b/.lil/changes/vllm-822.json new file mode 100644 index 000000000000..45d911b06b28 --- /dev/null +++ b/.lil/changes/vllm-822.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-822", + "category": "fix", + "summary": "Keep long GLM DCP profile prefills out of bounded KDA decode", + "models": ["GLM-5.3-Flash"], + "compatibility": "No configuration change is required. Valid MTP decode rows retain the B12X KDA decode path.", + "details": [ + "A row enters speculative KDA decode only when its length is one target token plus its declared draft count.", + "This prevents a zero-draft marker on a long memory-profile prefill from exceeding the decode kernel capacity during TP4/DCP4 startup." + ], + "pull_requests": [822], + "authors": ["Local Inference Lab"], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/822"], + "requires": [] +} From 3b4091131aaf14e0670a0b5c3d5665403f58b4dc Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 20:01:29 +0000 Subject: [PATCH 146/365] fix(glm): keep long profile prefills out of KDA decode Signed-off-by: Martin Vit (cherry picked from commit 65f73469a324f54f43b305590081dc3fd83daf0e) (cherry picked from commit 4d7f9023830a8dc9db2ed81196b843c3de3c9afa) --- .../v1/attention/test_gdn_metadata_builder.py | 39 +++++++++++++++++++ .../attention/backends/b12x_gdn_metadata.py | 6 ++- vllm/v1/attention/backends/gdn_attn.py | 10 ++++- 3 files changed, 53 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 2967f2f05e6d..20fea68b4c85 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -16,6 +16,7 @@ ) from vllm.config import SpeculativeConfig from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.b12x_gdn_metadata import B12xGdnMixedMetadata from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionMetadata, GDNAttentionMetadataBuilder, @@ -174,6 +175,26 @@ class GDNBuildTestCase: expected_num_prefill_tokens=0, expected_num_spec_decodes=0, ), + "long_prefill_with_zero_draft_profile_marker": GDNBuildTestCase( + seq_lens=[4112], + query_lens=[4096], + num_decode_draft_tokens=[0], + num_speculative_tokens=3, + expected_num_decodes=0, + expected_num_prefills=1, + expected_num_prefill_tokens=4096, + expected_num_spec_decodes=0, + ), + "single_token_zero_draft_decode": GDNBuildTestCase( + seq_lens=[33], + query_lens=[1], + num_decode_draft_tokens=[0], + num_speculative_tokens=3, + expected_num_decodes=1, + expected_num_prefills=0, + expected_num_prefill_tokens=0, + expected_num_spec_decodes=0, + ), # Multi-token prefill alongside spec decode — no decode to reclassify "spec_decode_with_real_prefill": GDNBuildTestCase( seq_lens=[100, 20], @@ -442,6 +463,24 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase): assert meta.num_spec_decodes == test_case.expected_num_spec_decodes +def test_b12x_mixed_metadata_keeps_zero_draft_profile_in_prefill() -> None: + builder = _create_gdn_builder(num_speculative_tokens=3) + builder._b12x_mixed = B12xGdnMixedMetadata( + max_tokens=4096, max_seqs=1, state_columns=4, device=DEVICE + ) + metadata = _build( + builder, + BatchSpec(seq_lens=[4112], query_lens=[4096]), + num_decode_draft_tokens=[0], + ) + + assert metadata.num_prefills == 1 + assert metadata.num_spec_decodes == 0 + assert metadata.b12x_mixed is not None + assert metadata.b12x_mixed._num_non_spec == 1 + assert metadata.b12x_mixed._num_spec == 0 + + def test_fresh_single_token_prompt_uses_prefill_state_initialization() -> None: builder = _create_gdn_builder() batch = BatchSpec(seq_lens=[1], query_lens=[1]) diff --git a/vllm/v1/attention/backends/b12x_gdn_metadata.py b/vllm/v1/attention/backends/b12x_gdn_metadata.py index 7f2db234446b..6642167c5e45 100644 --- a/vllm/v1/attention/backends/b12x_gdn_metadata.py +++ b/vllm/v1/attention/backends/b12x_gdn_metadata.py @@ -100,7 +100,11 @@ def stage( for row, (start, end) in enumerate(zip(starts, starts[1:])): if end == start: continue - if drafts is not None and drafts[row] >= 0 and end - start > 1: + if ( + drafts is not None + and drafts[row] >= 0 + and end - start == drafts[row] + 1 + ): if end - start > self.state_columns: raise ValueError("GDN verification exceeds planned state columns") spec_rows.append(row) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index c975ea3af5ee..2fa597abe42f 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -483,7 +483,15 @@ def build( # type: ignore[override] spec_sequence_masks = None num_spec_decodes = 0 else: - spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + # A speculative row contains exactly one target token followed by + # its draft tokens. Profiling may provide a zero-draft marker for + # a long prefill, which must not enter the bounded decode kernel. + query_lens_cpu = query_start_loc_cpu.diff() + spec_sequence_masks_cpu = ( + (num_decode_draft_tokens_cpu >= 0) + & (num_decode_draft_tokens_cpu <= self.num_spec) + & (query_lens_cpu == num_decode_draft_tokens_cpu + 1) + ) num_spec_decodes = spec_sequence_masks_cpu.sum().item() # A zero-draft varlen batch still recovers the previous step's # accepted recurrent state before consuming its bonus tokens. From 02e5a8d7bb1a7f757fa90b0b9df7b36fe09906df Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 20:02:36 +0000 Subject: [PATCH 147/365] docs(release): record GLM DCP profile classification fix Signed-off-by: Martin Vit (cherry picked from commit e36edd0ee43e839fc5cebae9dbd7ce05a3210b73) --- .lil/changes/vllm-822.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-822.json diff --git a/.lil/changes/vllm-822.json b/.lil/changes/vllm-822.json new file mode 100644 index 000000000000..45d911b06b28 --- /dev/null +++ b/.lil/changes/vllm-822.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-822", + "category": "fix", + "summary": "Keep long GLM DCP profile prefills out of bounded KDA decode", + "models": ["GLM-5.3-Flash"], + "compatibility": "No configuration change is required. Valid MTP decode rows retain the B12X KDA decode path.", + "details": [ + "A row enters speculative KDA decode only when its length is one target token plus its declared draft count.", + "This prevents a zero-draft marker on a long memory-profile prefill from exceeding the decode kernel capacity during TP4/DCP4 startup." + ], + "pull_requests": [822], + "authors": ["Local Inference Lab"], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/822"], + "requires": [] +} From 0aa1e083399f8e622827e651fac75fb9cf8ef347 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 21:07:58 +0000 Subject: [PATCH 148/365] fix(kda): keep decode output disjoint from live gate Signed-off-by: Martin Vit --- tests/models/test_glm5next_model.py | 70 +++++++++++++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 53 +++++++++----- 2 files changed, 105 insertions(+), 18 deletions(-) diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 445867df62dc..3a51de380c11 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -2399,6 +2399,76 @@ def make_layer(): assert layers[1]._b12x_kda_num_tokens.item() == 0 +def test_b12x_kda_decode_keeps_overlapping_gate_read_only(monkeypatch) -> None: + gate = torch.arange(2, dtype=torch.float32).view(2, 1, 1) + original_gate = gate.clone() + bound_outputs: list[torch.Tensor] = [] + + class FakeApi: + @staticmethod + def bind_kda(plan, **kwargs): + output = kwargs["output"] + bound_outputs.append(output) + if output.untyped_storage().data_ptr() == gate.untyped_storage().data_ptr(): + raise ValueError( + "mutable buffer output must not overlap read-only tensor raw_g" + ) + return SimpleNamespace(**kwargs) + + @staticmethod + def run_kda(binding, **kwargs): + torch.testing.assert_close(binding.raw_g, original_gate) + binding.output.fill_(7) + + monkeypatch.setattr( + kimi_gdn_linear_attn, + "get_forward_context", + lambda: SimpleNamespace(additional_kwargs={}), + ) + layer = KimiGatedDeltaNetAttention.__new__(KimiGatedDeltaNetAttention) + torch.nn.Module.__init__(layer) + layer._b12x_kda_api = FakeApi() + layer._b12x_kda_plan = SimpleNamespace( + caps=SimpleNamespace(max_state_slots=2), + scratch_specs=lambda: (), + ) + layer._b12x_kda_num_accepted_tokens = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_num_seqs = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_num_tokens = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_max_tokens = 2 + layer._b12x_kda_max_seqs = 1 + layer._b12x_kda_state_index_columns = 2 + layer.gate_lower_bound = -5.0 + layer.local_num_heads = 1 + layer.head_dim = 1 + layer.A_log = torch.ones(1) + layer.dt_bias = torch.ones(1) + layer.o_norm = SimpleNamespace(eps=1e-6, weight=torch.ones(1)) + layer.kv_cache = [None, torch.empty(2, 1, 1, 1)] + monkeypatch.setattr(layer, "_get_b12x_kda_workspace", lambda: torch.empty(1)) + + layer._run_b12x_kda_decode_post_conv( + metadata=SimpleNamespace(is_uniform_spec_decode=False), + mixed_qkv=torch.zeros(2, 3), + raw_g=gate, + raw_beta=torch.zeros(2, 1), + z=torch.zeros(2, 1, 1), + output=gate, + state_indices=torch.zeros(1, 2, dtype=torch.int32), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32), + num_accepted_tokens=None, + num_requests=1, + ) + + assert len(bound_outputs) == 2 + assert bound_outputs[0] is gate + assert ( + bound_outputs[1].untyped_storage().data_ptr() + != gate.untyped_storage().data_ptr() + ) + assert torch.equal(gate, torch.full_like(gate, 7)) + + @pytest.mark.parametrize("is_mtp_layer", [False, True]) def test_glm5next_sparse_mla_selects_b12x_backend( monkeypatch, is_mtp_layer: bool diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index b493d19baaf0..3cd031342aa7 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1356,30 +1356,47 @@ def _run_b12x_kda_decode_post_conv( num_tokens_tensor, ) = bound_metadata - binding = api.bind_kda( - plan, - scratch=scratch, - mixed_qkv=mixed_qkv, - raw_g=raw_g, - raw_beta=raw_beta, - z=z, - A_log=self.A_log, - dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), - norm_weight=self.o_norm.weight, - recurrent_state=self.kv_cache[1], - query_start_loc=query_start_loc, - num_accepted_tokens=accepted_tokens, - state_indices=state_indices[:num_requests, :state_columns], - num_seqs=num_seqs, - num_tokens=num_tokens_tensor, - output=output, - ) + def bind(destination: torch.Tensor) -> Any: + return api.bind_kda( + plan, + scratch=scratch, + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + z=z, + A_log=self.A_log, + dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), + norm_weight=self.o_norm.weight, + recurrent_state=self.kv_cache[1], + query_start_loc=query_start_loc, + num_accepted_tokens=accepted_tokens, + state_indices=state_indices[:num_requests, :state_columns], + num_seqs=num_seqs, + num_tokens=num_tokens_tensor, + output=destination, + ) + + destination = output + try: + binding = bind(destination) + except ValueError as error: + if str(error) != ( + "mutable buffer output must not overlap read-only tensor raw_g" + ): + raise + # An overlapping destination cannot be used while B12X reads the + # gate. Copy back only after recurrence has consumed that input; + # disjoint destinations keep the allocation-free path. + destination = torch.empty_like(output) + binding = bind(destination) api.run_kda( binding, lower_bound=self.gate_lower_bound, eps=self.o_norm.eps, scale=self.head_dim**-0.5, ) + if destination is not output: + output.copy_(destination) def forward( self, From 64b752317aecc2199d030a730f7a3db0bd147684 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 21:24:45 +0000 Subject: [PATCH 149/365] fix(kda): keep decode output disjoint from live gate Signed-off-by: Martin Vit --- .lil/changes/vllm-824.json | 16 +++++ tests/models/test_glm5next_model.py | 72 +++++++++++++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 55 +++++++++----- 3 files changed, 124 insertions(+), 19 deletions(-) create mode 100644 .lil/changes/vllm-824.json diff --git a/.lil/changes/vllm-824.json b/.lil/changes/vllm-824.json new file mode 100644 index 000000000000..a4f179084684 --- /dev/null +++ b/.lil/changes/vllm-824.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-824", + "category": "fix", + "summary": "Preserve the GLM KDA gate when a decode output aliases it", + "models": ["GLM-5.3-Flash"], + "compatibility": "No configuration change is required. Non-overlapping KDA outputs retain the allocation-free path.", + "details": [ + "When B12X rejects an output overlapping the read-only gate, vLLM binds a disjoint caller-owned output and copies the result back after recurrence.", + "Other B12X binding errors still propagate without a fallback." + ], + "pull_requests": [824], + "authors": ["Local Inference Lab"], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/824"], + "requires": [] +} diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index de03a5fa3b78..4e01f64490e0 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -2416,6 +2416,7 @@ def make_layer(): layer.dt_bias = torch.ones(1) layer.o_norm = SimpleNamespace(eps=1e-6, weight=torch.ones(1)) layer.kv_cache = [None, torch.empty(32, 1, 1, 1)] + layer.cache_config = SimpleNamespace(use_kda_recoverssm=False) return layer layers = [make_layer(), make_layer()] @@ -2475,6 +2476,77 @@ def make_layer(): assert layers[1]._b12x_kda_num_tokens.item() == 0 +def test_b12x_kda_decode_keeps_overlapping_gate_read_only(monkeypatch) -> None: + gate = torch.arange(2, dtype=torch.float32).view(2, 1, 1) + original_gate = gate.clone() + bound_outputs: list[torch.Tensor] = [] + + class FakeApi: + @staticmethod + def bind_kda(plan, **kwargs): + output = kwargs["output"] + bound_outputs.append(output) + if output.untyped_storage().data_ptr() == gate.untyped_storage().data_ptr(): + raise ValueError( + "mutable buffer output must not overlap read-only tensor raw_g" + ) + return SimpleNamespace(**kwargs) + + @staticmethod + def run_kda(binding, **kwargs): + torch.testing.assert_close(binding.raw_g, original_gate) + binding.output.fill_(7) + + monkeypatch.setattr( + kimi_gdn_linear_attn, + "get_forward_context", + lambda: SimpleNamespace(additional_kwargs={}), + ) + layer = KimiGatedDeltaNetAttention.__new__(KimiGatedDeltaNetAttention) + torch.nn.Module.__init__(layer) + layer._b12x_kda_api = FakeApi() + layer._b12x_kda_plan = SimpleNamespace( + caps=SimpleNamespace(max_state_slots=2), + scratch_specs=lambda: (), + ) + layer._b12x_kda_num_accepted_tokens = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_num_seqs = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_num_tokens = torch.zeros(1, dtype=torch.int32) + layer._b12x_kda_max_tokens = 2 + layer._b12x_kda_max_seqs = 1 + layer._b12x_kda_state_index_columns = 2 + layer.gate_lower_bound = -5.0 + layer.local_num_heads = 1 + layer.head_dim = 1 + layer.A_log = torch.ones(1) + layer.dt_bias = torch.ones(1) + layer.o_norm = SimpleNamespace(eps=1e-6, weight=torch.ones(1)) + layer.kv_cache = [None, torch.empty(2, 1, 1, 1)] + layer.cache_config = SimpleNamespace(use_kda_recoverssm=False) + monkeypatch.setattr(layer, "_get_b12x_kda_workspace", lambda: torch.empty(1)) + + layer._run_b12x_kda_decode_post_conv( + metadata=SimpleNamespace(is_uniform_spec_decode=False), + mixed_qkv=torch.zeros(2, 3), + raw_g=gate, + raw_beta=torch.zeros(2, 1), + z=torch.zeros(2, 1, 1), + output=gate, + state_indices=torch.zeros(1, 2, dtype=torch.int32), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32), + num_accepted_tokens=None, + num_requests=1, + ) + + assert len(bound_outputs) == 2 + assert bound_outputs[0] is gate + assert ( + bound_outputs[1].untyped_storage().data_ptr() + != gate.untyped_storage().data_ptr() + ) + assert torch.equal(gate, torch.full_like(gate, 7)) + + @pytest.mark.parametrize("is_mtp_layer", [False, True]) def test_glm5next_sparse_mla_selects_b12x_backend( monkeypatch, is_mtp_layer: bool diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index ad59b1ec1d6a..308318149620 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1381,25 +1381,40 @@ def _run_b12x_kda_decode_post_conv( num_tokens_tensor, ) = bound_metadata - binding = api.bind_kda( - plan, - scratch=scratch, - mixed_qkv=mixed_qkv, - raw_g=raw_g, - raw_beta=raw_beta, - z=z, - A_log=self.A_log, - dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), - norm_weight=self.o_norm.weight, - recurrent_state=self.kv_cache[1], - query_start_loc=query_start_loc, - num_accepted_tokens=accepted_tokens, - state_indices=state_indices[:num_requests, :state_columns], - num_seqs=num_seqs, - num_tokens=num_tokens_tensor, - output=output, - **self._b12x_recovery_records(), - ) + def bind(destination: torch.Tensor) -> Any: + return api.bind_kda( + plan, + scratch=scratch, + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + z=z, + A_log=self.A_log, + dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), + norm_weight=self.o_norm.weight, + recurrent_state=self.kv_cache[1], + query_start_loc=query_start_loc, + num_accepted_tokens=accepted_tokens, + state_indices=state_indices[:num_requests, :state_columns], + num_seqs=num_seqs, + num_tokens=num_tokens_tensor, + output=destination, + **self._b12x_recovery_records(), + ) + + destination = output + try: + binding = bind(destination) + except ValueError as error: + if str(error) != ( + "mutable buffer output must not overlap read-only tensor raw_g" + ): + raise + # An overlapping destination cannot be used while B12X reads the + # gate. Copy back only after recurrence has consumed that input; + # disjoint destinations keep the allocation-free path. + destination = torch.empty_like(output) + binding = bind(destination) api.run_kda( binding, lower_bound=self.gate_lower_bound, @@ -1407,6 +1422,8 @@ def _run_b12x_kda_decode_post_conv( scale=self.head_dim**-0.5, apply_output_norm=apply_output_norm, ) + if destination is not output: + output.copy_(destination) def _b12x_recovery_records(self) -> dict[str, torch.Tensor]: if not self.cache_config.use_kda_recoverssm: From a64214f6876e5f9b93f5c4263772a8d179b1c356 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:22:53 +0000 Subject: [PATCH 150/365] fix(router): preserve SM120 graph guard through fullgraph compilation --- .../kernels/test_gate_linear_rocm_dispatch.py | 7 +- .../kernels/test_gate_linear_sm120_compile.py | 64 +++++++++++++++++++ .../layers/fused_moe/router/gate_linear.py | 54 +++++++++++----- 3 files changed, 107 insertions(+), 18 deletions(-) create mode 100644 tests/kernels/test_gate_linear_sm120_compile.py diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index 1041459c676b..de002e020987 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -192,7 +192,7 @@ def test_sm120_enables_bf16_fp32_paths_without_datacenter_kernels(monkeypatch): assert gate.allow_ll_bf16_gemm assert not gate.allow_specialized_router_gemm - assert not gate.allow_dsv3_router_gemm + assert not gate.allow_fp32_router_gemm assert gate.allow_cublas_router_gemm @@ -276,6 +276,11 @@ def new_empty(self, shape): lambda: object(), ) monkeypatch.setattr(torch, "mm", lambda *args, **kwargs: expected) + monkeypatch.setattr( + torch.ops.vllm, + "sm120_cublas_router_gemm", + gate_linear_mod.sm120_cublas_router_gemm_impl, + ) output, output_bias = gate(x) diff --git a/tests/kernels/test_gate_linear_sm120_compile.py b/tests/kernels/test_gate_linear_sm120_compile.py new file mode 100644 index 000000000000..4dca4f808ad8 --- /dev/null +++ b/tests/kernels/test_gate_linear_sm120_compile.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SM120 router compilation and graph replay preserve FP32 logits.""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_device_capability((12, 0)), reason="requires SM120" +) + + +def _gate(monkeypatch): + for name in ( + "vllm.model_executor.layers.linear", + "vllm.model_executor.parameter", + ): + monkeypatch.setattr(f"{name}.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(f"{name}.get_tensor_model_parallel_world_size", lambda: 1) + gate = GateLinear( + 4096, 256, params_dtype=torch.bfloat16, out_dtype=torch.float32 + ).cuda() + with torch.no_grad(): + gate.weight.normal_(std=0.01) + assert gate._sm120_graph_pool_lifetime_guard + return gate + + +@pytest.mark.parametrize("compiled", [False, True]) +@torch.inference_mode() +def test_sm120_router_fullgraph_and_shared_pool_replay(monkeypatch, compiled): + assert torch.cuda.get_device_capability() == (12, 0) + torch.manual_seed(41) + gate = _gate(monkeypatch) + forward = ( + torch.compile(gate, fullgraph=True, backend="eager", dynamic=True) + if compiled + else gate + ) + pool = torch.cuda.graph_pool_handle() + captures = [] + for rows in (32, 64): + x = torch.randn(rows, 4096, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + result, bias = forward(x) + torch.cuda.current_stream().wait_stream(stream) + assert result.dtype == torch.float32 and bias is None + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + output, _ = forward(x) + captures.append((graph, x, output)) + # Replay in capture order while both differently sized graphs remain alive. + for _ in range(3): + for graph, x, output in captures: + x.normal_() + graph.replay() + reference = torch.nn.functional.linear(x.float(), gate.weight.float()) + torch.testing.assert_close(output, reference, atol=0.002, rtol=0.0001) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index b2ce28e38df0..5cf5787f18bf 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -208,26 +208,14 @@ def forward( # Tier 4: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: - graph_pool_guard = None if self._sm120_graph_pool_lifetime_guard: - from vllm.compilation.breakable_cudagraph import ( - BreakableCUDAGraphCapture, + # Keep the stream-capture query and allocation at runtime. + # Dynamo cannot trace the CUDA query's non-Tensor result. + return ( + torch.ops.vllm.sm120_cublas_router_gemm(x, self.weight), + None, ) - - if ( - BreakableCUDAGraphCapture.current() is not None - or torch.cuda.is_current_stream_capturing() - ): - # The former linear-plus-cast path reserved a BF16 router - # output before its FP32 result. Preserve that graph-pool - # address layout while using the fused FP32 cuBLAS output; - # auxiliary-stream graph nodes retain addresses across the - # differently sized captures that share this pool. - graph_pool_guard = x.new_empty( - (*x.shape[:-1], self.weight.shape[0]) - ) output = torch.mm(x, self.weight.T, out_dtype=torch.float32) - del graph_pool_guard return output, None # Tier 5: F.linear (ReplicatedLinear) @@ -242,6 +230,38 @@ def forward( _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS +def sm120_cublas_router_gemm_impl( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + graph_pool_guard = None + if ( + BreakableCUDAGraphCapture.current() is not None + or torch.cuda.is_current_stream_capturing() + ): + # Preserve the former BF16-output allocation before the fused FP32 + # result. Auxiliary-stream graph nodes retain these addresses across + # differently sized captures that share a graph memory pool. + graph_pool_guard = x.new_empty((*x.shape[:-1], weight.shape[0])) + output = torch.mm(x, weight.T, out_dtype=torch.float32) + del graph_pool_guard + return output + + +def sm120_cublas_router_gemm_fake( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], weight.shape[0]), dtype=torch.float32) + + +direct_register_custom_op( + op_name="sm120_cublas_router_gemm", + op_func=sm120_cublas_router_gemm_impl, + fake_impl=sm120_cublas_router_gemm_fake, +) + + def fp32_router_gemm_dispatch_impl( x: torch.Tensor, weight: torch.Tensor, From 6da84b79479d0b68610f346b9a8fcaa1e7cb1713 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:35:46 +0000 Subject: [PATCH 151/365] Make router eligibility fixtures independent of hardware --- tests/kernels/test_gate_linear_rocm_dispatch.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index de002e020987..aa061f258292 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -11,6 +11,8 @@ ``not bias`` so a biased gate cannot silently drop its bias term. """ +from types import SimpleNamespace + import pytest import torch @@ -62,7 +64,15 @@ def _make_gate( lambda: True, ) if is_rocm: - import vllm.platforms.rocm as rocm_platform + # The module resolves its architecture at import time. Mock that + # hardware probe too, so eligibility tests remain device-free. + with monkeypatch.context() as probe: + probe.setattr( + torch.cuda, + "get_device_properties", + lambda *args, **kwargs: SimpleNamespace(gcnArchName="gfx950"), + ) + import vllm.platforms.rocm as rocm_platform monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950) From dc71c88dfe6bed4ee64f26acb6f5e010705cad6a Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:46:17 +0000 Subject: [PATCH 152/365] Avoid caching a synthetic ROCm architecture in eligibility tests --- tests/kernels/test_gate_linear_rocm_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index aa061f258292..1c7d8df8d84f 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -70,7 +70,7 @@ def _make_gate( probe.setattr( torch.cuda, "get_device_properties", - lambda *args, **kwargs: SimpleNamespace(gcnArchName="gfx950"), + lambda *args, **kwargs: SimpleNamespace(gcnArchName=""), ) import vllm.platforms.rocm as rocm_platform From abf853f3d23bc7f8e0558e4199a605b3dc2b2774 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:49:18 +0000 Subject: [PATCH 153/365] Propagate MiMo target and MTP cache configuration --- .../quantization/test_mimo_v2_cache_config.py | 69 +++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 2 + vllm/model_executor/models/mimo_v2_mtp.py | 7 +- 3 files changed, 77 insertions(+), 1 deletion(-) create mode 100644 tests/models/quantization/test_mimo_v2_cache_config.py diff --git a/tests/models/quantization/test_mimo_v2_cache_config.py b/tests/models/quantization/test_mimo_v2_cache_config.py new file mode 100644 index 000000000000..06386dc7e612 --- /dev/null +++ b/tests/models/quantization/test_mimo_v2_cache_config.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiMo target and native draft constructors retain the requested cache policy.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.models.mimo_v2 as mimo +import vllm.model_executor.models.mimo_v2_mtp as mtp + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("kind", ["global", "sliding", "mtp"]) +@pytest.mark.parametrize("dtype", ["bfloat16", "fp8_e4m3"]) +def test_mimo_attention_receives_cache_policy(monkeypatch, kind, dtype): + observed = [] + + class Attention(torch.nn.Module): + def __init__(self, *, cache_config=None, **kwargs): + super().__init__() + observed.append(cache_config) + + for module in (mimo, mtp): + monkeypatch.setattr(module, "MiMoV2Attention", Attention) + monkeypatch.setattr(module, "MiMoV2MLP", lambda **kwargs: torch.nn.Identity()) + monkeypatch.setattr(module, "RMSNorm", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr(mtp, "LogitsProcessor", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr(mtp, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr( + mtp, "VocabParallelEmbedding", lambda *a, **k: torch.nn.Identity() + ) + config = SimpleNamespace( + hidden_size=32, + num_attention_heads=16, + num_key_value_heads=1, + head_dim=192, + v_head_dim=128, + swa_num_attention_heads=16, + swa_num_key_value_heads=2, + swa_head_dim=192, + swa_v_head_dim=128, + sliding_window_size=128, + attention_bias=False, + intermediate_size=64, + hidden_act="silu", + layernorm_epsilon=1e-6, + vocab_size=64, + hybrid_layer_pattern=[int(kind == "sliding")], + ) + cache_config = SimpleNamespace( + cache_dtype=dtype, calculate_kv_scales=True, kv_cache_dtype_skip_layers=["3"] + ) + vconfig = SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=config, hf_config=config), + quant_config=None, + cache_config=cache_config, + speculative_config=SimpleNamespace(num_speculative_tokens=1), + ) + if kind == "mtp": + mtp.MiMoV2MultiTokenPredictor(vllm_config=vconfig, prefix="model") + else: + mimo.MiMoV2FlashDecoderLayer(vllm_config=vconfig, prefix="model.layers.0") + assert len(observed) == 1 + # The whole policy must reach Attention, including scale calibration and + # intentional skip layers, not just a copied dtype string. + assert observed[0] is cache_config diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index c7ba39bb0dd9..09aa6879920b 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -385,6 +385,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: layer_id=layer_id, rope_theta=getattr(config, "swa_rope_theta", rope_theta), max_position_embeddings=max_position_embeddings, + cache_config=vllm_config.cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", @@ -402,6 +403,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: layer_id=layer_id, rope_theta=rope_theta, max_position_embeddings=max_position_embeddings, + cache_config=vllm_config.cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index fa2cabe28860..e6f93fbcc9d7 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -24,7 +24,7 @@ import torch.nn as nn from transformers import PretrainedConfig -from vllm.config import VllmConfig +from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -65,6 +65,7 @@ def __init__( config: PretrainedConfig, prefix: str, quant_config: QuantizationConfig | None = None, + cache_config: CacheConfig | None = None, ) -> None: super().__init__() @@ -100,6 +101,7 @@ def __init__( layer_id=0, rope_theta=swa_rope_theta, max_position_embeddings=getattr(config, "max_position_embeddings", 32768), + cache_config=cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", @@ -149,6 +151,7 @@ def __init__( num_mtp_layers: int, quant_config: QuantizationConfig | None, prefix: str, + cache_config: CacheConfig | None = None, ) -> None: super().__init__() self.layers = nn.ModuleDict( @@ -157,6 +160,7 @@ def __init__( config=config, prefix=f"{prefix}.{i}", quant_config=quant_config, + cache_config=cache_config, ) for i in range(num_mtp_layers) } @@ -183,6 +187,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config=config, num_mtp_layers=num_mtp_layers, quant_config=vllm_config.quant_config, + cache_config=vllm_config.cache_config, prefix=maybe_prefix(prefix, "mtp.layers"), ) From 73fd464f2a9101f5bb80b25068646c36e1e107d6 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 23:04:59 +0000 Subject: [PATCH 154/365] Preserve MiMo global attention when propagating cache policy --- .../quantization/test_mimo_v2_cache_config.py | 51 +++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 10 ++++ 2 files changed, 61 insertions(+) diff --git a/tests/models/quantization/test_mimo_v2_cache_config.py b/tests/models/quantization/test_mimo_v2_cache_config.py index 06386dc7e612..e5ad81279227 100644 --- a/tests/models/quantization/test_mimo_v2_cache_config.py +++ b/tests/models/quantization/test_mimo_v2_cache_config.py @@ -67,3 +67,54 @@ def __init__(self, *, cache_config=None, **kwargs): # The whole policy must reach Attention, including scale calibration and # intentional skip layers, not just a copied dtype string. assert observed[0] is cache_config + + +@pytest.mark.parametrize("window", [-1, 128]) +@pytest.mark.parametrize("skip_layer", [False, True]) +def test_mimo_preserves_per_layer_window_and_cache_spec( + monkeypatch, window, skip_layer +): + from vllm.config import CacheConfig, set_current_vllm_config + from vllm.v1.attention.backends.registry import AttentionBackendEnum + from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec + + monkeypatch.setattr(mimo, "get_tensor_model_parallel_world_size", lambda: 4) + for name in ("QKVParallelLinear", "RowParallelLinear", "get_rope"): + monkeypatch.setattr(mimo, name, lambda *a, **k: torch.nn.Identity()) + cache_config = CacheConfig( + cache_dtype="bfloat16", + block_size=16, + sliding_window=128, + kv_cache_dtype_skip_layers=["0"] if skip_layer else [], + ) + vconfig = SimpleNamespace( + model_config=SimpleNamespace(dtype=torch.bfloat16, is_mm_prefix_lm=False), + cache_config=cache_config, + attention_config=SimpleNamespace( + backend=AttentionBackendEnum.TRITON_ATTN_DIFFKV + ), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + with set_current_vllm_config(vconfig): + layer = mimo.MiMoV2Attention( + hidden_size=4096, + num_heads=64, + num_kv_heads=8 if window == 128 else 4, + head_dim=192, + v_head_dim=128, + sliding_window_size=window, + cache_config=cache_config, + prefix="model.layers.0.self_attn", + ) + spec = layer.attn.get_kv_cache_spec(vconfig) + assert cache_config.sliding_window == 128 + assert layer.attn.kv_cache_dtype == ("auto" if skip_layer else "bfloat16") + assert layer.attn.query_quant is None + assert spec.dtype == torch.bfloat16 + if window == -1: + assert layer.attn.sliding_window is None + assert type(spec) is FullAttentionSpec + else: + assert layer.attn.sliding_window == 128 + assert type(spec) is SlidingWindowSpec + assert spec.sliding_window == 128 diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 09aa6879920b..99d259170658 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from copy import copy from itertools import islice import torch @@ -292,6 +293,15 @@ def __init__( ) sliding_window = sliding_window_size if sliding_window_size > -1 else None + if ( + sliding_window is None + and cache_config is not None + and cache_config.sliding_window is not None + ): + # MiMo declares each layer's window explicitly. Do not let the + # model-wide SWA default turn a global layer into sliding attention. + cache_config = copy(cache_config) + cache_config.sliding_window = None # Use DiffKV backend when V has a different head dim than K. # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall From b2dffa2665825e5c8395baa70dc911ded3bbc301 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:04:37 +0000 Subject: [PATCH 155/365] perf(gdn): omit mixed worklists for uniform decode graphs Carry the selected graph contract into attention metadata. Uniform speculative Qwen graphs update only their live decode inputs; mixed and prefill graphs retain complete per-group worklists. Co-authored-by: Codex Signed-off-by: Martin Vit --- .../v1/attention/test_gdn_metadata_builder.py | 55 +++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model_state.py | 1 + vllm/v1/attention/backend.py | 3 + vllm/v1/attention/backends/gdn_attn.py | 18 ++++++ vllm/v1/worker/gpu/attn_utils.py | 2 + vllm/v1/worker/gpu/cudagraph_utils.py | 4 ++ vllm/v1/worker/gpu/input_batch.py | 4 ++ vllm/v1/worker/gpu/model_runner.py | 4 ++ 8 files changed, 91 insertions(+) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 2967f2f05e6d..2016fc0b6fd8 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -427,6 +427,61 @@ def test_uniform_spec_decode_falls_back_for_nonuniform_requests( assert not metadata.is_uniform_spec_decode +@pytest.mark.parametrize("uniform_graph", [False, True]) +def test_uniform_graph_skips_only_unused_b12x_mixed_worklists( + monkeypatch, uniform_graph: bool +): + """A graph's kernel contract, not uniform-looking rows, permits elision.""" + monkeypatch.setenv("VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH", "1") + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn." + "_resolve_gdn_prefill_backend", + lambda _: ("b12x", "b12x"), + ) + builders = [_create_gdn_builder(3, full_cuda_graph=True) for _ in range(2)] + for index, builder in enumerate(builders): + builder.vllm_config.cache_config.mamba_cache_mode = "align" + builder.mamba_aligned_state_indices = ( + torch.arange(8, dtype=torch.int32).reshape(2, 4) + 100 * index + ) + if uniform_graph: + + def unexpected_staging(*args, **kwargs): + raise AssertionError("Uniform decode must not stage mixed worklists") + + monkeypatch.setattr(builder._b12x_mixed, "stage", unexpected_staging) + monkeypatch.setattr( + builder._b12x_mixed, "copy_worklists_from", unexpected_staging + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[64, 64], query_lens=[4, 4]), BLOCK_SIZE, DEVICE + ).replace( + is_prefilling=torch.zeros(2, dtype=torch.bool), + uniform_decode_graph=uniform_graph, + ) + captured = [builder.build_for_cudagraph_capture(common) for builder in builders] + drafts = torch.full((2,), 3, dtype=torch.int32) + for values in ([1, 4], [3, 2]): + accepted = torch.tensor(values, dtype=torch.int32) + for builder in builders: + builder.mamba_aligned_state_indices.add_(8) + source = builders[0].build(0, common, accepted, drafts) + rebound = builders[1].update_block_table( + source, common.block_table_tensor, None + ) + for builder, metadata, graph_metadata in zip( + builders, [source, rebound], captured + ): + assert metadata.is_uniform_spec_decode + assert (metadata.b12x_mixed is None) is uniform_graph + torch.testing.assert_close( + graph_metadata.spec_state_indices_tensor, + builder.mamba_aligned_state_indices, + ) + torch.testing.assert_close(graph_metadata.num_accepted_tokens, accepted) + + @pytest.mark.parametrize( "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() ) diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index b131fdbb5192..82c9726f004e 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -392,6 +392,7 @@ def prepare_attn( model_specific_attn_metadata=model_metadata, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, + uniform_decode_graph=input_batch.uniform_decode_graph, ) if self.recoverssm is not None: self.recoverssm.record_step( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b43e5f507162..36b1549344cb 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -485,6 +485,9 @@ class CommonAttentionMetadata: _num_computed_tokens_cache: torch.Tensor | None = None _token_to_req_indices_cache: torch.Tensor | None = None + # True only when the selected graph excludes mixed/prefill kernels. + uniform_decode_graph: bool = False + def batch_size(self) -> int: return self.seq_lens.shape[0] diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 133983779240..123e6373c2cb 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -438,6 +438,14 @@ def build( # type: ignore[override] fast_build: bool = False, ) -> GDNAttentionMetadata: m = common_attn_metadata + if m.uniform_decode_graph and self._can_reuse_spec_inputs( + m, num_accepted_tokens, num_decode_draft_tokens_cpu + ): + assert num_accepted_tokens is not None + # Uniform graphs read the packed decode metadata, not the mixed + # prefill worklists. Mixed graphs keep their complete staging even + # when the request rows happen to be uniform on this invocation. + return self._build_uniform_spec_decode(m, num_accepted_tokens) mixed = getattr(self, "_b12x_mixed", None) if mixed is not None: mixed.stage( @@ -900,6 +908,16 @@ def update_block_table( del slot_mapping assert metadata.num_reqs > 0 assert metadata.seq_lens is not None + if ( + metadata.is_uniform_spec_decode + and metadata.b12x_mixed is None + and self._reuse_spec_decode_inputs + and self.mamba_aligned_state_indices is not None + ): + assert metadata.num_accepted_tokens is not None + return self._build_uniform_spec_decode( + metadata, metadata.num_accepted_tokens + ) mixed = None if metadata.b12x_mixed is not None: mixed = self._b12x_mixed diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index bcfee8b1e353..f955bf8c37c5 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -446,6 +446,7 @@ def build_attn_metadata( ubatch_idx: int = 0, fast_prefill: FastPrefillBatchMetadata | None = None, req_idx: np.ndarray | None = None, + uniform_decode_graph: bool = False, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -535,6 +536,7 @@ def build_attn_metadata( rswa_prefix_lens=rswa_prefix_lens, req_idx=req_idx, _token_to_req_indices_cache=token_to_req_indices, + uniform_decode_graph=uniform_decode_graph, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 83d654459a06..35878124f529 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1001,6 +1001,7 @@ def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: attn_groups, kv_cache_config, full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, + uniform_decode_graph=desc.uniform_token_count is not None, max_query_len=desc.max_query_len, pcp_manager=pcp_manager, ) @@ -1077,6 +1078,7 @@ def prepare_inputs_to_capture( full_cudagraph: bool, max_query_len: int | None = None, pcp_manager: "PCPManager | None" = None, + uniform_decode_graph: bool = False, ) -> AttentionState: input_batch = InputBatch.make_dummy( num_reqs, num_tokens, input_buffers, max_query_len=max_query_len @@ -1084,6 +1086,8 @@ def prepare_inputs_to_capture( if pcp_manager is not None: input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + input_batch.uniform_decode_graph = full_cudagraph and uniform_decode_graph + block_table_provider = pcp_manager or block_tables input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 32e6def84926..78e9f3095d69 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -123,6 +123,10 @@ class InputBatch: # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata). dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + # The selected full graph specializes uniform decode. This is not inferred + # from dummy query lengths: a mixed graph may be captured with uniform rows. + uniform_decode_graph: bool = False + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0038575c3096..c68b5dc3c66c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1630,6 +1630,10 @@ def prepare_inputs( has_structured_output_reqs=scheduler_output.has_structured_output_requests, prompt_lens=prompt_lens, fast_prefill=fast_prefill, + uniform_decode_graph=( + batch_desc.cg_mode == CUDAGraphMode.FULL + and batch_desc.uniform_token_count is not None + ), max_query_len=( int(num_scheduled_tokens_upper_bound.max()) if adaptive_verification is not None From 0089284980b03415a8bc9b11f9362ff2f04b0489 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:05:12 +0000 Subject: [PATCH 156/365] fix(spec-decode): inherit target CuMem allocator permission Keep the draft configuration compatible with the target native KV connector and expandable segments policy. Preserve both enabled and disabled target settings. Co-authored-by: llitz <22463547+llitz@users.noreply.github.com> Co-authored-by: Codex Signed-off-by: Martin Vit --- .../test_speculative_draft_hf_overrides.py | 20 +++++++++++++++++-- vllm/config/speculative.py | 1 + 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index bc5300e53fce..e519e2f92b33 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -181,6 +181,7 @@ def test_composed_override_is_picklable(): def _make_mtp_speculative_config( override: bool | None, checkpoint_value: bool, + enable_cumem_allocator: bool = False, ) -> SpeculativeConfig: draft_hf_config = _make_hf_config( architectures=["Qwen4ExpMTP"], @@ -199,10 +200,13 @@ def _make_mtp_speculative_config( max_model_len=128, quantization=None, hf_overrides={}, + enable_cumem_allocator=enable_cumem_allocator, ) - with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config): - return SpeculativeConfig( + with patch( + "vllm.config.speculative.ModelConfig", return_value=draft_model_config + ) as draft_constructor: + config = SpeculativeConfig( model="draft", method="mtp", num_speculative_tokens=1, @@ -210,6 +214,18 @@ def _make_mtp_speculative_config( target_model_config=target_model_config, target_parallel_config=ParallelConfig(), ) + assert ( + draft_constructor.call_args.kwargs["enable_cumem_allocator"] + is enable_cumem_allocator + ) + return config + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("enabled", [False, True]) +def test_mtp_draft_inherits_cumem_allocator_permission(enabled): + """The draft shares the process allocator used by native CPU KV offload.""" + _make_mtp_speculative_config(None, False, enable_cumem_allocator=enabled) @pytest.mark.cpu_test diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 591d029d1acd..e8e7d9c624c3 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1351,6 +1351,7 @@ def __post_init__(self): max_logprobs=self.target_model_config.max_logprobs, hf_overrides=draft_hf_overrides, config_format=self.target_model_config.config_format, + enable_cumem_allocator=self.target_model_config.enable_cumem_allocator, ) # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*) From ab89b608b74c8b06e67b3e0f4ab4b1e67e1e7809 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:04:37 +0000 Subject: [PATCH 157/365] perf(gdn): omit mixed worklists for uniform decode graphs Carry the selected graph contract into attention metadata. Uniform speculative Qwen graphs update only their live decode inputs; mixed and prefill graphs retain complete per-group worklists. Co-authored-by: Codex Signed-off-by: Martin Vit --- .../v1/attention/test_gdn_metadata_builder.py | 55 +++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model_state.py | 1 + vllm/v1/attention/backend.py | 3 + vllm/v1/attention/backends/gdn_attn.py | 18 ++++++ vllm/v1/worker/gpu/attn_utils.py | 2 + vllm/v1/worker/gpu/cudagraph_utils.py | 4 ++ vllm/v1/worker/gpu/input_batch.py | 4 ++ vllm/v1/worker/gpu/model_runner.py | 4 ++ 8 files changed, 91 insertions(+) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 20fea68b4c85..b490a052a8b8 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -448,6 +448,61 @@ def test_uniform_spec_decode_falls_back_for_nonuniform_requests( assert not metadata.is_uniform_spec_decode +@pytest.mark.parametrize("uniform_graph", [False, True]) +def test_uniform_graph_skips_only_unused_b12x_mixed_worklists( + monkeypatch, uniform_graph: bool +): + """A graph's kernel contract, not uniform-looking rows, permits elision.""" + monkeypatch.setenv("VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH", "1") + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn." + "_resolve_gdn_prefill_backend", + lambda _: ("b12x", "b12x"), + ) + builders = [_create_gdn_builder(3, full_cuda_graph=True) for _ in range(2)] + for index, builder in enumerate(builders): + builder.vllm_config.cache_config.mamba_cache_mode = "align" + builder.mamba_aligned_state_indices = ( + torch.arange(8, dtype=torch.int32).reshape(2, 4) + 100 * index + ) + if uniform_graph: + + def unexpected_staging(*args, **kwargs): + raise AssertionError("Uniform decode must not stage mixed worklists") + + monkeypatch.setattr(builder._b12x_mixed, "stage", unexpected_staging) + monkeypatch.setattr( + builder._b12x_mixed, "copy_worklists_from", unexpected_staging + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[64, 64], query_lens=[4, 4]), BLOCK_SIZE, DEVICE + ).replace( + is_prefilling=torch.zeros(2, dtype=torch.bool), + uniform_decode_graph=uniform_graph, + ) + captured = [builder.build_for_cudagraph_capture(common) for builder in builders] + drafts = torch.full((2,), 3, dtype=torch.int32) + for values in ([1, 4], [3, 2]): + accepted = torch.tensor(values, dtype=torch.int32) + for builder in builders: + builder.mamba_aligned_state_indices.add_(8) + source = builders[0].build(0, common, accepted, drafts) + rebound = builders[1].update_block_table( + source, common.block_table_tensor, None + ) + for builder, metadata, graph_metadata in zip( + builders, [source, rebound], captured + ): + assert metadata.is_uniform_spec_decode + assert (metadata.b12x_mixed is None) is uniform_graph + torch.testing.assert_close( + graph_metadata.spec_state_indices_tensor, + builder.mamba_aligned_state_indices, + ) + torch.testing.assert_close(graph_metadata.num_accepted_tokens, accepted) + + @pytest.mark.parametrize( "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() ) diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index b131fdbb5192..82c9726f004e 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -392,6 +392,7 @@ def prepare_attn( model_specific_attn_metadata=model_metadata, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, + uniform_decode_graph=input_batch.uniform_decode_graph, ) if self.recoverssm is not None: self.recoverssm.record_step( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b43e5f507162..36b1549344cb 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -485,6 +485,9 @@ class CommonAttentionMetadata: _num_computed_tokens_cache: torch.Tensor | None = None _token_to_req_indices_cache: torch.Tensor | None = None + # True only when the selected graph excludes mixed/prefill kernels. + uniform_decode_graph: bool = False + def batch_size(self) -> int: return self.seq_lens.shape[0] diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 2fa597abe42f..b31a76f2f268 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -446,6 +446,14 @@ def build( # type: ignore[override] fast_build: bool = False, ) -> GDNAttentionMetadata: m = common_attn_metadata + if m.uniform_decode_graph and self._can_reuse_spec_inputs( + m, num_accepted_tokens, num_decode_draft_tokens_cpu + ): + assert num_accepted_tokens is not None + # Uniform graphs read the packed decode metadata, not the mixed + # prefill worklists. Mixed graphs keep their complete staging even + # when the request rows happen to be uniform on this invocation. + return self._build_uniform_spec_decode(m, num_accepted_tokens) mixed = getattr(self, "_b12x_mixed", None) if mixed is not None: mixed.stage( @@ -916,6 +924,16 @@ def update_block_table( del slot_mapping assert metadata.num_reqs > 0 assert metadata.seq_lens is not None + if ( + metadata.is_uniform_spec_decode + and metadata.b12x_mixed is None + and self._reuse_spec_decode_inputs + and self.mamba_aligned_state_indices is not None + ): + assert metadata.num_accepted_tokens is not None + return self._build_uniform_spec_decode( + metadata, metadata.num_accepted_tokens + ) mixed = None if metadata.b12x_mixed is not None: mixed = self._b12x_mixed diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index bcfee8b1e353..f955bf8c37c5 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -446,6 +446,7 @@ def build_attn_metadata( ubatch_idx: int = 0, fast_prefill: FastPrefillBatchMetadata | None = None, req_idx: np.ndarray | None = None, + uniform_decode_graph: bool = False, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -535,6 +536,7 @@ def build_attn_metadata( rswa_prefix_lens=rswa_prefix_lens, req_idx=req_idx, _token_to_req_indices_cache=token_to_req_indices, + uniform_decode_graph=uniform_decode_graph, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 83d654459a06..35878124f529 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1001,6 +1001,7 @@ def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: attn_groups, kv_cache_config, full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, + uniform_decode_graph=desc.uniform_token_count is not None, max_query_len=desc.max_query_len, pcp_manager=pcp_manager, ) @@ -1077,6 +1078,7 @@ def prepare_inputs_to_capture( full_cudagraph: bool, max_query_len: int | None = None, pcp_manager: "PCPManager | None" = None, + uniform_decode_graph: bool = False, ) -> AttentionState: input_batch = InputBatch.make_dummy( num_reqs, num_tokens, input_buffers, max_query_len=max_query_len @@ -1084,6 +1086,8 @@ def prepare_inputs_to_capture( if pcp_manager is not None: input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + input_batch.uniform_decode_graph = full_cudagraph and uniform_decode_graph + block_table_provider = pcp_manager or block_tables input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 32e6def84926..78e9f3095d69 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -123,6 +123,10 @@ class InputBatch: # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata). dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + # The selected full graph specializes uniform decode. This is not inferred + # from dummy query lengths: a mixed graph may be captured with uniform rows. + uniform_decode_graph: bool = False + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8a809ee0fe3e..393897f349c6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1646,6 +1646,10 @@ def prepare_inputs( has_structured_output_reqs=scheduler_output.has_structured_output_requests, prompt_lens=prompt_lens, fast_prefill=fast_prefill, + uniform_decode_graph=( + batch_desc.cg_mode == CUDAGraphMode.FULL + and batch_desc.uniform_token_count is not None + ), max_query_len=( int(num_scheduled_tokens_upper_bound.max()) if adaptive_verification is not None From 92859aa1292b8fc1849891f68af558e4ed4fc2a5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:07:53 +0000 Subject: [PATCH 158/365] docs(release): describe speculative CuMem compatibility Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-834.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-834.json diff --git a/.lil/changes/vllm-834.json b/.lil/changes/vllm-834.json new file mode 100644 index 000000000000..462297ab5233 --- /dev/null +++ b/.lil/changes/vllm-834.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-834", + "category": "fix", + "summary": "Preserve the target CuMem allocator setting for speculative drafts", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "MTP can use native CPU KV offload with expandable segments when the target enables the CuMem allocator.", + "Connector safety checks remain enabled; this does not qualify generic Qwen CPU offload." + ], + "pull_requests": [834], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/834"] +} From f07d0d676c0bd46367cb397471109f16676cfab4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:20:55 +0000 Subject: [PATCH 159/365] style(spec-decode): wrap allocator permission argument Co-authored-by: Codex Signed-off-by: Martin Vit --- vllm/config/speculative.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index e8e7d9c624c3..df159bc65303 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1351,7 +1351,9 @@ def __post_init__(self): max_logprobs=self.target_model_config.max_logprobs, hf_overrides=draft_hf_overrides, config_format=self.target_model_config.config_format, - enable_cumem_allocator=self.target_model_config.enable_cumem_allocator, + enable_cumem_allocator=( + self.target_model_config.enable_cumem_allocator + ), ) # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*) From 6bce997226fc683b911644226f629d59f925ddab Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:36:55 +0000 Subject: [PATCH 160/365] docs(release): describe Qwen uniform graph metadata elision Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-835.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-835.json diff --git a/.lil/changes/vllm-835.json b/.lil/changes/vllm-835.json new file mode 100644 index 000000000000..c44d72b3d28f --- /dev/null +++ b/.lil/changes/vllm-835.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-835", + "category": "performance", + "summary": "Reduce CPU preparation for Qwen MTP CUDA graphs", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Uniform speculative decode skips unused B12X mixed-prefill worklists while keeping live recurrent state and acceptance inputs.", + "Mixed batches and prefill retain complete metadata staging. No precision, quantization or all-reduce setting changes." + ], + "pull_requests": [835], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/835"] +} From 0c43e61d83dcfa61352252f1ef62bdf9f62d6214 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:36:55 +0000 Subject: [PATCH 161/365] docs(release): describe Qwen uniform graph metadata elision Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-835.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-835.json diff --git a/.lil/changes/vllm-835.json b/.lil/changes/vllm-835.json new file mode 100644 index 000000000000..c44d72b3d28f --- /dev/null +++ b/.lil/changes/vllm-835.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-835", + "category": "performance", + "summary": "Reduce CPU preparation for Qwen MTP CUDA graphs", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Uniform speculative decode skips unused B12X mixed-prefill worklists while keeping live recurrent state and acceptance inputs.", + "Mixed batches and prefill retain complete metadata staging. No precision, quantization or all-reduce setting changes." + ], + "pull_requests": [835], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/835"] +} From 088aae168779ee3c2e6cf424b5a6b48a699f9502 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:05:12 +0000 Subject: [PATCH 162/365] fix(spec-decode): inherit target CuMem allocator permission Keep the draft configuration compatible with the target native KV connector and expandable segments policy. Preserve both enabled and disabled target settings. Co-authored-by: llitz <22463547+llitz@users.noreply.github.com> Co-authored-by: Codex Signed-off-by: Martin Vit --- .../test_speculative_draft_hf_overrides.py | 20 +++++++++++++++++-- vllm/config/speculative.py | 1 + 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index b81f1f4e292f..bee54eba2bcd 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -218,6 +218,7 @@ def test_composed_override_is_picklable(): def _make_mtp_speculative_config( override: bool | None, checkpoint_value: bool, + enable_cumem_allocator: bool = False, ) -> SpeculativeConfig: draft_hf_config = _make_hf_config( architectures=["Qwen4ExpMTP"], @@ -236,10 +237,13 @@ def _make_mtp_speculative_config( max_model_len=128, quantization=None, hf_overrides={}, + enable_cumem_allocator=enable_cumem_allocator, ) - with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config): - return SpeculativeConfig( + with patch( + "vllm.config.speculative.ModelConfig", return_value=draft_model_config + ) as draft_constructor: + config = SpeculativeConfig( model="draft", method="mtp", num_speculative_tokens=1, @@ -247,6 +251,18 @@ def _make_mtp_speculative_config( target_model_config=target_model_config, target_parallel_config=ParallelConfig(), ) + assert ( + draft_constructor.call_args.kwargs["enable_cumem_allocator"] + is enable_cumem_allocator + ) + return config + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("enabled", [False, True]) +def test_mtp_draft_inherits_cumem_allocator_permission(enabled): + """The draft shares the process allocator used by native CPU KV offload.""" + _make_mtp_speculative_config(None, False, enable_cumem_allocator=enabled) @pytest.mark.cpu_test diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cb79f8999373..c80700ffa653 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1394,6 +1394,7 @@ def __post_init__(self): max_logprobs=self.target_model_config.max_logprobs, hf_overrides=draft_hf_overrides, config_format=self.target_model_config.config_format, + enable_cumem_allocator=self.target_model_config.enable_cumem_allocator, ) # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*) From 14d30c7b6345bfb632864ec1803f2abcfb15c371 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:07:53 +0000 Subject: [PATCH 163/365] docs(release): describe speculative CuMem compatibility Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-834.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-834.json diff --git a/.lil/changes/vllm-834.json b/.lil/changes/vllm-834.json new file mode 100644 index 000000000000..462297ab5233 --- /dev/null +++ b/.lil/changes/vllm-834.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-834", + "category": "fix", + "summary": "Preserve the target CuMem allocator setting for speculative drafts", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "MTP can use native CPU KV offload with expandable segments when the target enables the CuMem allocator.", + "Connector safety checks remain enabled; this does not qualify generic Qwen CPU offload." + ], + "pull_requests": [834], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/834"] +} From 96b79aa073c2059462b613876366cf693d1c0c27 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 01:20:55 +0000 Subject: [PATCH 164/365] style(spec-decode): wrap allocator permission argument Co-authored-by: Codex Signed-off-by: Martin Vit --- vllm/config/speculative.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c80700ffa653..13fe2e87c765 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1394,7 +1394,9 @@ def __post_init__(self): max_logprobs=self.target_model_config.max_logprobs, hf_overrides=draft_hf_overrides, config_format=self.target_model_config.config_format, - enable_cumem_allocator=self.target_model_config.enable_cumem_allocator, + enable_cumem_allocator=( + self.target_model_config.enable_cumem_allocator + ), ) # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*) From 6235b993c0267a98201fc9776b1511f2514c6a49 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:00:06 +0000 Subject: [PATCH 165/365] docs(release): preserve the Qwen positional override fragment Signed-off-by: Martin Vit --- .lil/changes/vllm-777.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 .lil/changes/vllm-777.json diff --git a/.lil/changes/vllm-777.json b/.lil/changes/vllm-777.json new file mode 100644 index 000000000000..b28ed461166d --- /dev/null +++ b/.lil/changes/vllm-777.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-777", + "category": "fix", + "summary": "Preserve Qwen positional overrides in MTP draft configuration", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "YaRN deployments should verify their configured target and draft context limits before serving.", + "details": [ + "Dictionary-based target Hugging Face overrides reach the internally constructed MTP draft after its model type and architecture are normalized." + ], + "pull_requests": [777], + "requires": [] +} From ceba0ea68330e04a9ac68a6a55f72eae08a028c6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:00:06 +0000 Subject: [PATCH 166/365] docs(release): preserve the KDA alias fix release fragment Signed-off-by: Martin Vit --- .lil/changes/vllm-824.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-824.json diff --git a/.lil/changes/vllm-824.json b/.lil/changes/vllm-824.json new file mode 100644 index 000000000000..a4f179084684 --- /dev/null +++ b/.lil/changes/vllm-824.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-824", + "category": "fix", + "summary": "Preserve the GLM KDA gate when a decode output aliases it", + "models": ["GLM-5.3-Flash"], + "compatibility": "No configuration change is required. Non-overlapping KDA outputs retain the allocation-free path.", + "details": [ + "When B12X rejects an output overlapping the read-only gate, vLLM binds a disjoint caller-owned output and copies the result back after recurrence.", + "Other B12X binding errors still propagate without a fallback." + ], + "pull_requests": [824], + "authors": ["Local Inference Lab"], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/824"], + "requires": [] +} From fe68924c7644c0175e0cdb14d6d61dcc51a974f4 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Wed, 2 Sep 2026 16:31:45 -0700 Subject: [PATCH 167/365] [Bugfix][KV Offload] SimpleCPUOffload: skip non-prefix-cacheable groups Scratch KV cache groups such as the QSA key ring are not hash-addressable and cannot be offloaded. Exclude them from lazy target sizing and from scheduler store/load transfer selection using KVCacheConfig.prefix_cacheable_group_ids. Add a QSA-shaped scheduler regression covering eager store/load and lazy target sizing. Co-authored-by: Claude Fable 5.1 Signed-off-by: Zhewen Li (cherry picked from commit 3cc857e33d3c693ba4beef2eb8a453a99d969877) (cherry picked from commit a598d29dabc857f0feb74d40fcc88a9951392ac1) (cherry picked from commit 56a456a56d9c94782e6832e4812630ac768e4db7) --- tests/v1/simple_kv_offload/test_scheduler.py | 120 +++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 10 ++ 2 files changed, 130 insertions(+) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 901e076e0950..4678bf7498d9 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -41,6 +41,7 @@ register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -2843,3 +2844,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True + + +QSA_RING_BLOCK_SIZE = 4 + + +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: + register_all_kvcache_specs(vllm_config=None) + full_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring_spec = CircularBufferSpec( + block_size=QSA_RING_BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + ) + groups = [ + KVCacheGroupSpec(["layer_0"], full_spec), + KVCacheGroupSpec(["qsa_ring"], ring_spec), + ] + tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, + ) + for group in groups + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _make_qsa_scheduler( + lazy: bool = False, +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + sched = SimpleCPUOffloadScheduler( + vllm_config=_make_vllm_config(), + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + lazy_offload=lazy, + ) + gpu_pool = BlockPool( + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE + ) + sched.bind_gpu_block_pool(gpu_pool) + return sched, gpu_pool + + +def test_qsa_ring_group_is_never_stored_or_loaded() -> None: + sched, gpu_pool = _make_qsa_scheduler() + assert sched.prefix_cacheable_group_ids == {0} + + num_blocks = 2 + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + ring_block = gpu_pool.get_new_blocks(1)[0] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} + assert ring_block.block_id not in meta.store_gpu_blocks + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-qsa-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + ring_block2 = gpu_pool.get_new_blocks(1)[0] + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} + assert ring_block2.block_id not in meta2.load_gpu_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + without_ring = KVCacheConfig( + num_blocks=16, + kv_cache_tensors=with_ring.kv_cache_tensors[:1], + kv_cache_groups=with_ring.kv_cache_groups[:1], + ) + max_batched = 64 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + with_ring, max_batched + ) + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + without_ring, max_batched + ) + assert target_with == target_without diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index bc7e6c3a13b4..825911fc6a46 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -165,6 +165,9 @@ def __init__( kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks + self.prefix_cacheable_group_ids = frozenset( + kv_cache_config.prefix_cacheable_group_ids + ) self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU # Find the full attention kv group for prefix cache matching. self.fa_gidx = -1 @@ -299,6 +302,8 @@ def _estimate_lazy_target_blocks( target = 0 for g in kv_cache_config.kv_cache_groups: spec = g.kv_cache_spec + if not spec.prefix_cacheable: + continue # Only full attention is sharded across DCP ranks; replicated specs # (mamba, sliding window, chunked-local) keep their own block size. block_size = spec.block_size * dcp_world_size_for_kv_cache_spec( @@ -443,6 +448,9 @@ def update_state_after_alloc( # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): + if g not in self.prefix_cacheable_group_ids: + cpu_hit_blocks.append([]) + continue g_block_size = self.group_block_sizes[g] n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -831,6 +839,8 @@ def _select_eager_blocks_to_store( for g, group_gpu_ids in enumerate(block_ids_by_group): if len(gpu_block_ids) >= num_free: break + if g not in self.prefix_cacheable_group_ids: + continue group_manager = self.cpu_coordinator.single_type_managers[g] if not group_manager.has_positionally_stable_blocks: continue From 5ca070bd94493de74623c6eeba961a051fc04818 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 21 Sep 2026 19:24:54 +0000 Subject: [PATCH 168/365] docs(release): attribute QSA ring offload import Identify the upstream SimpleCPUOffloadConnector source and its scope in the integration container changelog. Signed-off-by: Martin Vit (cherry picked from commit f6485322f97357a21f6e92bd5f2be5f9b8555e3e) --- .lil/changes/vllm-upstream-55033.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-upstream-55033.json diff --git a/.lil/changes/vllm-upstream-55033.json b/.lil/changes/vllm-upstream-55033.json new file mode 100644 index 000000000000..9000cc06c9d8 --- /dev/null +++ b/.lil/changes/vllm-upstream-55033.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-upstream-55033", + "category": "fix", + "summary": "Exclude Qwen QSA scratch cache from Simple CPU offload", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "Applies only when SimpleCPUOffloadConnector is enabled. The regular OffloadingConnector remains unqualified for this model.", + "details": [ + "The optional Simple CPU KV connector excludes non-prefix-cacheable QSA ring blocks from lazy target sizing, external load selection, and eager stores.", + "The change preserves the original author and code from vllm-project/vllm#55033." + ], + "pull_requests": [], + "authors": ["Zhewen Li"], + "evidence": ["https://github.com/vllm-project/vllm/pull/55033"], + "requires": [] +} From eafcc4191a263e372f7372457409c8721fd616a6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:01:19 +0000 Subject: [PATCH 169/365] docs(release): preserve integration changelog contributor policy Signed-off-by: Martin Vit --- .lil/README.md | 33 +++++++++++++++++++++++++++++++++ AGENTS.md | 7 +++++++ 2 files changed, 40 insertions(+) create mode 100644 .lil/README.md diff --git a/.lil/README.md b/.lil/README.md new file mode 100644 index 000000000000..c86c03acb3fb --- /dev/null +++ b/.lil/README.md @@ -0,0 +1,33 @@ +# Integration container changelog fragments + +A runtime-affecting change that will be pushed or imported into an +`integration/*` branch must include one versioned JSON fragment in +`.lil/changes/`. The community container publisher compares these fragments +between exact component revisions and creates the container release notes. + +The complete schema, examples, correction policy, and publication flow are in +the [integration container changelog runbook](https://github.com/local-inference-lab/blackwell-llm-docker/blob/main/docs/integration-release-changelog.md). + +Use a component-prefixed identifier that matches the filename. For example, +`.lil/changes/vllm-816.json` contains: + +```json +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-816", + "category": "fix", + "summary": "Merge QSA selection across DCP ranks", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "DCP1, DCP2, and DCP4 use the same distributed selection contract." + ], + "pull_requests": [816], + "requires": ["b12x-402"] +} +``` + +The fragment belongs in the implementation PR or in the commit series that +imports the reviewed change. Preserve contributor attribution. Once a container +release has published a fragment, do not edit or delete it; add a corrective +fragment instead. diff --git a/AGENTS.md b/AGENTS.md index a1060975b0b4..ec1a0cea9c96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -156,3 +156,10 @@ vulnerability process. - **Editing these instructions**: [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) — Rules for modifying AGENTS.md or any domain-specific guide it references. + +## LIL integration container changelog + +Before pushing or importing a runtime-affecting change into an +`integration/*` branch, read [`.lil/README.md`](.lil/README.md) and add the +required `.lil/changes/.json` fragment to the same commit series. Preserve +published fragments unchanged; corrections use a new fragment. From 7fc33e97b371e4631af2cffebfd54b67fef12365 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:17:16 +0000 Subject: [PATCH 170/365] docs: record canonical Qwen and Spark contracts in beta release notes Signed-off-by: Martin Vit --- .../vllm-karmic-canonical-sync-20260922.json | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .lil/changes/vllm-karmic-canonical-sync-20260922.json diff --git a/.lil/changes/vllm-karmic-canonical-sync-20260922.json b/.lil/changes/vllm-karmic-canonical-sync-20260922.json new file mode 100644 index 000000000000..553bf740dfbe --- /dev/null +++ b/.lil/changes/vllm-karmic-canonical-sync-20260922.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-karmic-canonical-sync-20260922", + "category": "compatibility", + "summary": "Carry canonical Qwen tensor-parallel preparation and Spark launch contracts into beta", + "models": ["Qwen3.8-Flash-Next", "DeepSeek-V4-Flash", "DeepSeek-V4.1-Flash", "GLM-5.3-Flash"], + "compatibility": "No user action required.", + "details": [ + "Distribute eligible Qwen residual-mixing projections across tensor-parallel ranks.", + "Preserve per-head GDN normalization geometry during warmup and use the declared B12X QSA capacity API.", + "Retain canonical Spark switched-RoCE and DeepSeek TP3 launcher settings." + ], + "authors": ["lukealonso"], + "requires": [] +} From 3e7b8f223d9adfac967377f3e578d90d76c0b81e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:17:18 +0000 Subject: [PATCH 171/365] test(glm5next): retain selector release and rebinding assertions Signed-off-by: Martin Vit --- tests/models/test_glm5next_pooled_indexer.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 21248d24b3b2..3be850cc9911 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -821,10 +821,15 @@ def test_glm53_selector_constructs_with_matching_preparation_heads( assert indexer.indexer_op._num_q_heads == indexer.scratch.q_fp8.shape[1] == 32 assert indexer._physical_selection_plan is None indexer._physical_selection_plan = object() + indexer.indexer_op._plans[("decode", 8)] = object() indexer.unbind_main_kv_cache() assert indexer._physical_selection_plan is None assert indexer._index_cache is None assert indexer._main_cache_num_blocks == 0 + assert indexer.indexer_op._index_cache is None + assert indexer.indexer_op._plans == {} + indexer.bind_main_kv_cache(main) + assert indexer.indexer_op._index_cache is indexer._index_cache def test_glm53_packed_tail_accepts_nvfp4_main_record() -> None: From a3dd23009f26f02ea66f388680511ece5d873646 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 04:28:01 +0000 Subject: [PATCH 172/365] docs: describe canonical GDN metadata fusion in beta releases Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-karmic-gdn-metadata-fusion.json | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .lil/changes/vllm-karmic-gdn-metadata-fusion.json diff --git a/.lil/changes/vllm-karmic-gdn-metadata-fusion.json b/.lil/changes/vllm-karmic-gdn-metadata-fusion.json new file mode 100644 index 000000000000..8d21b527aa65 --- /dev/null +++ b/.lil/changes/vllm-karmic-gdn-metadata-fusion.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-karmic-gdn-metadata-fusion", + "category": "performance", + "summary": "Fuse recurrent metadata copies across cache groups", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Retain captured buffer addresses while copying worklists and refreshing recurrent-state indices in one GPU launch.", + "Mixed batches use the fused refresh; uniform decode graphs retain the separate metadata-elision optimization." + ], + "authors": ["lukealonso"], + "pull_requests": [], + "requires": [] +} From aa86abd04b6dee28403283e1f4359cf2f2134ecd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:41:56 +0000 Subject: [PATCH 173/365] Bound InstantTensor staging and retain deferred checkpoint ownership Signed-off-by: Martin Vit --- .../instanttensor_loader/test_weight_utils.py | 597 +++++++++++++++++- .../model_loader/test_registry.py | 80 +++ .../model_loader/test_reload.py | 149 +++++ tests/models/kimi_k3/test_weight_loading.py | 47 ++ .../model_loader/default_loader.py | 56 +- .../model_loader/reload/layerwise.py | 116 +++- .../model_loader/weight_utils.py | 484 +++++++++++++- vllm/model_executor/weight_transfer.py | 5 + vllm/v1/worker/gpu_worker.py | 6 + 9 files changed, 1497 insertions(+), 43 deletions(-) diff --git a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py index d97cd93caca4..f3651bfb4198 100644 --- a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py @@ -2,14 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob +import inspect +import json +import struct import sys from contextlib import contextmanager from types import SimpleNamespace import pytest import torch +from safetensors.torch import save_file import vllm.model_executor.model_loader.weight_utils as weight_utils +from vllm.model_executor.model_loader.reload.layerwise import ( + _own_deferred_accelerator_tensors, +) from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, instanttensor_weights_iterator, @@ -18,13 +25,68 @@ from vllm.platforms import current_platform -def test_instanttensor_requests_owned_tensors(monkeypatch): +def _safetensors_tensor_metadata(filename): + with open(filename, "rb") as checkpoint: + header_size = struct.unpack(" None: self.post_load_called = True +class _SourceLifetimeQuantMethod(QuantizeMethodBase): + """Observes checkpoint-source ownership when online processing starts.""" + + uses_meta_device = True + + def __init__(self, source_refs): + self.source_refs = source_refs + self.sources_alive_at_process = None + + def create_weights(self, layer, *weight_args, **extra_weight_attrs): + pass + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + gc.collect() + self.sources_alive_at_process = [ + source_ref() is not None for source_ref in self.source_refs + ] + + +class _DeferredOnlineQuantAttention(_ReloadableAttentionLayer): + """Attention layer whose checkpoint tensor is processed after loading.""" + + def __init__(self): + torch.nn.Module.__init__(self) + self.source_refs = [] + self.quant_method = _SourceLifetimeQuantMethod(self.source_refs) + + def tracking_weight_loader(param, loaded_weight): + self.source_refs.append(ref(loaded_weight)) + default_weight_loader(param, loaded_weight) + + weight = torch.nn.Parameter(torch.empty(2, 2, device="meta")) + weight.weight_loader = tracking_weight_loader + self.register_parameter("weight", weight) + self.post_load_called = False + initialize_online_processing(self) + + +def test_attention_first_load_releases_sources_before_online_quantization(): + model_config = types.SimpleNamespace(dtype=torch.float32) + layer = _DeferredOnlineQuantAttention() + model = torch.nn.Sequential(layer) + + layer.weight.weight_loader(layer.weight, torch.full((2, 2), 7.0)) + + assert layer.source_refs[0]() is not None + finalize_layerwise_reload(model, model_config) + + assert layer.quant_method.sources_alive_at_process + assert not any(layer.quant_method.sources_alive_at_process) + assert layer.post_load_called + assert torch.equal(layer.weight, torch.full((2, 2), 7.0)) + + +def test_initial_online_processing_loads_into_materialized_parameters(): + quant_method = _RecordingQuantMethod() + layer = _LateBiasLayer(quant_method) + loaded_weight = torch.full((4, 2), 2.0) + loaded_bias = torch.full((4,), 3.0) + + layer.weight.weight_loader(layer.weight, loaded_weight) + + assert not layer.weight.is_meta + assert torch.equal(layer.weight, loaded_weight) + assert not get_layerwise_info(layer).loaded_weights + assert quant_method.bias_at_process is None + + layer.bias.weight_loader(layer.bias, loaded_bias) + + assert torch.equal(quant_method.bias_at_process, loaded_bias) + assert not get_layerwise_info(layer).loaded_weights + + +def test_initial_online_processing_preserves_deferred_writer_accounting(): + from vllm.model_executor.weight_transfer import weight_transfer + + class DeferredWriter: + def __init__(self): + self.pending = [] + + def __call__(self, destination, source): + self.pending.append((destination, source)) + return True + + def flush(self): + for destination, source in self.pending: + destination.copy_(source) + self.pending.clear() + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(torch.empty(4, 2, device="meta")) + layer.weight.weight_loader = default_weight_loader + initialize_online_processing(layer) + expected = torch.arange(8, dtype=torch.float32).reshape(4, 2) + writer = DeferredWriter() + + with weight_transfer(writer): + layer.weight.weight_loader(layer.weight, expected) + + assert not layer.weight.is_meta + assert not get_layerwise_info(layer).can_load() + assert not writer.pending + torch.testing.assert_close(layer.weight, expected) + + +def test_online_processing_finalizes_checkpoint_omitted_padding(): + class RecordingWeightQuantMethod(QuantizeMethodBase): + uses_meta_device = True + + def __init__(self): + self.weight_at_process = None + + def create_weights(self, layer, *args, **kwargs): + raise NotImplementedError + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + self.weight_at_process = layer.weight.detach().clone() + + def shard_loader(param, loaded_weight, shard_id): + start = shard_id * loaded_weight.shape[0] + param.data[start : start + loaded_weight.shape[0]].copy_(loaded_weight) + + quant_method = RecordingWeightQuantMethod() + layer = torch.nn.Module() + layer.quant_method = quant_method + weight = torch.nn.Parameter(torch.empty(6, 2, device="meta")) + weight.weight_loader = shard_loader + layer.register_parameter("weight", weight) + initialize_online_processing(layer) + layer._vllm_online_processing_unloaded = {"weight": 4} + + first = torch.full((2, 2), 3.0) + second = torch.full((2, 2), 7.0) + layer.weight.weight_loader(layer.weight, first, 0) + assert quant_method.weight_at_process is None + layer.weight.weight_loader(layer.weight, second, 1) + + expected = torch.cat((first, second, torch.zeros(2, 2))) + assert torch.equal(quant_method.weight_at_process, expected) + assert not get_layerwise_info(layer).can_load() + + def test_move_metatensors(): tensor = torch.empty((1, 2, 3)) meta_tensor = to_meta_tensor(tensor) diff --git a/tests/models/kimi_k3/test_weight_loading.py b/tests/models/kimi_k3/test_weight_loading.py index b24ad17d9a22..f135d247b172 100644 --- a/tests/models/kimi_k3/test_weight_loading.py +++ b/tests/models/kimi_k3/test_weight_loading.py @@ -21,6 +21,11 @@ def __init__(self) -> None: self.tensor_a = nn.Parameter(torch.zeros(1)) self.tensor_c = nn.Parameter(torch.zeros(1)) self.finalized_values: list[tuple[float, float]] = [] + self.workspace_reservations = 0 + + def reserve_attn_res_workspace(self) -> None: + assert self.finalized_values == [(1.0, 3.0)] + self.workspace_reservations += 1 def load_weights(self, weights): params = dict(self.named_parameters()) @@ -69,4 +74,46 @@ def test_interleaved_composite_weights_finalize_kimi_once_after_loading() -> Non model.process_weights_after_loading() assert language_model.model.finalized_values == [(1.0, 3.0)] + assert language_model.model.workspace_reservations == 1 assert model.vision_tower.tensor_b.item() == 2.0 + + +@pytest.mark.parametrize("source_ndim", [1, 4]) +@pytest.mark.parametrize("rank", [0, 9]) +def test_kda_head_parameters_zero_only_checkpoint_absent_tp_tail( + monkeypatch, source_ndim, rank +): + from vllm.models.kimi_k3.nvidia import kda + + monkeypatch.setattr(kda, "get_tensor_model_parallel_rank", lambda: rank) + source = torch.arange(96, dtype=torch.float32) + 1 + param = nn.Parameter(torch.full((10,), float("nan")), requires_grad=False) + param.allow_tp_padding = True + loaded = source if source_ndim == 1 else source.view(1, 1, 96, 1) + kda.a_log_weight_loader(0)(param, loaded) + expected = torch.nn.functional.pad(source, (0, 4))[rank * 10 : (rank + 1) * 10] + torch.testing.assert_close(param, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("with_decode_copy", [False, True]) +def test_kda_packed_conv_loading_preserves_qkv_boundaries_with_tp_tail( + with_decode_copy, +): + from vllm.models.kimi_k3.nvidia.kda import _make_decode_conv1d_weight_loader + + source = torch.arange(96 * 3, dtype=torch.float32).view(96, 1, 3) + param = nn.Parameter(torch.full((30, 1, 3), float("nan")), requires_grad=False) + param.allow_tp_padding = True + decode = torch.full((3, 3, 10), float("nan")) if with_decode_copy else None + loader = _make_decode_conv1d_weight_loader([100] * 3, 10, 9, decode) + for shard in (2, 0, 1): + weight = source + 1000 * shard + loader(param, weight, shard) + expected = torch.cat((weight[90:], torch.zeros(4, 1, 3))) + torch.testing.assert_close( + param[shard * 10 : (shard + 1) * 10], expected, rtol=0, atol=0 + ) + if decode is not None: + torch.testing.assert_close( + decode[shard], expected.squeeze(1).T, rtol=0, atol=0 + ) diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index afc93737e15a..583c9e31b99b 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import glob +import json import os import time from collections.abc import Callable, Generator, Iterable @@ -95,6 +96,8 @@ def __init__(self, load_config: LoadConfig): "enable_multithread_load", "num_threads", "enable_weights_track", + "instanttensor_priority_weight_name_prefixes", + "instanttensor_small_checkpoint_max_bytes", } unexpected_keys = set(extra_config.keys()) - allowed_keys @@ -119,6 +122,30 @@ def __init__(self, load_config: LoadConfig): f"num_threads must be a positive integer, got {num_threads!r}" ) + priority_prefixes = extra_config.get( + "instanttensor_priority_weight_name_prefixes" + ) + if priority_prefixes is not None and not ( + isinstance(priority_prefixes, list) + and priority_prefixes + and all(isinstance(prefix, str) and prefix for prefix in priority_prefixes) + ): + raise ValueError( + "instanttensor_priority_weight_name_prefixes must be a " + "non-empty list of non-empty strings" + ) + small_checkpoint_max_bytes = extra_config.get( + "instanttensor_small_checkpoint_max_bytes" + ) + if small_checkpoint_max_bytes is not None and not ( + isinstance(small_checkpoint_max_bytes, int) + and not isinstance(small_checkpoint_max_bytes, bool) + and small_checkpoint_max_bytes > 0 + ): + raise ValueError( + "instanttensor_small_checkpoint_max_bytes must be a positive integer" + ) + self.enable_weights_track: bool | None = extra_config.get( "enable_weights_track", None ) @@ -259,7 +286,11 @@ def _prepare_weights( return hf_folder, hf_weights_files, use_safetensors def _safetensors_weights_iterator( - self, hf_weights_files: list[str], source: "Source" + self, + hf_weights_files: list[str], + source: "Source", + *, + indexed_tensor_files: dict[str, str] | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: extra_config = self.load_config.model_loader_extra_config if self.load_config.load_format == "fastsafetensors": @@ -273,6 +304,13 @@ def _safetensors_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, weight_name_prefixes=source.weight_name_prefixes, + indexed_tensor_files=indexed_tensor_files, + priority_weight_name_prefixes=extra_config.get( + "instanttensor_priority_weight_name_prefixes" + ), + small_checkpoint_max_bytes=extra_config.get( + "instanttensor_small_checkpoint_max_bytes" + ), ) if extra_config.get("enable_multithread_load"): return multi_thread_safetensors_weights_iterator( @@ -310,6 +348,16 @@ def _get_weights_iterator( ) if source.file_weight_filter is not None and not use_safetensors: raise ValueError("Checkpoint file-backed weights require safetensors files") + indexed_tensor_files: dict[str, str] | None = None + if use_safetensors and self.load_config.load_format == "instanttensor": + index_path = os.path.join(hf_folder, SAFE_WEIGHTS_INDEX_NAME) + if os.path.isfile(index_path): + with open(index_path, encoding="utf-8") as index_handle: + weight_map = json.load(index_handle)["weight_map"] + indexed_tensor_files = { + name: os.path.abspath(os.path.join(hf_folder, filename)) + for name, filename in weight_map.items() + } if self.load_config.load_format == "npcache": # Currently np_cache only support *.bin checkpoints assert use_safetensors is False @@ -323,7 +371,7 @@ def _get_weights_iterator( elif use_safetensors: if source.file_weight_filter is None: weights_iterator = self._safetensors_weights_iterator( - hf_weights_files, source + hf_weights_files, source, indexed_tensor_files=indexed_tensor_files ) else: load_format = self.load_config.load_format @@ -342,7 +390,9 @@ def _get_weights_iterator( tensor_order = "name" weights_iterator = file_backed_safetensors_weights_iterator( ordered_files, - lambda files: self._safetensors_weights_iterator(files, source), + lambda files: self._safetensors_weights_iterator( + files, source, indexed_tensor_files=indexed_tensor_files + ), source.file_weight_filter, weight_name_prefixes=source.weight_name_prefixes, local_expert_ids=self.local_expert_ids, diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 0804eb279ef9..0aa1dc406ff6 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -12,7 +12,10 @@ from vllm.model_executor.layers.attention import is_deferred_attention_layer from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.weight_transfer import flush_weight_transfers +from vllm.model_executor.weight_transfer import ( + flush_weight_transfers, + is_weight_transfer_active, +) from .meta import ( SKIP_LOAD_TENSORS, @@ -54,6 +57,60 @@ LOADING_LAYERS: WeakSet[torch.nn.Module] = WeakSet() +_ONLINE_PROCESSING_UNLOADED_ATTR = "_vllm_online_processing_unloaded" + + +def _online_processing_load_numel_total(layer: torch.nn.Module) -> int: + """Count checkpoint elements required before processing a layer. + + A layer can allocate padded tensor tails that have no corresponding + checkpoint payload. The layer declares those omitted element counts in a + ``dict[str, int]`` named ``_vllm_online_processing_unloaded``. + + Args: + layer: Layer whose registered tensors and unloaded-tail annotations + define the expected checkpoint payload. + + Returns: + Number of tensor elements that checkpoint weight loaders must provide. + + Raises: + TypeError: If the unloaded-tail annotation is not a dictionary. + ValueError: If an annotated tensor is absent or its unloaded element + count is outside the tensor's bounds. + """ + total = get_layer_size(layer) + unloaded = getattr(layer, _ONLINE_PROCESSING_UNLOADED_ATTR, {}) + if not isinstance(unloaded, dict): + raise TypeError( + f"{_ONLINE_PROCESSING_UNLOADED_ATTR} must be a dict, " + f"got {type(unloaded).__name__}" + ) + for name, numel in unloaded.items(): + tensor = get_layer_tensors(layer).get(name) + if tensor is None: + raise ValueError(f"Annotated unloaded tensor {name!r} is missing") + if not isinstance(numel, int) or not 0 <= numel <= tensor.numel(): + raise ValueError( + f"Invalid unloaded element count for {name}: {numel} " + f"(tensor numel={tensor.numel()})" + ) + total -= numel + return total + + +def _zero_online_processing_unloaded(layer: torch.nn.Module) -> None: + """Zero checkpoint-omitted parameter tails before online processing. + + Args: + layer: Materialized layer with unloaded-tail annotations. + """ + unloaded = getattr(layer, _ONLINE_PROCESSING_UNLOADED_ATTR, {}) + for name, numel in unloaded.items(): + if numel: + getattr(layer, name).data.view(-1)[-numel:].zero_() + + def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo: """ Get information related to restoring and layerwise processing. If no previous @@ -133,7 +190,7 @@ def initialize_online_processing(layer: torch.nn.Module): # Track loading progress to determine when to process/copy info.load_numel = 0 - info.load_numel_total = get_layer_size(layer) + info.load_numel_total = _online_processing_load_numel_total(layer) _wrap_parameters_weight_loader(layer) @@ -147,6 +204,28 @@ def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: tensor.weight_loader = make_online_process_loader(layer, name) +def _own_deferred_accelerator_tensors(bound_args: inspect.BoundArguments) -> None: + """Give deferred weight-loader arguments independent device storage. + + Streaming model loaders may yield views into reusable staging storage. + Layerwise processing replays bound arguments only after every checkpoint + shard for a layer arrives, so accelerator tensors must be cloned before + they enter the deferred queue. This also covers views created by name or + shard mapping because custom attributes on the source tensor are not + guaranteed to propagate to those views. + + Args: + bound_args: Normalized weight-loader arguments to update in place. + """ + for name, value in tuple(bound_args.arguments.items()): + if ( + name != "param" + and isinstance(value, torch.Tensor) + and value.device.type not in ("meta", "cpu") + ): + bound_args.arguments[name] = value.clone() + + def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable: """Create a wrapped weight loader that defers processing.""" info = get_layerwise_info(layer) @@ -177,16 +256,28 @@ def online_process_loader(*args, **kwargs): # Re-run on each load: layers may register parameters later (e.g., `bias`). # Wrap late parameters and refresh `load_numel_total` so processing waits # until all parameters are loaded. - info.load_numel_total = get_layer_size(layer) + info.load_numel_total = _online_processing_load_numel_total(layer) _wrap_parameters_weight_loader(layer) # Bind and normalize arguments bound_args = loader_signature.bind(*args, **kwargs) bound_args.apply_defaults() - # Buffer loaded weights, track loading progress - info.loaded_weights.append((param_name, bound_args)) - num_loaded, ret = get_numel_loaded(original_loader, bound_args) + direct_load = ( + info.kernel_tensors is None + and not is_deferred_attention_layer(layer) + and not is_weight_transfer_active() + ) + if direct_load: + # Initial loading consumes each source before its staging slot is + # reused. Reloads preserve kernel storage and defer their writes. + materialize_layer(layer, info) + bound_args.arguments["param"] = getattr(layer, param_name) + num_loaded, ret = get_numel_loaded(original_loader, bound_args) + else: + _own_deferred_accelerator_tensors(bound_args) + info.loaded_weights.append((param_name, bound_args)) + num_loaded, ret = get_numel_loaded(original_loader, bound_args) info.load_numel += num_loaded logger.debug( @@ -353,12 +444,19 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): param.weight_loader = _get_original_loader(param) # Load all buffered weights into materialized layer (using original loaders) - for name, args in info.loaded_weights: + loaded_args = None + for name, loaded_args in info.loaded_weights: param = getattr(layer, name) - args.arguments["param"] = param - param.weight_loader(*args.args, **args.kwargs) + loaded_args.arguments["param"] = param + param.weight_loader(*loaded_args.args, **loaded_args.kwargs) flush_weight_transfers() + if info.kernel_tensors is None: + # All initial-load copies are complete. Quantization may use the memory + # formerly held by checkpoint sources for its temporary workspace. + info.loaded_weights.clear() + loaded_args = None + _zero_online_processing_unloaded(layer) # Process weights (quantization, repacking, etc.) quant_method = getattr(layer, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index c669d54bacb2..cee0ba0453a5 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -14,7 +14,8 @@ import time from collections import defaultdict from collections.abc import Callable, Generator, Iterable, Sequence -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass from pathlib import Path from typing import IO, Any, Literal @@ -1327,13 +1328,35 @@ def _make_loader(nogds: bool) -> "ParallelLoader": pl.close() +@dataclass(frozen=True) +class _InstantTensorCpuFallback: + position: int + name: str + filename: str + + +@dataclass(frozen=True) +class _InstantTensorSelection: + gpu_tensor_count: int + selected_tensor_count: int + cpu_fallbacks: tuple[_InstantTensorCpuFallback, ...] + + def instanttensor_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, weight_name_prefixes: Sequence[str] | None = None, + *, + indexed_tensor_files: dict[str, str] | None = None, + priority_weight_name_prefixes: Sequence[str] | None = None, + small_checkpoint_max_bytes: int | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: - """Iterate over the weights in the model safetensor files - using instanttensor library.""" + """Iterate over weights in model safetensor files with InstantTensor. + + When an index is available, restrict the physical I/O layout before + opening the loader. This avoids loading stale tensors from base shards + when a composed checkpoint remaps their names to overlay shards. + """ try: import instanttensor except ImportError as e: @@ -1344,6 +1367,113 @@ def instanttensor_weights_iterator( if not current_platform.is_cuda(): raise ValueError("InstantTensor requires NVIDIA GPUs") + if small_checkpoint_max_bytes is not None and ( + isinstance(small_checkpoint_max_bytes, bool) or small_checkpoint_max_bytes <= 0 + ): + raise ValueError( + "InstantTensor small-checkpoint threshold must be a positive integer" + ) + + if priority_weight_name_prefixes: + if not all( + isinstance(prefix, str) and prefix + for prefix in priority_weight_name_prefixes + ): + raise ValueError( + "InstantTensor priority weight prefixes must be non-empty strings" + ) + priority_names_by_file: dict[str, set[str]] = defaultdict(set) + tensor_files: Iterable[tuple[str, str]] + if indexed_tensor_files is not None: + tensor_files = indexed_tensor_files.items() + else: + # The target and every speculative-model source share one loader + # configuration. Header discovery lets an unindexed draft ignore + # target-only priorities without disabling them for the target. + discovered_tensor_files: list[tuple[str, str]] = [] + for filename in hf_weights_files: + absolute_filename = os.path.abspath(filename) + with safe_open(filename, framework="pt", device="cpu") as reader: + discovered_tensor_files.extend( + (name, absolute_filename) for name in reader.offset_keys() + ) + tensor_files = discovered_tensor_files + + for name, filename in tensor_files: + if not name.startswith(tuple(priority_weight_name_prefixes)): + continue + if weight_name_prefixes and not _matches_weight_name_prefixes( + name, weight_name_prefixes + ): + continue + priority_names_by_file[os.path.abspath(filename)].add(name) + if not priority_names_by_file: + logger.info_once( + "InstantTensor priority prefixes %s match no tensors in this " + "checkpoint source; all selected tensors use the normal " + "loading schedule.", + tuple(priority_weight_name_prefixes), + ) + else: + logger.info_once( + "InstantTensor loads %d tensors from %d/%d checkpoint shards " + "through CPU safetensors before GPU staging because they match " + "priority prefixes %s.", + sum(len(names) for names in priority_names_by_file.values()), + len(priority_names_by_file), + len(hf_weights_files), + tuple(priority_weight_name_prefixes), + ) + priority_names: set[str] = set() + for filename in hf_weights_files: + names = priority_names_by_file.get(os.path.abspath(filename)) + if not names: + continue + with safe_open(filename, framework="pt", device="cpu") as reader: + physical_names = list(reader.offset_keys()) + missing_names = names.difference(physical_names) + if missing_names: + raise RuntimeError( + "Safetensors index maps priority tensors to a shard " + "that does not contain them: " + f"{sorted(missing_names)[:3]} in {filename}" + ) + for name in physical_names: + if name not in names: + continue + priority_names.add(name) + yield name, reader.get_tensor(name) + expected_priority_names = set().union(*priority_names_by_file.values()) + if priority_names != expected_priority_names: + missing_names = expected_priority_names.difference(priority_names) + raise RuntimeError( + "InstantTensor priority tensors reference checkpoint shards " + f"outside the selected file set: {sorted(missing_names)[:3]}" + ) + else: + priority_names = set() + + if ( + small_checkpoint_max_bytes is not None + and indexed_tensor_files is None + and not priority_names + ): + checkpoint_size = _get_checkpoints_size_bytes(hf_weights_files) + if checkpoint_size <= small_checkpoint_max_bytes: + logger.info_once( + "Loading the %.2f GiB unindexed checkpoint through CPU " + "safetensors because it does not exceed the configured %.2f GiB " + "InstantTensor small-checkpoint threshold.", + checkpoint_size / 1024**3, + small_checkpoint_max_bytes / 1024**3, + ) + yield from safetensors_weights_iterator( + hf_weights_files, + use_tqdm_on_load, + weight_name_prefixes=weight_name_prefixes, + ) + return + try: world_group = get_world_group() except AssertionError: @@ -1353,38 +1483,336 @@ def instanttensor_weights_iterator( process_group = world_group.device_group if world_group.world_size > 1 else None device = current_platform.current_device() - - # copy=True yields tensors that own their memory, staying valid after the - # context exits or InstantTensor reuses its buffer. - with instanttensor.safe_open( - hf_weights_files, + copy_setting = os.getenv("INSTANTTENSOR_COPY", "1") + if copy_setting not in {"0", "1"}: + raise ValueError(f"INSTANTTENSOR_COPY must be 0 or 1, got {copy_setting!r}") + copy_tensors = copy_setting == "1" + + configured_buffer_size = os.getenv("INSTANTTENSOR_BUFFER_SIZE") + max_gpu_tensor_size: int | None = None + if configured_buffer_size is not None: + try: + max_gpu_tensor_size = int(configured_buffer_size) + except ValueError as e: + raise ValueError( + "INSTANTTENSOR_BUFFER_SIZE must be an integer number of bytes" + ) from e + if max_gpu_tensor_size <= 0: + raise ValueError("INSTANTTENSOR_BUFFER_SIZE must be greater than zero") + + restrict_before_io = ( + indexed_tensor_files is not None + or bool(weight_name_prefixes) + or max_gpu_tensor_size is not None + or bool(priority_names) + ) + open_kwargs: dict[str, Any] = {} + if restrict_before_io: + open_kwargs["load_now"] = False + + # copy=True is the safe default. copy=False exposes views into + # InstantTensor's reusable buffer for synchronous consumers. + instant_open = instanttensor.safe_open( + list(hf_weights_files), framework="pt", device=device, process_group=process_group, - copy=True, - ) as f: - # Track bytes so the bar reports load throughput (GB/s). - pbar = tqdm( - total=f.total_tensor_size, - desc="Loading safetensors using InstantTensor loader", - disable=not enable_tqdm(use_tqdm_on_load), - bar_format=_BAR_FORMAT, - position=tqdm._get_free_pos(), - unit="B", - unit_scale=True, - unit_divisor=1024, - mininterval=1.0, + copy=copy_tensors, + **open_kwargs, + ) + selection: _InstantTensorSelection | None = None + if restrict_before_io: + selection = _restrict_instanttensor_to_selected_ranges( + instant_open, + indexed_tensor_files=indexed_tensor_files, + weight_name_prefixes=weight_name_prefixes, + max_tensor_size=max_gpu_tensor_size, + excluded_tensor_names=priority_names, ) + if selection.selected_tensor_count == 0: + return + + cpu_fallbacks = selection.cpu_fallbacks if selection is not None else () + if not cpu_fallbacks: + with instant_open as f: + # Track bytes so the bar reports load throughput (GB/s). + pbar = tqdm( + total=f.total_tensor_size, + desc="Loading safetensors using InstantTensor loader", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + position=tqdm._get_free_pos(), + unit="B", + unit_scale=True, + unit_divisor=1024, + mininterval=1.0, + ) + try: + tensor: torch.Tensor | None = None + for name, tensor in f.tensors(): + pbar.update(tensor.numel() * tensor.element_size()) + if weight_name_prefixes and not _matches_weight_name_prefixes( + name, weight_name_prefixes + ): + continue + if not copy_tensors: + # Parameter loaders consume borrowed tensors synchronously. + # A deferred loader must own the tensor before advancing. + tensor._vllm_instanttensor_borrowed = True + yield name, tensor + # A generator frame retains its final DLPack-backed view. Drop it + # before the InstantTensor context releases the staging ring. + tensor = None + finally: + pbar.close() + return + + assert max_gpu_tensor_size is not None + assert selection is not None + logger.info_once( + "Loading %d tensors larger than the %d-byte InstantTensor buffer " + "through CPU safetensors", + len(cpu_fallbacks), + max_gpu_tensor_size, + ) + fallback_by_position = {fallback.position: fallback for fallback in cpu_fallbacks} + with ExitStack() as stack: + fallback_readers: dict[str, Any] = {} + for cpu_fallback in cpu_fallbacks: + if cpu_fallback.filename not in fallback_readers: + fallback_readers[cpu_fallback.filename] = stack.enter_context( + safe_open(cpu_fallback.filename, framework="pt", device="cpu") + ) + + pbar = None + gpu_tensors: Any = iter(()) + if selection.gpu_tensor_count: + instant_reader = stack.enter_context(instant_open) + pbar = tqdm( + total=instant_reader.total_tensor_size, + desc="Loading safetensors using InstantTensor loader", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + position=tqdm._get_free_pos(), + unit="B", + unit_scale=True, + unit_divisor=1024, + mininterval=1.0, + ) + gpu_tensors = iter(instant_reader.tensors()) + try: - for name, tensor in f.tensors(): - pbar.update(tensor.numel() * tensor.element_size()) - if weight_name_prefixes and not _matches_weight_name_prefixes( - name, weight_name_prefixes - ): + gpu_tensor: torch.Tensor | None = None + for position in range(selection.selected_tensor_count): + scheduled_fallback = fallback_by_position.get(position) + if scheduled_fallback is not None: + yield ( + scheduled_fallback.name, + fallback_readers[scheduled_fallback.filename].get_tensor( + scheduled_fallback.name + ), + ) continue - yield name, tensor + try: + name, gpu_tensor = next(gpu_tensors) + except StopIteration as e: + raise RuntimeError( + "InstantTensor produced fewer tensors than the selected " + "checkpoint schedule" + ) from e + if pbar is not None: + pbar.update(gpu_tensor.numel() * gpu_tensor.element_size()) + if not copy_tensors: + gpu_tensor._vllm_instanttensor_borrowed = True + yield name, gpu_tensor + + if next(gpu_tensors, None) is not None: + raise RuntimeError( + "InstantTensor produced more tensors than the selected " + "checkpoint schedule" + ) + gpu_tensor = None finally: - pbar.close() + if pbar is not None: + pbar.close() + + +def _restrict_instanttensor_to_selected_ranges( + instant_open: Any, + *, + indexed_tensor_files: dict[str, str] | None, + weight_name_prefixes: Sequence[str] | None, + max_tensor_size: int | None = None, + excluded_tensor_names: set[str] | None = None, +) -> _InstantTensorSelection: + """Restrict an InstantTensor loader to requested checkpoint ranges. + + Args: + instant_open: Unopened InstantTensor safetensors loader whose metadata + can be restricted before I/O starts. + indexed_tensor_files: Mapping from tensor names to their canonical + checkpoint files. ``None`` accepts tensors from every input file. + weight_name_prefixes: Model parameter prefixes selected for loading. + ``None`` or an empty sequence accepts every prefix. + max_tensor_size: Largest tensor payload, in bytes, retained in the GPU + staging path. Larger selected tensors are assigned to CPU loading. + excluded_tensor_names: Tensor names already emitted by an earlier + loading pass. Excluded names are removed before InstantTensor opens + its GPU staging context. + + Returns: + Selected GPU tensor count, total selected tensor count, and CPU + fallbacks with their positions in checkpoint order. + """ + required_attrs = ( + "filename", + "ordered_tensor_metadatas", + "tensor_offsets", + "tensor_sizes", + "total_tensor_size", + "tensor_name_to_index", + "loader_handle", + "_determine_buffer_size", + ) + missing = [name for name in required_attrs if not hasattr(instant_open, name)] + if missing: + raise RuntimeError( + "Installed InstantTensor does not expose the metadata layout needed " + f"for index-aware loading (missing: {', '.join(missing)})" + ) + if instant_open.loader_handle is not None: + raise RuntimeError("InstantTensor selection must be applied before opening I/O") + + filenames = list(instant_open.filename) + metadata = list(instant_open.ordered_tensor_metadatas) + offsets = list(instant_open.tensor_offsets) + selected_filenames: list[str] = [] + selected_metadata: list[tuple[str, dict[str, Any]]] = [] + selected_offsets: list[tuple[int, int]] = [] + cpu_fallbacks: list[_InstantTensorCpuFallback] = [] + selected_position = 0 + excluded_selected_count = 0 + metadata_pos = 0 + offset_pos = 0 + + for filename in filenames: + filename_abs = os.path.abspath(filename) + with safe_open(filename, framework="pt") as physical_file: + physical_names = list(physical_file.offset_keys()) + tensor_count = len(physical_names) + file_metadata = metadata[metadata_pos : metadata_pos + tensor_count] + file_offsets = offsets[offset_pos : offset_pos + tensor_count + 1] + if len(file_metadata) != tensor_count or len(file_offsets) != tensor_count + 1: + raise RuntimeError( + "InstantTensor metadata/offset count does not match safetensors " + f"header for {filename}" + ) + if [name for name, _ in file_metadata] != physical_names: + raise RuntimeError( + "InstantTensor tensor order does not match safetensors offset order " + f"for {filename}" + ) + + keep = [] + for item_index, name in enumerate(physical_names): + indexed_here = indexed_tensor_files is None or ( + indexed_tensor_files.get(name) == filename_abs + ) + prefix_matches = not weight_name_prefixes or _matches_weight_name_prefixes( + name, weight_name_prefixes + ) + otherwise_selected = indexed_here and prefix_matches + excluded_here = ( + otherwise_selected + and excluded_tensor_names is not None + and name in excluded_tensor_names + ) + excluded_selected_count += int(excluded_here) + selected_here = otherwise_selected and not excluded_here + item_data_offsets = file_metadata[item_index][1]["data_offsets"] + tensor_size = int(item_data_offsets[1]) - int(item_data_offsets[0]) + use_cpu_fallback = ( + selected_here + and max_tensor_size is not None + and tensor_size > max_tensor_size + ) + keep.append(selected_here and not use_cpu_fallback) + if selected_here: + if use_cpu_fallback: + cpu_fallbacks.append( + _InstantTensorCpuFallback( + position=selected_position, + name=name, + filename=filename, + ) + ) + selected_position += 1 + + run_start = 0 + while run_start < tensor_count: + while run_start < tensor_count and not keep[run_start]: + run_start += 1 + if run_start == tensor_count: + break + run_end = run_start + 1 + while run_end < tensor_count and keep[run_end]: + run_end += 1 + + logical_file_index = len(selected_filenames) + selected_filenames.append(filename) + selected_metadata.extend(file_metadata[run_start:run_end]) + selected_offsets.extend( + (logical_file_index, int(file_offsets[index][1])) + for index in range(run_start, run_end) + ) + selected_offsets.append((logical_file_index, int(file_offsets[run_end][1]))) + run_start = run_end + + metadata_pos += tensor_count + offset_pos += tensor_count + 1 + + if metadata_pos != len(metadata) or offset_pos != len(offsets): + raise RuntimeError( + "InstantTensor layout contains unaccounted metadata or offsets" + ) + if not selected_metadata and not cpu_fallbacks and excluded_selected_count == 0: + raise RuntimeError("InstantTensor index/prefix selection matched no tensors") + + selected_names = [name for name, _ in selected_metadata] + if len(selected_names) != len(set(selected_names)): + raise RuntimeError( + "InstantTensor index-aware selection still contains duplicate tensor names" + ) + fallback_names = [fallback.name for fallback in cpu_fallbacks] + if len(fallback_names) != len(set(fallback_names)): + raise RuntimeError( + "InstantTensor CPU fallback selection contains duplicate tensor names" + ) + overlap = set(selected_names).intersection(fallback_names) + if overlap: + raise RuntimeError( + f"InstantTensor GPU and CPU selections overlap: {sorted(overlap)[:3]}" + ) + selected_sizes = [ + int(item["data_offsets"][1]) - int(item["data_offsets"][0]) + for _, item in selected_metadata + ] + if selected_metadata: + instant_open.filename = selected_filenames + instant_open.ordered_tensor_metadatas = selected_metadata + instant_open.tensor_name_to_index = { + name: index for index, name in enumerate(selected_names) + } + instant_open.tensor_offsets = selected_offsets + instant_open.tensor_sizes = selected_sizes + instant_open.total_tensor_size = sum(selected_sizes) + instant_open._determine_buffer_size(None) + + return _InstantTensorSelection( + gpu_tensor_count=len(selected_metadata), + selected_tensor_count=selected_position, + cpu_fallbacks=tuple(cpu_fallbacks), + ) def pt_weights_iterator( diff --git a/vllm/model_executor/weight_transfer.py b/vllm/model_executor/weight_transfer.py index a8cbe230341b..25be5b1ed178 100644 --- a/vllm/model_executor/weight_transfer.py +++ b/vllm/model_executor/weight_transfer.py @@ -52,6 +52,11 @@ def copy_weight(destination: torch.Tensor, source: torch.Tensor) -> torch.Tensor return destination +def is_weight_transfer_active() -> bool: + """Whether checkpoint copies use a scoped writer with deferred completion.""" + return _writer.get() is not None + + def flush_weight_transfers() -> None: """Complete queued checkpoint reads before numerical weight preparation.""" writer = _writer.get() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index eb10ea6f28bd..587a105ae499 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -474,6 +474,12 @@ def init_device(self): # Set random seed. set_random_seed(self.model_config.seed) + if current_platform.is_cuda(): + # cuBLAS allocates outside the Torch cache. Initialize its + # per-thread handle before weights can fill physical memory, + # and include its retained storage in the admission snapshot. + torch.cuda.current_blas_handle() + # Now take memory snapshot after NCCL is initialized gc.collect() torch.accelerator.empty_cache() From 88904938a59062d3698fc25b15fe743fab3db6a4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:41:57 +0000 Subject: [PATCH 174/365] Preserve parallel-draft precision and mixed DCP cache geometry Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 (1M context) --- .../test_speculative_draft_hf_overrides.py | 31 + tests/kernels/core/test_rotary_embedding.py | 82 ++ .../quantization/test_marlin_tile_padding.py | 13 + tests/kernels/test_bf16_linear_staging.py | 149 ++++ .../kimi_k3/test_dspark_compact_rope.py | 110 +++ tests/models/kimi_k3/test_dspark_dcp_cache.py | 78 ++ tests/models/test_dflash2_mla.py | 148 ++++ tests/models/test_dspark_mla.py | 386 +++++++++- tests/models/test_qwen3_dflash_streaming.py | 368 +++++++++ tests/models/test_qwen3_dspark_tp_padding.py | 102 +++ .../test_dflash2_mla_config.py | 103 +++ .../test_speculators_dspark_config.py | 48 ++ tests/v1/core/test_kv_cache_utils.py | 192 +++++ .../offloading_connector/test_scheduler.py | 83 ++- tests/v1/kv_connector/unit/test_config.py | 19 + tests/v1/spec_decode/test_dflash2.py | 111 +++ tests/v1/spec_decode/test_dflash_dcp.py | 57 +- tests/v1/worker/test_attn_utils.py | 105 ++- tests/v1/worker/test_dspark_greedy_markov.py | 53 ++ tests/v1/worker/test_dspark_utils.py | 68 ++ tests/v1/worker/test_mamba_utils.py | 66 +- tests/v1/worker/test_workspace.py | 67 +- vllm/config/model.py | 9 +- vllm/config/speculative.py | 6 + vllm/config/vllm.py | 26 +- .../kv_transfer/kv_connector/v1/__init__.py | 4 + .../kv_transfer/kv_connector/v1/base.py | 28 +- .../kv_connector/v1/offloading/scheduler.py | 10 +- .../kv_connector/v1/offloading_connector.py | 14 +- .../kernels/linear/bf16_staging.py | 62 ++ .../kernels/linear/mxfp8/staged.py | 221 ++++++ .../quantization/utils/marlin_utils_fp8.py | 16 +- .../layers/rotary_embedding/compact.py | 82 ++ vllm/model_executor/models/config.py | 71 ++ vllm/model_executor/models/qwen3_dflash.py | 404 +++++++++- vllm/model_executor/models/qwen3_dspark.py | 32 +- vllm/model_executor/models/registry.py | 4 + vllm/models/kimi_k3/nvidia/dflash2_mla.py | 155 ++++ vllm/models/kimi_k3/nvidia/dspark_mla.py | 702 +++++++++++++++++- .../model_arch_config_convertor.py | 16 + vllm/v1/attention/backends/triton_attn.py | 3 + vllm/v1/core/kv_cache_utils.py | 54 +- vllm/v1/kv_cache_interface.py | 7 +- vllm/v1/worker/gpu/attn_utils.py | 23 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 5 +- .../gpu/spec_decode/dflash/speculator.py | 29 +- .../gpu/spec_decode/dspark/speculator.py | 28 +- .../v1/worker/gpu/spec_decode/dspark/utils.py | 16 +- vllm/v1/worker/mamba_utils.py | 24 +- 49 files changed, 4386 insertions(+), 104 deletions(-) create mode 100644 tests/kernels/test_bf16_linear_staging.py create mode 100644 tests/models/kimi_k3/test_dspark_compact_rope.py create mode 100644 tests/models/kimi_k3/test_dspark_dcp_cache.py create mode 100644 tests/models/test_dflash2_mla.py create mode 100644 tests/models/test_qwen3_dflash_streaming.py create mode 100644 tests/models/test_qwen3_dspark_tp_padding.py create mode 100644 tests/transformers_utils/test_dflash2_mla_config.py create mode 100644 vllm/model_executor/kernels/linear/bf16_staging.py create mode 100644 vllm/model_executor/kernels/linear/mxfp8/staged.py create mode 100644 vllm/model_executor/layers/rotary_embedding/compact.py create mode 100644 vllm/models/kimi_k3/nvidia/dflash2_mla.py diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 7c3acf19874f..ac8db68ffa98 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -86,6 +86,37 @@ def test_none_overrides_fall_back_to_arch_mapping(): assert composed is SpeculativeConfig.hf_config_override +@pytest.mark.cpu_test +@pytest.mark.parametrize("method", ["dspark", "dflash"]) +def test_draft_online_quantization_preserves_projection_exclusions(method): + from vllm.config.quantization import resolve_quantization_config + + overrides = {"linear": "mxfp8", "ignore": ["re:.*fused_qkv_a_proj$"]} + target = MagicMock( + model="target", max_model_len=128, quantization="mxfp4", hf_overrides={} + ) + with ( + patch( + "vllm.config.speculative.ModelConfig", + side_effect=RuntimeError("draft constructor boundary"), + ) as constructor, + pytest.raises(RuntimeError, match="draft constructor boundary"), + ): + SpeculativeConfig( + model="draft", + method=method, + num_speculative_tokens=7, + quantization="mxfp8", + quantization_config=overrides, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + assert constructor.call_args.kwargs["quantization_config"] == ( + resolve_quantization_config("mxfp8", overrides) + ) + assert target.quantization == "mxfp4" + + @pytest.mark.cpu_test def test_callable_overrides_reach_the_draft_config(): """A callable override (config-to-config transform) composes with the diff --git a/tests/kernels/core/test_rotary_embedding.py b/tests/kernels/core/test_rotary_embedding.py index 1cbb5dbd1881..67fce19c2ff3 100644 --- a/tests/kernels/core/test_rotary_embedding.py +++ b/tests/kernels/core/test_rotary_embedding.py @@ -77,3 +77,85 @@ def test_rotary_embedding_opcheck( query.flatten(start_dim=-2), key.flatten(start_dim=-2) if key is not None else None, ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("is_neox_style", [True, False]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@torch.inference_mode() +def test_compact_rotary_preserves_full_table_and_changed_input_graphs( + default_vllm_config, is_neox_style, dtype +): + """Global positions above 500k retain exact table values and rotated Q/K.""" + from vllm.model_executor.layers.rotary_embedding.compact import ( + CompactRotaryEmbedding, + ) + + capacity, head_dim, max_position = 4096, 128, 1048576 + with torch.device("cuda:0"): + reference = RotaryEmbedding( + head_dim, head_dim, max_position, 2000000, is_neox_style, dtype + ) + compact = CompactRotaryEmbedding( + head_dim, + head_dim, + max_position, + 2000000, + is_neox_style, + dtype, + capacity=capacity, + ) + positions = torch.randint(max_position, (capacity,), device="cuda:0") + positions[:8] = torch.tensor( + [0, 1, 4095, 500000, 524287, 524288, 950000, 1048575], device="cuda:0" + ) + query = torch.randn(capacity, head_dim * 2, dtype=dtype, device="cuda:0") + key = torch.randn(capacity, head_dim, dtype=dtype, device="cuda:0") + pointers = [buffer.data_ptr() for buffer in compact.buffers()] + assert ( + sum(t.nbytes for t in compact.buffers()) < reference.cos_sin_cache.nbytes / 100 + ) + for rows in (1, 8, 129, capacity): + q, k = query[:rows].clone(), key[:rows].clone() + expected_q, expected_k = reference.forward_cuda(positions[:rows], q, k) + q, k = query[:rows].clone(), key[:rows].clone() + actual_q, actual_k = compact.forward_cuda(positions[:rows], q, k) + torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=0) + torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=0) + torch.testing.assert_close( + compact.cos_sin_cache[:rows], + reference.cos_sin_cache[positions[:rows]], + rtol=0, + atol=0, + ) + + rows = 8 + q, k = query[:rows].clone(), key[:rows].clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + compact.forward_cuda(positions[:rows], q, k) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + compact.forward_cuda(positions[:rows], q, k) + positions[:rows].add_(123).remainder_(max_position) + expected_q, expected_k = reference.forward_cuda( + positions[:rows], query[:rows].clone(), key[:rows].clone() + ) + q.copy_(query[:rows]) + k.copy_(key[:rows]) + compact.cos_sin_cache.fill_(float("nan")) + compact.freqs_workspace.fill_(float("nan")) + allocated = torch.accelerator.memory_allocated() + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() == allocated + torch.testing.assert_close(q, expected_q, rtol=0, atol=0) + torch.testing.assert_close(k, expected_k, rtol=0, atol=0) + assert pointers == [buffer.data_ptr() for buffer in compact.buffers()] + + with pytest.raises(ValueError, match="exceed prepared capacity"): + compact.materialize( + torch.zeros(capacity + 1, device="cuda:0", dtype=torch.int64) + ) diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c96..1b48cc05b0ea 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -449,6 +449,19 @@ def test_mxfp8_marlin_padded_round_trip(shape): assert output.shape == (8, size_n) torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2) + padded_n, _ = marlin_repacked_nk(layer.weight, num_bits=8) + caller_output = torch.full((8, padded_n), float("nan"), dtype=dtype, device="cuda") + actual = apply_mxfp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=size_n, + size_k=size_k, + output=caller_output, + ) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + @pytest.mark.skipif( _gpu_marlin_unsupported(), diff --git a/tests/kernels/test_bf16_linear_staging.py b/tests/kernels/test_bf16_linear_staging.py new file mode 100644 index 000000000000..2daaac2de4d7 --- /dev/null +++ b/tests/kernels/test_bf16_linear_staging.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.kernels.linear.bf16_staging import Bf16InputAccumulator + + +def make_accumulator(device="cpu"): + weight = torch.arange(24, device=device, dtype=torch.bfloat16).reshape(3, 8) + layer = SimpleNamespace( + input_size=8, + weight=weight, + bias=None, + quant_method=SimpleNamespace( + apply=lambda layer, x, bias: torch.nn.functional.linear( + x, layer.weight, bias + ) + ), + ) + return Bf16InputAccumulator( + layer, torch.empty(129, 3, device=device, dtype=torch.bfloat16), 4 + ) + + +@pytest.mark.parametrize("rows", [1, 8, 129]) +def test_bf16_staging_preserves_bits_and_releases_slices(rows): + accumulator = make_accumulator() + source = torch.arange(rows * 8, dtype=torch.bfloat16).reshape(rows, 8) + expected = accumulator.layer.quant_method.apply(accumulator.layer, source, None) + accumulator.begin(rows) + reusable = torch.empty(rows, 4, dtype=torch.bfloat16) + for chunk in source.split(4, dim=-1): + reusable.copy_(chunk) + accumulator.append(reusable) + reusable.fill_(float("nan")) + torch.testing.assert_close(accumulator.values[:rows], source, rtol=0, atol=0) + torch.testing.assert_close(accumulator.finish(), expected, rtol=0, atol=0) + with pytest.raises(RuntimeError, match="incomplete"): + accumulator.finish() + + +def test_bf16_staging_rejects_partial_or_overfilled_inputs(): + accumulator = make_accumulator() + source = torch.ones(8, 4, dtype=torch.bfloat16) + for rows in (0, 130): + with pytest.raises(ValueError, match="capacity"): + accumulator.begin(rows) + with pytest.raises(ValueError, match="geometry"): + accumulator.append(source) + accumulator.begin(8) + with pytest.raises(ValueError, match="dtype"): + accumulator.append(source.float()) + accumulator.append(source) + with pytest.raises(RuntimeError, match="incomplete"): + accumulator.finish() + accumulator.append(source) + with pytest.raises(ValueError, match="geometry"): + accumulator.append(source) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_bf16_staging_graph_replays_changed_inputs(): + accumulator = make_accumulator("cuda") + source = torch.randn(8, 8, device="cuda", dtype=torch.bfloat16) + + def run(): + accumulator.begin(8) + for chunk in source.split(4, dim=-1): + accumulator.append(chunk) + return accumulator.finish() + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + pointer = accumulator.values.data_ptr() + for scale in (0.125, 3.0, -0.5): + source.mul_(scale) + graph.replay() + expected = accumulator.layer.quant_method.apply(accumulator.layer, source, None) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + assert accumulator.values.data_ptr() == pointer + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", [1, 8, 129, 4096]) +@torch.inference_mode() +def test_dflash_tp16_marlin_staging_matches_bf16_concatenation(rows): + from torch import nn + + from vllm.model_executor.kernels.linear.mxfp8.marlin import ( + MarlinMxfp8LinearKernel, + ) + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + prepare_mxfp8_layer_for_marlin, + ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + torch.manual_seed(231) + width, slices, local_output = 7168, 6, 448 + layer = nn.Module() + layer.input_size = layer.input_size_per_partition = width * slices + layer.output_size_per_partition = local_output + layer.bias = None + weight = ( + torch.randn(local_output, width * slices, device="cuda", dtype=torch.bfloat16) + / 64 + ) + values, scales = mxfp8_e4m3_quantize(weight) + layer.weight = nn.Parameter(values, requires_grad=False) + layer.weight_scale = nn.Parameter(scales, requires_grad=False) + prepare_mxfp8_layer_for_marlin(layer) + kernel = object.__new__(MarlinMxfp8LinearKernel) + layer.quant_method = SimpleNamespace(apply=kernel.apply_weights) + accumulator = Bf16InputAccumulator( + layer, + torch.empty(rows, local_output, device="cuda", dtype=torch.bfloat16), + width, + ) + source = torch.randn(rows, width * slices, device="cuda", dtype=torch.bfloat16) + + def run(): + accumulator.begin(rows) + for part in source.split(width, dim=-1): + accumulator.append(part) + return accumulator.finish() + + expected = kernel.apply_weights(layer, source) + actual = run() + torch.testing.assert_close(accumulator.values, source, rtol=0, atol=0) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + for scale in (-0.25, 2.0): + source.mul_(scale) + graph.replay() + expected = kernel.apply_weights(layer, source) + torch.testing.assert_close(result, expected, rtol=0, atol=0) diff --git a/tests/models/kimi_k3/test_dspark_compact_rope.py b/tests/models/kimi_k3/test_dspark_compact_rope.py new file mode 100644 index 000000000000..50278686c37a --- /dev/null +++ b/tests/models/kimi_k3/test_dspark_compact_rope.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +from torch import nn + +from vllm.models.kimi_k3.nvidia import dspark_mla + + +class _Rotary(nn.Module): + def __init__(self) -> None: + super().__init__() + self.scaling_factor = 2.0 + self.mscale = 1.25 + self.head_size = 4 + self.is_neox_style = False + self.register_buffer( + "cos_sin_cache", + torch.arange(64, dtype=torch.float32).view(16, 4), + ) + + def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: + assert scaling_factor == self.scaling_factor + return torch.tensor([0.5, 0.25], dtype=torch.float32) + + +class _Attention(nn.Module): + def __init__(self, rotary: _Rotary) -> None: + super().__init__() + self.rotary_emb = rotary + + +class _Layer(nn.Module): + def __init__(self, rotary: _Rotary) -> None: + super().__init__() + self.self_attn = _Attention(rotary) + + +def _draft_model(rotary: _Rotary) -> dspark_mla.K3DSparkModel: + model = object.__new__(dspark_mla.K3DSparkModel) + nn.Module.__init__(model) + model.layers = nn.ModuleList([_Layer(rotary), _Layer(rotary)]) + model._max_num_context_tokens = 8 + return model + + +def test_compact_rope_materializes_requested_global_positions() -> None: + positions = torch.tensor([2, 7], dtype=torch.int64) + inv_freq = torch.tensor([0.5, 0.25], dtype=torch.float32) + freqs = torch.empty((4, 2), dtype=torch.float32) + cache = torch.empty((4, 4), dtype=torch.float32) + + result = dspark_mla._fill_compact_rope_cache( + positions, + inv_freq, + freqs, + cache, + mscale=1.25, + ) + + expected_freqs = positions[:, None] * inv_freq[None, :] + expected = torch.cat((expected_freqs.cos(), expected_freqs.sin()), dim=-1) * 1.25 + torch.testing.assert_close(result, expected) + assert result.untyped_storage().data_ptr() == cache.untyped_storage().data_ptr() + + +def test_compact_rope_rejects_more_positions_than_workspace() -> None: + with pytest.raises(ValueError, match="workspace is too small"): + dspark_mla._fill_compact_rope_cache( + torch.arange(3), + torch.ones(2), + torch.empty((2, 2)), + torch.empty((2, 4)), + mscale=1.0, + ) + + +def test_compact_rope_releases_unshared_full_position_table() -> None: + rotary = _Rotary() + original_bytes = rotary.cos_sin_cache.nbytes + model = _draft_model(rotary) + + model._init_compact_rope() + model._compact_rope_enabled = True + remapped_positions, compact_cache = model._get_rope_inputs( + torch.tensor([2, 7], dtype=torch.int64) + ) + + assert original_bytes == 256 + assert rotary.cos_sin_cache.shape == (0, 4) + torch.testing.assert_close(remapped_positions, torch.tensor([0, 1])) + expected_freqs = torch.tensor([[1.0, 0.5], [3.5, 1.75]]) + expected_cache = ( + torch.cat((expected_freqs.cos(), expected_freqs.sin()), dim=-1) * 1.25 + ) + torch.testing.assert_close(compact_cache, expected_cache) + + +def test_compact_rope_retains_target_owned_position_table() -> None: + rotary = _Rotary() + target = nn.Module() + target.rotary = rotary + model = _draft_model(rotary) + + with dspark_mla.protect_k3_compact_rope_sources(target): + model._init_compact_rope() + + assert rotary.cos_sin_cache.shape == (16, 4) + assert not dspark_mla._COMPACT_ROPE_PROTECTED_IDS diff --git a/tests/models/kimi_k3/test_dspark_dcp_cache.py b/tests/models/kimi_k3/test_dspark_dcp_cache.py new file mode 100644 index 000000000000..788ce094c88a --- /dev/null +++ b/tests/models/kimi_k3/test_dspark_dcp_cache.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.models.kimi_k3.nvidia import mla + + +def _cache_spec( + monkeypatch, + *, + non_causal: bool, + dcp_size: int, + draft_kv_window: int = 0, +): + attention = object.__new__(mla.MultiHeadLatentAttention) + attention.kv_cache_dtype = "auto" + attention.head_size = 640 + attention.non_causal_multi_token_decode = non_causal + attention.draft_kv_window = draft_kv_window + config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=64), + model_config=object(), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_size, + ), + ) + monkeypatch.setattr( + mla, + "kv_cache_dtype_str_to_dtype", + lambda _cache_dtype, _model_config: torch.bfloat16, + ) + monkeypatch.setattr(mla, "get_kv_quant_mode", lambda _cache_dtype: None) + return mla.MultiHeadLatentAttention.get_kv_cache_spec(attention, config) + + +def test_external_k3_dspark_cache_is_replicated_under_target_dcp( + monkeypatch, +) -> None: + monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + + spec = _cache_spec(monkeypatch, non_causal=True, dcp_size=16) + + assert spec.dcp_replicated + + +def test_explicit_draft_sharding_disables_cache_replication(monkeypatch) -> None: + monkeypatch.setenv("VLLM_DCP_SHARD_DRAFT", "1") + + spec = _cache_spec(monkeypatch, non_causal=True, dcp_size=16) + + assert not spec.dcp_replicated + + +def test_target_k3_cache_is_not_marked_as_draft_replication(monkeypatch) -> None: + monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + + spec = _cache_spec(monkeypatch, non_causal=False, dcp_size=16) + + assert not spec.dcp_replicated + + +def test_bounded_dspark_cache_preserves_non_causal_mode(monkeypatch) -> None: + monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + + spec = _cache_spec( + monkeypatch, + non_causal=True, + dcp_size=16, + draft_kv_window=65_536, + ) + + assert isinstance(spec, mla.SlidingWindowMLASpec) + assert spec.sliding_window == 65_536 + assert spec.non_causal_multi_token_decode + assert spec.dcp_replicated diff --git a/tests/models/test_dflash2_mla.py b/tests/models/test_dflash2_mla.py new file mode 100644 index 000000000000..f23d5b8dcb1e --- /dev/null +++ b/tests/models/test_dflash2_mla.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm.model_executor.models.registry import ModelRegistry +from vllm.models.kimi_k3.nvidia.dflash2_mla import ( + DFlash2KimiK3DecoderLayer, + DFlash2KimiK3ForCausalLM, + DFlash2KimiK3Model, +) +from vllm.models.kimi_k3.nvidia.dspark_mla import ( + K3DSparkForCausalLM, + K3DSparkModel, +) + +pytestmark = pytest.mark.cpu_test + + +def test_mla_dflash2_registry_and_head_contract(): + assert ( + ModelRegistry._try_load_model_cls("DFlash2KimiK3Model") + is DFlash2KimiK3ForCausalLM + ) + assert DFlash2KimiK3ForCausalLM.model_cls is DFlash2KimiK3Model + assert DFlash2KimiK3Model.decoder_layer_cls is DFlash2KimiK3DecoderLayer + assert not DFlash2KimiK3Model.has_markov_head + assert K3DSparkModel.has_markov_head + assert not DFlash2KimiK3ForCausalLM.has_own_embed_tokens + assert not DFlash2KimiK3ForCausalLM.has_own_lm_head + + +def test_mla_dflash2_maps_context_and_norm_checkpoint_names(monkeypatch): + names = [ + "fc.weight", + "hidden_norm.weight", + "norm.weight", + "layers.0.self_attn.q_a_proj.weight", + "candidate_selector.predecessor_codebook", + ] + tensor = torch.ones(1) + expected = [ + "context_proj.weight", + "context_norm.weight", + "final_norm.weight", + *names[3:], + ] + + def load(self, weights): + weights = list(weights) + assert [name for name, _ in weights] == expected + assert all(weight is tensor for _, weight in weights) + return set(expected) + + monkeypatch.setattr(K3DSparkForCausalLM, "load_weights", load) + model = object.__new__(DFlash2KimiK3ForCausalLM) + assert model.load_weights((name, tensor) for name in names) == set(expected) + + +def test_mla_dflash2_uses_distributed_topk_without_full_logit_gather(): + processor = Mock() + head = object() + states = torch.zeros(3, 8) + model = SimpleNamespace( + lm_head=head, + model=SimpleNamespace(candidate_selector=SimpleNamespace(top_k=16)), + candidate_logits_processor=processor, + ) + result = DFlash2KimiK3ForCausalLM.compute_candidates(model, states) + processor.get_top_k_tokens.assert_called_once_with(head, states, 16) + assert result is processor.get_top_k_tokens.return_value + + +@pytest.mark.parametrize("index,expected", [(0, 4096), (3, 4096), (4, 0)]) +def test_mla_dflash2_preserves_checkpoint_attention_windows( + default_vllm_config, monkeypatch, index, expected +): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + monkeypatch.setenv("VLLM_DSPARK_DRAFT_KV_WINDOW", "2048") + config = SimpleNamespace( + is_causal=False, + layer_types=["sliding_attention"] * 4 + ["full_attention"], + dflash_config={}, + sliding_window=4096, + ) + assert DFlash2KimiK3DecoderLayer.get_attention_window(config, index) == expected + + +def test_mla_dflash2_rejects_unsupported_causal_layer(): + config = SimpleNamespace(is_causal=True, layer_types=None, dflash_config={}) + with pytest.raises(ValueError, match="requires non-causal"): + DFlash2KimiK3DecoderLayer.get_attention_window(config, 0) + + +def test_candidate_selector_does_not_prepare_markov_argmax(): + model = SimpleNamespace(model=SimpleNamespace(markov_head=None)) + assert not K3DSparkForCausalLM.supports_local_draft_argmax(model) + + +@pytest.mark.parametrize("tp", [8, 9, 10, 12, 16]) +def test_padded_draft_mlp_preserves_unquantized_projection( + default_vllm_config, monkeypatch, tp +): + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + from vllm.models.kimi_k3.nvidia import model as kimi_model + + monkeypatch.setattr(kimi_model, "get_tensor_model_parallel_world_size", lambda: tp) + monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: tp) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: tp) + generator = torch.Generator().manual_seed(717) + hidden, intermediate = 64, 160 + gate = torch.randn(intermediate, hidden, generator=generator) / hidden**0.5 + up = torch.randn(intermediate, hidden, generator=generator) / hidden**0.5 + down = torch.randn(hidden, intermediate, generator=generator) / intermediate**0.5 + x = torch.randn(3, hidden, generator=generator) + functional = torch.nn.functional + expected = functional.linear( + functional.silu(functional.linear(x, gate)) * functional.linear(x, up), down + ) + result = torch.zeros_like(expected) + for rank in range(tp): + monkeypatch.setattr( + linear, "get_tensor_model_parallel_rank", lambda rank=rank: rank + ) + monkeypatch.setattr( + parameter, "get_tensor_model_parallel_rank", lambda rank=rank: rank + ) + mlp = kimi_model.KimiMLP( + hidden, intermediate, "silu", reduce_results=False, pad_intermediate=True + ) + weight = mlp.gate_up_proj.weight + weight.weight_loader(weight, gate, 0) + weight.weight_loader(weight, up, 1) + weight_down = mlp.down_proj.weight + weight_down.weight_loader(weight_down, down) + local_gate, local_up = functional.linear(x, weight).chunk(2, dim=-1) + result += functional.linear(functional.silu(local_gate) * local_up, weight_down) + local_width = weight_down.shape[1] + valid = max(0, min(local_width, intermediate - rank * local_width)) + assert torch.count_nonzero(weight_down[:, valid:]) == 0 + assert torch.count_nonzero(weight[valid:local_width]) == 0 + assert torch.count_nonzero(weight[local_width + valid :]) == 0 + torch.testing.assert_close(result, expected, atol=2e-6, rtol=2e-5) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 288d8a778806..360ac713296d 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -97,6 +98,61 @@ def fail_collective(*args, **kwargs): assert bias.shape == (2, 128) +@pytest.mark.parametrize("replicate_w1", [False, True]) +def test_dspark_sharded_markov_weights_and_local_sampling( + monkeypatch, default_vllm_config, replicate_w1 +): + from vllm.model_executor.layers import vocab_parallel_embedding as embedding + + monkeypatch.setenv("VLLM_DSPARK_SHARD_MARKOV_HEAD", "1") + monkeypatch.setenv("VLLM_DSPARK_REPLICATE_MARKOV_W1", str(int(replicate_w1))) + monkeypatch.setattr(embedding, "get_tensor_model_parallel_rank", lambda: 11) + monkeypatch.setattr(embedding, "get_tensor_model_parallel_world_size", lambda: 12) + head = DSparkMarkovHead(190, 190, 8, prefix="markov_head") + assert head.markov_w2.weight.shape == (16, 8) + assert head.markov_w1.weight.shape == (190 if replicate_w1 else 16, 8) + processor = LogitsProcessor(190) + assert head.local_bias(torch.ones(2, 8), processor).shape == (2, 16) + + model = K3DSparkForCausalLM.__new__(K3DSparkForCausalLM) + nn.Module.__init__(model) + model.model = nn.Module() + model.model.markov_head = head + model.lm_head = embedding.ParallelLMHead(190, 8, prefix="lm_head") + model.logits_processor = processor + model._argmax_enabled = True + model._argmax_capacity = 8 + model._argmax_plan = object() + model._argmax_output = torch.empty(8, dtype=torch.int64) + calls = [] + + def sample(base, bias, out, *, plan): + assert plan is model._argmax_plan + calls.append((base.clone(), bias.clone())) + return out.fill_(7) + + model._argmax_runtime = SimpleNamespace(fused_add_argmax=sample) + assert model.supports_local_draft_argmax() + base = torch.zeros(2, 16, dtype=torch.bfloat16) + bias = torch.zeros_like(base) + base[:, -2:] = 1000 + assert model.sample_local_draft_logits(base, bias).tolist() == [7, 7] + assert torch.isneginf(calls[0][0][:, -2:]).all() + assert torch.isneginf(calls[0][1][:, -2:]).all() + + processor.scale = 0.5 + assert not model.supports_local_draft_argmax() + processor.scale = 1 + head.markov_w2.num_embeddings_padded += 64 + assert not model.supports_local_draft_argmax() + + +def test_dspark_sharded_head_rejects_gathered_topk(monkeypatch, default_vllm_config): + monkeypatch.setenv("VLLM_DSPARK_SHARD_MARKOV_HEAD", "1") + with pytest.raises(ValueError, match="top-k requires replicated"): + DSparkMarkovHead(128, 128, 8, "head", retain_weight_for_gather=True) + + @pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) def test_v41_dspark_retains_each_markov_embedding_during_graph_replay( default_vllm_config, @@ -164,6 +220,7 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): class DummyModule(nn.Module): def __init__(self, *args, **kwargs): super().__init__() + self.weight = nn.Parameter(torch.empty(0)) def make_markov_head(*args, **kwargs): markov_head_calls.append((args, kwargs)) @@ -177,7 +234,7 @@ def make_context_kv_proj(*args, **kwargs): monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule) monkeypatch.setattr(dspark_mla, "MergedColumnParallelLinear", make_context_kv_proj) monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule) - monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule) + monkeypatch.setattr(K3DSparkModel, "decoder_layer_cls", DummyModule) monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head) config = SimpleNamespace( @@ -216,6 +273,331 @@ def make_context_kv_proj(*args, **kwargs): ] +@pytest.mark.cpu_test +@pytest.mark.parametrize( + ("tp_size", "expect_sharded"), + [(1, False), (8, True), (12, False), (16, True)], +) +def test_k3_dspark_context_projection_uses_divisible_tp_geometry( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + expect_sharded: bool, +) -> None: + context_projection_calls = [] + + class DummyModule(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.weight = nn.Parameter(torch.empty(0)) + + def make_sharded_projection(*args, **kwargs): + context_projection_calls.append((args, kwargs)) + return DummyModule() + + monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None) + monkeypatch.setattr(dspark_mla, "ColumnParallelLinear", make_sharded_projection) + monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule) + monkeypatch.setattr(dspark_mla, "MergedColumnParallelLinear", DummyModule) + monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule) + monkeypatch.setattr(K3DSparkModel, "decoder_layer_cls", DummyModule) + monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", DummyModule) + + config = SimpleNamespace( + target_hidden_size=7168, + num_target_layers=5, + hidden_size=7168, + kv_lora_rank=512, + qk_rope_head_dim=64, + rms_norm_eps=1e-6, + num_hidden_layers=1, + vocab_size=163840, + draft_vocab_size=163840, + markov_rank=256, + ) + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=config) + ), + parallel_config=SimpleNamespace(tensor_parallel_size=tp_size), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16), + ) + + model = K3DSparkModel( + vllm_config=vllm_config, + start_layer_id=0, + prefix="model", + ) + + assert model.context_proj_sharded is expect_sharded + assert model._streamed_context_states.dtype == model.context_proj.weight.dtype + assert model._streamed_context_states.device == model.context_proj.weight.device + assert bool(context_projection_calls) is expect_sharded + if expect_sharded: + args, kwargs = context_projection_calls[0] + assert args == (7168 * 5, 7168) + assert kwargs["gather_output"] is True + assert kwargs["quant_config"] is None + + +@pytest.mark.cpu_test +def test_k3_dspark_streams_auxiliary_projection_into_tp_local_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + torch.manual_seed(0) + model = object.__new__(K3DSparkModel) + nn.Module.__init__(model) + model.config = SimpleNamespace( + target_hidden_size=4, + hidden_size=2, + target_layer_ids=(0, 1), + ) + model.context_proj_sharded = True + model.context_proj = nn.Linear(8, 2, bias=False, dtype=torch.bfloat16) + model.context_norm = SimpleNamespace( + weight=torch.tensor([0.75, 1.25], dtype=torch.bfloat16), + variance_epsilon=1e-5, + ) + model.context_kv_proj = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16) + model._max_num_context_tokens = 1024 + model._streamed_aux_layer_ids = (1, 2) + model._streamed_aux_scratch = None + model._streamed_aux_tokens = 0 + model._streamed_aux_index = 0 + model._streamed_aux_generation = 0 + model._completed_stream_generation = 0 + model._consumed_stream_generation = 0 + model._completed_stream_result = None + model._context_local_width = 2 + model._context_local_start = 0 + model.register_buffer( + "_streamed_context_states", + torch.empty(1024, 2, dtype=torch.bfloat16), + persistent=False, + ) + scratch = torch.empty(1024, 4, dtype=torch.bfloat16) + model.bind_auxiliary_stream_scratch(scratch) + monkeypatch.setattr( + dspark_mla, + "tensor_model_parallel_all_reduce_in_place", + lambda tensor: tensor, + ) + + first = torch.randn(1024, 4, dtype=torch.bfloat16) + second = torch.randn(1024, 4, dtype=torch.bfloat16) + residual = torch.randn(1024, 4, dtype=torch.bfloat16) + combined = torch.cat((first, second + residual), dim=-1) + projected = torch.nn.functional.linear(combined, model.context_proj.weight) + expected = projected * torch.rsqrt( + projected.float().square().mean(dim=-1, keepdim=True) + 1e-5 + ).to(projected.dtype) + expected *= model.context_norm.weight + + assert model.can_stream_auxiliary_states((1, 2), first) + with torch.inference_mode(): + model.begin_auxiliary_stream(first) + model.accumulate_auxiliary_state(first, None) + model.accumulate_auxiliary_state(second, residual) + output = model.finish_auxiliary_stream() + + # Streaming replaces one BF16 GEMM with an ordered sum of BF16 GEMMs, so + # cancellation can move individual low-magnitude elements by several + # ulps. Preserve the projected vector rather than requiring elementwise + # equality between the two legal BF16 accumulation orders. + mean_absolute_error = (output.float() - expected.float()).abs().mean() + cosine_similarity = torch.nn.functional.cosine_similarity( + output.float().flatten(), expected.float().flatten(), dim=0 + ) + assert mean_absolute_error.item() < 2e-3 + assert cosine_similarity.item() > 0.9999 + assert model.is_streamed_context_states([output]) + assert not model.is_streamed_context_states([output]) + assert output.data_ptr() == model._streamed_context_states.data_ptr() + + with torch.inference_mode(): + model.begin_auxiliary_stream(first) + model.accumulate_auxiliary_state(first, None) + model.accumulate_auxiliary_state(second, residual) + next_output = model.finish_auxiliary_stream() + + assert not model.is_streamed_context_states([output]) + assert model.is_streamed_context_states([next_output]) + + +@pytest.mark.cpu_test +def test_k3_dspark_streamed_auxiliary_normalization_uses_global_hidden_width( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = object.__new__(K3DSparkModel) + nn.Module.__init__(model) + model.config = SimpleNamespace( + target_hidden_size=4, + hidden_size=4, + target_layer_ids=(0, 1), + ) + model.context_proj_sharded = True + model.context_proj = nn.Linear(8, 2, bias=False, dtype=torch.bfloat16) + model.context_norm = SimpleNamespace( + weight=torch.tensor([0.5, 1.0, 1.5, 2.0], dtype=torch.bfloat16), + variance_epsilon=1e-5, + ) + model.context_kv_proj = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16) + model._max_num_context_tokens = 1024 + model._streamed_aux_layer_ids = (1, 2) + model._streamed_aux_scratch = None + model._streamed_aux_tokens = 0 + model._streamed_aux_index = 0 + model._streamed_aux_generation = 0 + model._completed_stream_generation = 0 + model._consumed_stream_generation = 0 + model._completed_stream_result = None + model._context_local_width = 2 + model._context_local_start = 2 + model.register_buffer( + "_streamed_context_states", + torch.empty(1024, 2, dtype=torch.bfloat16), + persistent=False, + ) + model.bind_auxiliary_stream_scratch(torch.empty(1024, 4, dtype=torch.bfloat16)) + + first = torch.randn(1024, 4, dtype=torch.bfloat16) + second = torch.randn(1024, 4, dtype=torch.bfloat16) + combined = torch.cat((first, second), dim=-1) + local_projection = torch.nn.functional.linear(combined, model.context_proj.weight) + other_rank_projection = torch.randn_like(local_projection) + + def add_other_rank_squared_norm(tensor: torch.Tensor) -> torch.Tensor: + tensor.add_(other_rank_projection.float().square().sum(dim=-1, keepdim=True)) + return tensor + + monkeypatch.setattr( + dspark_mla, + "tensor_model_parallel_all_reduce_in_place", + add_other_rank_squared_norm, + ) + expected_scale = torch.rsqrt( + ( + local_projection.float().square().sum(dim=-1, keepdim=True) + + other_rank_projection.float().square().sum(dim=-1, keepdim=True) + ) + / model.config.hidden_size + + model.context_norm.variance_epsilon + ).to(local_projection.dtype) + expected = local_projection * expected_scale * model.context_norm.weight[2:4] + + with torch.inference_mode(): + model.begin_auxiliary_stream(first) + model.accumulate_auxiliary_state(first, None) + model.accumulate_auxiliary_state(second, None) + output = model.finish_auxiliary_stream() + + torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.cpu_test +def test_k3_dspark_disables_streaming_for_incompatible_scratch_width() -> None: + model = object.__new__(K3DSparkModel) + nn.Module.__init__(model) + model.config = SimpleNamespace(target_hidden_size=8) + model._max_num_context_tokens = 16 + model._streamed_aux_scratch = torch.empty(16, 8) + + assert not model.bind_auxiliary_stream_scratch(torch.empty(16, 4)) + assert model._streamed_aux_scratch is None + + +@pytest.mark.cpu_test +def test_k3_dspark_preserves_fallback_when_target_has_no_stream_hook() -> None: + wrapper = object.__new__(K3DSparkForCausalLM) + nn.Module.__init__(wrapper) + wrapper.model = SimpleNamespace(bind_auxiliary_stream_scratch=Mock()) + target_model = SimpleNamespace( + get_language_model=lambda: SimpleNamespace(model=object()) + ) + + wrapper.bind_target_auxiliary_stream( + target_model, + torch.empty(16, 8), + ) + + wrapper.model.bind_auxiliary_stream_scratch.assert_not_called() + + +@pytest.mark.cpu_test +def test_k3_dspark_projects_streamed_context_one_layer_at_a_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = object.__new__(K3DSparkModel) + nn.Module.__init__(model) + model._context_local_start = 2 + model._context_local_width = 2 + model._context_kv_width = 3 + model._context_kv_lora_rank = 2 + model._context_rope_dim = 1 + model._context_rms_norm_eps = 1e-5 + model._context_kv_norm_weights = torch.ones(2, 2) + model.context_kv_proj = nn.Linear(4, 6, bias=False) + model.context_kv_proj.weight.data.copy_(torch.arange(24).view(6, 4)) + cache_updates = [] + + class Attention: + def __init__(self): + self.rotary_emb = SimpleNamespace(head_size=1, is_neox_style=True) + self.impl = SimpleNamespace( + do_kv_cache_update=lambda *args: cache_updates.append(args) + ) + self.kv_cache = object() + self.kv_cache_dtype = "bf16" + self._k_scale = 1.0 + + model.layers = [ + SimpleNamespace(self_attn=Attention()), + SimpleNamespace(self_attn=Attention()), + ] + model._get_rope_inputs = lambda positions: (positions, torch.empty(0)) + reduced = [] + + def record_all_reduce(tensor): + reduced.append(tensor.clone()) + return tensor + + monkeypatch.setattr( + dspark_mla, + "tensor_model_parallel_all_reduce_in_place", + record_all_reduce, + ) + rms_norm_io = [] + + def record_rms_norm(output, input_, weight, eps): + rms_norm_io.append((output, input_)) + output.copy_(input_) + + monkeypatch.setattr(dspark_mla.ops, "rms_norm", record_rms_norm) + monkeypatch.setattr(dspark_mla.ops, "rotary_embedding", lambda *args: None) + context = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + positions = torch.tensor([5, 6]) + first_slots = torch.tensor([0, 1]) + + model._precompute_streamed_context_kv( + context, + positions, + [first_slots, None], + ) + + expected = [] + for layer_index in range(2): + weight = model.context_kv_proj.weight[ + layer_index * 3 : (layer_index + 1) * 3, 2:4 + ] + expected.append(torch.nn.functional.linear(context, weight)) + torch.testing.assert_close(reduced[0], expected[0]) + torch.testing.assert_close(reduced[1], expected[1]) + assert all(output.data_ptr() != input_.data_ptr() for output, input_ in rms_norm_io) + assert len(cache_updates) == 1 + torch.testing.assert_close(cache_updates[0][0], expected[0][:, :2]) + torch.testing.assert_close(cache_updates[0][1], expected[0][:, 2:].view(2, 1, 1)) + assert cache_updates[0][3] is first_slots + + def test_context_kv_weights_are_loaded_as_merged_linear_shards(): weights = [ ( @@ -361,7 +743,7 @@ def test_v41_context_graph_replay_matches_checkpoint_projection( config.scheduler_config.max_num_batched_tokens = 8 config.scheduler_config.max_num_seqs = 8 config.compilation_config.max_cudagraph_capture_size = 8 - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device("cuda", torch.accelerator.current_device_index()) generator = torch.Generator(device=device).manual_seed(1831) class NativeLinear(nn.Module): diff --git a/tests/models/test_qwen3_dflash_streaming.py b/tests/models/test_qwen3_dflash_streaming.py new file mode 100644 index 000000000000..107d697b94e4 --- /dev/null +++ b/tests/models/test_qwen3_dflash_streaming.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.qwen3_dflash import ( + DFlashQwen3ForCausalLM, + DFlashQwen3Model, +) + + +class _Accumulator: + input_width = 8 + slice_width = 4 + max_tokens = 2048 + + def __init__(self, output: torch.Tensor): + self.output = output + self.tokens = 0 + self.inputs: list[torch.Tensor] = [] + + def begin(self, tokens: int) -> None: + self.tokens = tokens + self.inputs.clear() + + def append(self, source: torch.Tensor) -> None: + self.inputs.append(source.clone()) + + def finish(self) -> torch.Tensor: + self.output[: self.tokens].fill_(7) + return self.output[: self.tokens] + + +def test_padded_auxiliary_projection_trims_after_one_gather(monkeypatch): + from vllm.model_executor.models import qwen3_dflash + + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.config = SimpleNamespace(hidden_size=5) + model._aux_projection_tp_size = 2 + local = torch.arange(6, dtype=torch.bfloat16).view(2, 3) + calls = [] + + def gather(value, dim): + calls.append(value) + return torch.cat((value, value), dim=dim) + + monkeypatch.setattr(qwen3_dflash, "tensor_model_parallel_all_gather", gather) + result = model._gather_auxiliary_projection(local) + assert len(calls) == 1 + torch.testing.assert_close(result, torch.cat((local, local), dim=-1)[:, :5]) + + +def test_padded_auxiliary_projection_trims_small_inputs(): + fc = nn.Linear(8, 6, bias=False) + fc.input_size = 8 + wrapper = object.__new__(DFlashQwen3ForCausalLM) + nn.Module.__init__(wrapper) + wrapper.model = SimpleNamespace( + config=SimpleNamespace(hidden_size=5), + use_aux_hidden_state=True, + fc=fc, + _streamed_aux_accumulator=None, + ) + source = torch.arange(24, dtype=torch.float32).reshape(3, 8) + torch.testing.assert_close(wrapper.combine_hidden_states(source), fc(source)[:, :5]) + + +def test_dflash_streams_auxiliary_states_and_claims_result_once() -> None: + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.config = SimpleNamespace(hidden_size=4) + model.use_aux_hidden_state = True + model._target_hidden_size = 4 + model._streamed_aux_layer_ids = (2, 4) + model._streamed_aux_tokens = 0 + model._streamed_aux_index = 0 + model._streamed_aux_generation = 0 + model._completed_stream_generation = 0 + model._consumed_stream_generation = 0 + model._completed_stream_result = None + scratch = torch.empty(2048, 4) + accumulator = _Accumulator(scratch) + model.bind_auxiliary_stream(accumulator, scratch) + + first = torch.ones(1024, 4) + second = torch.full((1024, 4), 2.0) + residual = torch.full((1024, 4), 3.0) + assert model.can_stream_auxiliary_states((2, 4), first) + + model.begin_auxiliary_stream(first) + model.accumulate_auxiliary_state(first, None) + model.accumulate_auxiliary_state(second, residual) + result = model.finish_auxiliary_stream() + + torch.testing.assert_close(accumulator.inputs[0], first) + torch.testing.assert_close(accumulator.inputs[1], second + residual) + assert model.is_streamed_context_states([result]) + assert not model.is_streamed_context_states([result]) + + +def test_dflash_narrows_output_scratch_for_target_residual_staging() -> None: + """A narrower target state may reuse the draft output buffer safely.""" + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.config = SimpleNamespace(hidden_size=6) + model.use_aux_hidden_state = True + model._target_hidden_size = 4 + model._streamed_aux_layer_ids = (2, 4) + model._streamed_aux_tokens = 0 + model._streamed_aux_index = 0 + model._streamed_aux_generation = 0 + model._completed_stream_generation = 0 + model._consumed_stream_generation = 0 + model._completed_stream_result = None + scratch = torch.empty(2048, 6) + accumulator = _Accumulator(scratch) + model.bind_auxiliary_stream(accumulator, scratch) + + primary = torch.full((1024, 4), 2.0) + residual = torch.full((1024, 4), 3.0) + model.begin_auxiliary_stream(primary) + model.accumulate_auxiliary_state(primary, residual) + + assert len(accumulator.inputs) == 1 + torch.testing.assert_close( + accumulator.inputs[0], + torch.full((1024, 4), 5.0), + ) + + +def test_dflash_rejects_output_scratch_narrower_than_target_state() -> None: + """Streaming must fall back when target residuals cannot fit the scratch.""" + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.config = SimpleNamespace(hidden_size=4) + model._target_hidden_size = 6 + scratch = torch.empty(2048, 4) + + with pytest.raises(ValueError, match="at least 6 wide"): + model.bind_auxiliary_stream(_Accumulator(scratch), scratch) + + +@pytest.mark.parametrize("rows", [1, 8, 1023, 1024, 2048]) +def test_dflash_large_concatenated_inputs_reuse_staged_projection(rows) -> None: + """Every large fallback uses staging; its FC needs only small-row scratch.""" + fc = nn.Linear(8, 4, bias=False) + fc.input_size = 8 + accumulator = _Accumulator(torch.empty(2048, 4)) + wrapper = object.__new__(DFlashQwen3ForCausalLM) + nn.Module.__init__(wrapper) + wrapper.model = SimpleNamespace( + config=SimpleNamespace(hidden_size=4), + use_aux_hidden_state=True, + fc=fc, + _streamed_aux_accumulator=accumulator, + _gather_auxiliary_projection=lambda output: output, + ) + source = torch.arange(rows * 8, dtype=torch.float32).reshape(rows, 8) + output = wrapper.combine_hidden_states(source) + if rows < 1024: + assert not accumulator.inputs + torch.testing.assert_close(output, fc(source)) + else: + assert len(accumulator.inputs) == 2 + torch.testing.assert_close(torch.cat(accumulator.inputs, dim=-1), source) + assert output.data_ptr() == accumulator.output.data_ptr() + + +def test_dflash_preparation_bounds_nonstaged_fc_capacity() -> None: + """Prepared FC capacity covers exactly the rows not handled by staging.""" + from vllm.utils.b12x import B12xWorkload + + workload = B12xWorkload( + stage="weights", + token_counts=(1, 8, 1024, 4096), + fixed_token_counts=(1, 8, 1024), + output_dtype=torch.bfloat16, + max_tokens=4096, + max_seqs=128, + max_model_len=950000, + ) + accumulator = SimpleNamespace(preparation_unit=lambda name: name) + linear = SimpleNamespace(layer_name="draft.fc", unit=lambda w, name: w) + wrapper = SimpleNamespace( + model=SimpleNamespace( + _streamed_aux_accumulator=accumulator, + fc=SimpleNamespace(b12x_linear=linear), + ) + ) + small, staged = DFlashQwen3ForCausalLM.get_b12x_preparation_units( + wrapper, wrapper, workload + ) + assert small.max_tokens == 1023 + assert small.token_counts == (1, 8, 1023) + assert small.fixed_token_counts == (1, 8) + assert staged == "dflash.auxiliary" + + +def test_dflash_bf16_staging_keeps_local_projection_and_single_gather(monkeypatch): + from vllm.model_executor.models import qwen3_dflash + + monkeypatch.setattr(qwen3_dflash.envs, "VLLM_DFLASH_AUX_BF16_STAGING", True) + monkeypatch.setattr(qwen3_dflash.envs, "VLLM_DFLASH_AUX_MXFP8_STREAMING", False) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.config = SimpleNamespace(hidden_size=4) + model.use_aux_hidden_state = True + model._target_hidden_size = 4 + model._aux_projection_tp_size = 2 + model._streamed_aux_layer_ids = (2, 4) + model._streamed_aux_generation = 0 + weight = torch.arange(16, dtype=torch.bfloat16).reshape(2, 8) / 16 + calls = [] + + def local_linear(layer, x, bias): + calls.append("linear") + return torch.nn.functional.linear(x, weight, bias) + + def gather(local_output, dim): + calls.append("gather") + assert local_output.shape[1] == 2 + return torch.cat([local_output, local_output], dim=dim) + + monkeypatch.setattr(qwen3_dflash, "tensor_model_parallel_all_gather", gather) + model.fc = SimpleNamespace( + input_size=8, bias=None, quant_method=SimpleNamespace(apply=local_linear) + ) + wrapper = object.__new__(DFlashQwen3ForCausalLM) + nn.Module.__init__(wrapper) + wrapper.model = model + bound = [] + target = SimpleNamespace(set_aux_hidden_state_projector=bound.append) + scratch = torch.empty(2048, 4, dtype=torch.bfloat16) + wrapper.bind_target_auxiliary_stream(target, scratch) + assert bound == [model] + first = torch.full((1024, 4), 0.125, dtype=torch.bfloat16) + second = torch.full_like(first, -0.25) + model.begin_auxiliary_stream(first) + model.accumulate_auxiliary_state(first, None) + model.accumulate_auxiliary_state(second, None) + actual = model.finish_auxiliary_stream() + expected_local = torch.nn.functional.linear(torch.cat([first, second], -1), weight) + torch.testing.assert_close( + actual, torch.cat([expected_local, expected_local], -1), rtol=0, atol=0 + ) + assert calls == ["linear", "gather"] + assert ( + wrapper.get_b12x_preparation_units(wrapper, SimpleNamespace(stage="weights")) + == () + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "slices,width,output_width,capacity", + [(2, 128, 64, 257), (6, 7168, 7168, 4096), (6, 7168, 448, 4096)], +) +def test_prepared_mxfp8_staging_matches_concatenation_and_replays( + slices, + width, + output_width, + capacity, +): + """Slice reuse preserves MXFP8 bytes, one-GEMM math and graph ownership.""" + from b12x._lib.runtime_control import kernel_resolution_guard + from b12x.gemm import blockscaled + from b12x.preparation import PreparationSession + + from vllm.model_executor.kernels.linear.mxfp8.staged import ( + B12xMxfp8InputAccumulator, + ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + torch.manual_seed(314159) + device = torch.device("cuda:0") + weight = ( + torch.randn(output_width, slices * width, device=device, dtype=torch.bfloat16) + / 64 + ) + values, scales = mxfp8_e4m3_quantize(weight) + packed = blockscaled.pack_weight(values, scales) + del weight, values, scales + layer = SimpleNamespace( + b12x_mxfp8_packed_weight=packed, b12x_activation_mode="quantized" + ) + output = torch.empty(capacity, output_width, dtype=torch.bfloat16, device=device) + accumulator = B12xMxfp8InputAccumulator(layer, output, width) + unit = accumulator.preparation_unit("staged-test") + source = ( + torch.randn(slices, capacity, width, dtype=torch.bfloat16, device=device) / 8 + ) + + def run(rows): + accumulator.begin(rows) + for i in range(slices): + accumulator.append(source[i, :rows]) + return accumulator.finish() + + def check(rows, expected=None): + joined = torch.cat([source[i, :rows] for i in range(slices)], dim=-1) + ref_values, ref_scales = mxfp8_e4m3_quantize(joined) + torch.testing.assert_close( + accumulator.values[:rows].view(torch.uint8), + ref_values.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + accumulator.scales[:rows], ref_scales, rtol=0, atol=0 + ) + actual = output[:rows] + assert torch.isfinite(actual).all() and torch.count_nonzero(actual) > 0 + if expected is not None: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + dequantized_input = ref_values.float() * torch.exp2( + ref_scales.float() - 127 + ).repeat_interleave(32, dim=1) + dequantized_weight = packed.weight.values.float() * torch.exp2( + packed.weight.scale_rows.view(torch.uint8) + .reshape(output_width, -1) + .float() + - 127 + ).repeat_interleave(32, dim=1) + reference = dequantized_input @ dequantized_weight.t() + torch.testing.assert_close(actual.float(), reference, rtol=0.01, atol=0.003) + + with PreparationSession(device=device, autotune=True, compile_workers=2) as session: + session.prepare(unit.requests) + session.freeze() + pointers = [t.data_ptr() for t in (output, accumulator.values, accumulator.mma)] + for rows in (1, 8, 129, capacity): + with kernel_resolution_guard("staged input live counts"): + run(rows) + check(rows) + + rows = 129 + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run(rows) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run(rows) + source.mul_(0.5) + run(rows) + expected = output[:rows].clone() + torch.accelerator.synchronize() + output.fill_(float("nan")) + accumulator.values.view(torch.uint8).fill_(0xFF) + accumulator.mma.fill_(0xFF) + allocated = torch.accelerator.memory_allocated() + with kernel_resolution_guard("staged input graph replay"): + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() == allocated + check(rows, expected) + assert pointers == [ + t.data_ptr() for t in (output, accumulator.values, accumulator.mma) + ] diff --git a/tests/models/test_qwen3_dspark_tp_padding.py b/tests/models/test_qwen3_dspark_tp_padding.py new file mode 100644 index 000000000000..17ab0763d24e --- /dev/null +++ b/tests/models/test_qwen3_dspark_tp_padding.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Projection loading and GQA invariants for padded DSpark tensor parallelism.""" + +from types import MethodType, SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F + +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.models.qwen3_dflash import ( + _dflash_padded_width, + _enable_dflash_tp_padding, +) + + +def _weight(shape): + parameter = nn.Parameter( + torch.full(shape, float("nan"), dtype=torch.float64), requires_grad=False + ) + parameter.input_dim = 1 + parameter.output_dim = 0 + layer = nn.Module() + layer.register_parameter("weight", parameter) + layer.quant_method = UnquantizedLinearMethod() + _enable_dflash_tp_padding(layer) + return parameter + + +@pytest.mark.parametrize("tp,padded_kv", [(8, 16), (10, 20), (12, 24)]) +def test_padded_gqa_projection_matches_unsharded_attention(tp, padded_kv): + """Real checkpoint heads contribute once; absent heads contribute zero.""" + generator = torch.Generator().manual_seed(123) + hidden, heads, kv_heads, dim, rows = 4, 96, 16, 2, 3 + x = torch.randn(rows, hidden, generator=generator, dtype=torch.float64) + weights = [ + torch.randn(h * dim, hidden, generator=generator, dtype=x.dtype) + for h in (heads, kv_heads, kv_heads) + ] + output_weight = torch.randn(hidden, heads * dim, generator=generator, dtype=x.dtype) + + def attention(q, k, v, num_heads, num_kv): + q = q.view(rows, num_heads, dim).transpose(0, 1) + k = k.view(rows, num_kv, dim).transpose(0, 1) + v = v.view(rows, num_kv, dim).transpose(0, 1) + out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True) + return out.transpose(0, 1).reshape(rows, num_heads * dim) + + q, k, v = (F.linear(x, w) for w in weights) + reference = F.linear(attention(q, k, v, heads, kv_heads), output_weight) + local_kv = padded_kv // tp + local_q = local_kv * (heads // kv_heads) + widths = [local_q * dim, local_kv * dim, local_kv * dim] + partials = [] + for rank in range(tp): + qkv = _weight((sum(widths), hidden)) + projection = SimpleNamespace( + tp_rank=rank, + num_heads=local_q, + num_kv_heads=local_kv, + num_kv_head_replicas=1, + head_size=dim, + v_head_size=dim, + ) + projection.validate_shard_id = MethodType( + QKVParallelLinear.validate_shard_id, projection + ) + for key, weight in zip(("q", "k", "v"), weights): + QKVParallelLinear.weight_loader(projection, qkv, weight, key) + q, k, v = F.linear(x, qkv).split(widths, dim=-1) + out_weight = _weight((hidden, local_q * dim)) + RowParallelLinear.weight_loader( + SimpleNamespace(tp_rank=rank), out_weight, output_weight + ) + result = F.linear(attention(q, k, v, local_q, local_kv), out_weight) + assert torch.isfinite(result).all() + if rank * local_kv >= kv_heads: + assert torch.count_nonzero(qkv) == 0 + assert torch.count_nonzero(out_weight) == 0 + assert torch.count_nonzero(result) == 0 + partials.append(result) + torch.testing.assert_close(sum(partials), reference, rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("width", [7168, 14336]) +def test_draft_projection_padding_retains_supported_shapes(width): + assert _dflash_padded_width(width, 16) == width + padded = _dflash_padded_width(width, 10) + assert padded >= width and padded % (10 * 128) == 0 + + +def test_draft_padding_rejects_serialized_quantization(): + layer = nn.Module() + layer.quant_method = object() + with pytest.raises(ValueError, match="BF16/FP16 checkpoints"): + _enable_dflash_tp_padding(layer) diff --git a/tests/transformers_utils/test_dflash2_mla_config.py b/tests/transformers_utils/test_dflash2_mla_config.py new file mode 100644 index 000000000000..5ce90381a2da --- /dev/null +++ b/tests/transformers_utils/test_dflash2_mla_config.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + Qwen3ModelArchConfigConvertor, +) + + +def _write_draft(path, *, mla=True): + path.mkdir() + config = { + "architectures": ["DFlash2DraftModel"], + "model_type": "qwen3", + "hidden_size": 7168, + "intermediate_size": 14336, + "num_hidden_layers": 5, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "q_lora_rank": 1536, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "vocab_size": 163840, + "max_position_embeddings": 32768, + "torch_dtype": "bfloat16", + "is_causal": False, + "layer_types": ["sliding_attention"] * 4 + ["full_attention"], + "sliding_window": 4096, + "dflash_config": { + "attention_mode": "mla" if mla else "gqa", + "block_size": 8, + "conv_kernel_size": 2, + "conv_group_size": 16, + "selector_rank": 256, + "selector_top_k": 16, + "target_layer_ids": [19, 37, 66, 78, 90], + "mask_token_id": 163592, + }, + } + (path / "config.json").write_text(json.dumps(config)) + + +@pytest.mark.parametrize("tp", [8, 9, 10, 12, 16]) +def test_mla_dflash2_uses_padded_heads_and_latent_kv(tmp_path, tp): + path = tmp_path / "draft" + _write_draft(path) + model = ModelConfig(model=str(path), tokenizer_mode="skip", runner="draft") + assert model.architectures == ["DFlash2KimiK3Model"] + assert model.is_deepseek_mla and model.use_mla + parallel = ParallelConfig( + tensor_parallel_size=tp, distributed_executor_backend="external_launcher" + ) + model.verify_with_parallel_config(parallel) + model.verify_with_parallel_config(parallel) + assert model.get_head_size() == 576 + assert model.get_num_kv_heads(parallel) == 1 + assert model.get_num_attention_heads(parallel) == (64 + tp - 1) // tp + assert model.hf_text_config.original_num_attention_heads == 64 + assert model.hf_text_config.intermediate_size == 14336 + model.verify_with_parallel_config(ParallelConfig()) + assert model.hf_text_config.num_attention_heads == 64 + + +def test_mla_dflash2_eagle_wrapper_preserves_architecture(tmp_path): + path = tmp_path / "draft" + _write_draft(path) + target = ModelConfig( + model=str(path), tokenizer_mode="skip", runner="draft", max_model_len=32768 + ) + spec = SpeculativeConfig( + model=str(path), + method="dflash", + num_speculative_tokens=7, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + draft = spec.draft_model_config + assert draft.architectures == ["DFlash2KimiK3Model"] + assert draft.hf_config.model_type == "eagle" + assert draft.use_mla and draft.get_head_size() == 576 + + +def test_gqa_dflash2_architecture_is_unchanged(tmp_path): + path = tmp_path / "draft" + _write_draft(path, mla=False) + model = ModelConfig(model=str(path), tokenizer_mode="skip", runner="draft") + assert model.architectures == ["DFlash2DraftModel"] + assert not model.use_mla + + +def test_non_dflash_qwen_architecture_is_unchanged(): + from transformers import Qwen3Config + + config = Qwen3Config(architectures=["Qwen3ForCausalLM"]) + convertor = Qwen3ModelArchConfigConvertor(config, config) + assert convertor.get_architectures() == ["Qwen3ForCausalLM"] + assert not convertor.is_deepseek_mla() diff --git a/tests/transformers_utils/test_speculators_dspark_config.py b/tests/transformers_utils/test_speculators_dspark_config.py index ce32b6e5dfad..cade7b8888ce 100644 --- a/tests/transformers_utils/test_speculators_dspark_config.py +++ b/tests/transformers_utils/test_speculators_dspark_config.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + import pytest +from vllm.config import ModelConfig, ParallelConfig from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig pytestmark = pytest.mark.skip_global_cleanup @@ -56,3 +59,48 @@ def test_dspark_updater_maps_bonus_anchor_semantics() -> None: assert config["sample_from_anchor"] is False assert config["dspark_bonus_anchor"] is True + + +@pytest.mark.parametrize( + "tp,heads,kv_heads", + [ + (1, 96, 16), + (8, 96, 16), + (10, 120, 20), + (12, 144, 24), + (16, 96, 16), + (32, 96, 16), + ], +) +def test_dspark_parallel_config_preserves_complete_gqa_groups( + tmp_path, tp, heads, kv_heads +): + """Padding must not change which learned KV head serves a query head.""" + config = { + "architectures": ["Qwen3DSparkModel"], + "model_type": "qwen3", + "hidden_size": 7168, + "intermediate_size": 14336, + "num_hidden_layers": 5, + "num_attention_heads": 96, + "num_key_value_heads": 16, + "head_dim": 64, + "vocab_size": 163840, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", + } + (tmp_path / "config.json").write_text(json.dumps(config)) + model = ModelConfig(model=str(tmp_path), tokenizer_mode="skip", runner="draft") + parallel = ParallelConfig( + tensor_parallel_size=tp, distributed_executor_backend="external_launcher" + ) + for _ in range(2): + model.verify_with_parallel_config(parallel) + assert model.hf_text_config.num_attention_heads == heads + assert model.hf_text_config.num_key_value_heads == kv_heads + assert model.get_num_attention_heads(parallel) == heads // tp + assert model.get_num_kv_heads(parallel) == max(1, kv_heads // tp) + assert model.get_head_size() == 64 + model.verify_with_parallel_config(ParallelConfig()) + assert model.hf_text_config.num_attention_heads == 96 + assert model.hf_text_config.num_key_value_heads == 16 diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 36316cbeb558..df827532f7f2 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3727,6 +3727,191 @@ def test_glm_dspark_cache_groups_preserve_recurrent_target_and_draft_layers(): assert draft_group.kv_cache_spec == draft +@pytest.mark.parametrize( + "group_size,model_type,non_causal,method,expected_width", + [ + (0, "kimi_k3", True, "dspark", 24), + (3, "kimi_k3", True, "dspark", 3), + (6, "kimi_k3", True, "dspark", 6), + (6, "other", True, "dspark", 24), + (6, "kimi_k3", False, "dspark", 24), + (0, "kimi_k3", False, "dflash", 24), + (3, "kimi_k3", False, "dflash", 3), + (6, "kimi_k3", False, "dflash", 6), + ], +) +def test_kimi_parallel_draft_full_depth_cache_admission( + monkeypatch, group_size, model_type, non_causal, method, expected_width +): + """A replicated draft tail must not reserve 24-layer target-sized blocks.""" + monkeypatch.setenv("VLLM_K3_KV_GROUP_SIZE", str(group_size)) + mla = MLAAttentionSpec( + block_size=768, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="fp8", + ) + recurrent = MambaSpec( + block_size=12288, + shapes=((2304, 10), (6, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + page_size_padded=mla.page_size_bytes, + mamba_cache_mode="align", + num_speculative_blocks=7, + ) + draft = SlidingWindowMLASpec( + block_size=768, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="fp8", + sliding_window=32768, + dcp_replicated=True, + non_causal_multi_token_decode=non_causal, + ) + config = _grouping_config() + if method == "dflash": + draft = SlidingWindowSpec( + block_size=768, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=4096, + extra_retained_tokens=4096, + dcp_replicated=True, + ) + config.cache_config.kv_cache_layout = "BLHNC" + config.speculative_config = SimpleNamespace( + method=method, + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(model_type="k3_dspark") + ), + use_eagle=lambda: True, + use_eagle_block_drop=lambda: False, + ) + config.cache_config.mamba_cache_mode = "align" + config.model_config = SimpleNamespace( + max_model_len=950000, + hf_config=SimpleNamespace(model_type=model_type), + get_num_layers=lambda _: 93, + ) + config.attention_config = SimpleNamespace(hisparse_config=None) + config.parallel_config = SimpleNamespace( + decode_context_parallel_size=16, + prefill_context_parallel_size=1, + pipeline_parallel_size=1, + ) + config.max_in_flight_tokens = 8192 + config.use_request_boundary_checkpoints = False + full_layers = set(range(3, 92, 4)) | {92} + specs = { + f"model.layers.{i}.attn": mla if i in full_layers else recurrent + for i in range(93) + } + draft_layers = 6 if method == "dflash" else 5 + specs.update( + {f"draft.layers.{i}.attn": draft for i in range(93, 93 + draft_layers)} + ) + groups = get_kv_cache_groups(config, specs) + allocated = kv_cache_utils.get_kv_cache_config_from_groups( + config, groups, 1325000000 + ) + pool_stride = kv_cache_utils._pool_bytes_per_block(groups) + needed = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) + + bounded = expected_width < 24 + assert len(groups) == {3: 33, 6: 17, 24: 5}[expected_width] + assert pool_stride == mla.page_size_bytes * expected_width + assert (needed + pool_stride <= 1325000000) == bounded + assert set(name for group in groups for name in group.layer_names) == set(specs) + assert sum(len(group.layer_names) for group in groups) == len(specs) + scheduler = generate_scheduler_kv_cache_config([allocated]) + assert get_kv_cache_capacity(config, scheduler) == get_kv_cache_capacity( + config, allocated + ) + for group in groups: + original = specs[group.layer_names[0]] + assert group.kv_cache_spec == original + if bounded: + assert get_kv_cache_capacity(config, allocated)[0] >= 950000 + if expected_width == 3 and method == "dspark": + # This budget rounds into 60 allocator mapping blocks of 20 MiB; + # a six-layer group needs a 61st mapping block for the same context. + assert needed + pool_stride <= 1255000000 + compact = kv_cache_utils.get_kv_cache_config_from_groups( + config, groups, 1255000000 + ) + assert get_kv_cache_capacity(config, compact)[0] >= 950000 + + +@pytest.mark.parametrize("group_size", [0, 3]) +def test_kimi_mla_dflash2_full_context_cache_admission(monkeypatch, group_size): + """A full-context draft layer must not inherit a 24-layer pool stride.""" + monkeypatch.setenv("VLLM_K3_KV_GROUP_SIZE", str(group_size)) + config = _grouping_config() + config.model_config = SimpleNamespace( + max_model_len=950000, + hf_config=SimpleNamespace(model_type="kimi_k3"), + get_num_layers=lambda _: 93, + ) + config.parallel_config = SimpleNamespace( + decode_context_parallel_size=10, + prefill_context_parallel_size=1, + pipeline_parallel_size=1, + ) + config.speculative_config = SimpleNamespace( + method="dflash", use_eagle=lambda: True, use_eagle_block_drop=lambda: False + ) + config.attention_config = SimpleNamespace(hisparse_config=None) + config.cache_config.mamba_cache_mode = "align" + config.max_in_flight_tokens = 8192 + config.use_request_boundary_checkpoints = False + mla = MLAAttentionSpec( + block_size=1280, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="fp8", + ) + recurrent = MambaSpec( + block_size=12800, + shapes=((3840, 10), (10, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + page_size_padded=mla.page_size_bytes, + mamba_cache_mode="align", + num_speculative_blocks=7, + ) + full_draft = replace(mla, dcp_replicated=True, non_causal_multi_token_decode=True) + sliding_draft = SlidingWindowMLASpec( + block_size=1280, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="fp8", + sliding_window=4096, + dcp_replicated=True, + non_causal_multi_token_decode=True, + ) + full_layers = set(range(3, 92, 4)) | {92} + specs = { + f"model.layers.{i}.attn": mla if i in full_layers else recurrent + for i in range(93) + } + specs.update({f"draft.layers.{i}.attn": sliding_draft for i in range(93, 97)}) + specs["draft.layers.97.attn"] = full_draft + groups = get_kv_cache_groups(config, specs) + needed = kv_cache_utils._max_memory_usage_bytes_from_groups(config, groups) + stride = kv_cache_utils._pool_bytes_per_block(groups) + assert (needed + stride <= 4 * 1024**3) == (group_size == 3) + assert stride == mla.page_size_bytes * (group_size or 24) + # Grouping may change allocation width, never the attention window or owner. + for group in groups: + for name in group.layer_names: + assert group.kv_cache_spec == specs[name] + assert {name for group in groups for name in group.layer_names} == set(specs) + + def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): block_size = 544 full_spec = FullAttentionSpec( @@ -4981,6 +5166,13 @@ def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): assert "could be identified as the draft model's" not in caplog_vllm.text +def test_independent_dspark_cache_does_not_require_eagle_block_drop(caplog_vllm): + config = _spec_decode_grouping_config() + config.speculative_config.use_eagle_block_drop = lambda: False + get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) + assert "could be identified as the draft model's" not in caplog_vllm.text + + def test_no_warning_when_draft_group_is_identified(caplog_vllm): get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 9693c36bea98..d5946489a465 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, call @@ -57,6 +58,7 @@ KVCacheConfig, KVCacheGroupSpec, MambaSpec, + SlidingWindowMLASpec, SlidingWindowSpec, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry @@ -195,7 +197,57 @@ def test_swa_offload_window_covers_unaligned_hit( manager.complete_load(job.keys, state.req_context) -def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: +@pytest.mark.parametrize("dcp_size", [1, 8, 16]) +@pytest.mark.parametrize("replicated", [False, True]) +def test_offload_store_retains_draft_window_in_group_coordinates(dcp_size, replicated): + """Replicated draft pages must not inherit the target's DCP token width.""" + config = _make_vllm_config( + tensor_parallel_size=dcp_size, decode_context_parallel_size=dcp_size + ) + config.speculative_config = None + config.cache_config.block_size = 768 + config.cache_config.prefix_cache_retention_interval = 0 + groups = [ + KVCacheGroupSpec( + ["target"], + FullAttentionSpec( + block_size=768, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["draft"], + SlidingWindowMLASpec( + block_size=768, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=32768, + dcp_replicated=replicated, + ), + ), + ] + caches = KVCacheConfig(num_blocks=128, kv_cache_tensors=[], kv_cache_groups=groups) + spec = MockOffloadingSpec(build_offloading_config(config, caches)) + scheduler = OffloadingConnectorScheduler(spec, config, caches) + group = scheduler.config.kv_group_configs[1] + prompt_tokens = 78114 + end = prompt_tokens // group.tokens_per_block + mask = scheduler._reachable_store_block_mask( + group, 0, end, end, (prompt_tokens - 1,) + ) + alignment = 768 * dcp_size + boundary = (prompt_tokens - 1) // alignment * alignment + last_block = boundary // group.tokens_per_block + needed = (32768 - 2) // group.tokens_per_block + 1 + assert mask is not None + assert [index for index, keep in enumerate(mask) if keep] == list( + range(last_block - needed, last_block) + ) + + +def _make_partial_tail_scheduler( + *, recurrent_groups: int = 1, recurrent_first: bool = False +) -> OffloadingConnectorScheduler: vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True}) vllm_config.cache_config.prefix_match_unit = 4 vllm_config.speculative_config = None @@ -203,6 +255,14 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: enable_kv_cache_events=True, publisher="null" ) kv_cache_config = _make_mamba_hybrid_kv_cache_config() + full_group, recurrent_group = kv_cache_config.kv_cache_groups + recurrent = [ + replace(recurrent_group, layer_names=[f"mamba_layer_{index}"]) + for index in range(recurrent_groups) + ] + kv_cache_config.kv_cache_groups = ( + [*recurrent, full_group] if recurrent_first else [full_group, *recurrent] + ) spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) @@ -371,14 +431,22 @@ def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse): assert meta.jobs_to_flush == {job_id} -def test_normal_store_excludes_align_mode_mamba_sources(): - scheduler = _make_partial_tail_scheduler() +@pytest.mark.parametrize( + "recurrent_groups,recurrent_first", [(1, False), (1, True), (3, True)] +) +def test_normal_store_excludes_align_mode_mamba_sources( + recurrent_groups, recurrent_first +): + scheduler = _make_partial_tail_scheduler( + recurrent_groups=recurrent_groups, recurrent_first=recurrent_first + ) request = _make_partial_tail_request(scheduler) request.num_computed_tokens = 0 request.status = RequestStatus.RUNNING req_status = scheduler._req_status["req"] - req_status.group_states[0].block_ids[:] = [11] - req_status.group_states[1].block_ids[:] = [99] + full_group_idx = recurrent_groups if recurrent_first else 0 + for index, state in enumerate(req_status.group_states): + state.block_ids[:] = [11 if index == full_group_idx else 90 + index] req_status.update_offload_keys() scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) @@ -395,7 +463,10 @@ def test_normal_store_excludes_align_mode_mamba_sources(): src_spec = job.src_spec assert isinstance(src_spec, GPULoadStoreSpec) assert src_spec.block_ids.tolist() == [11] - assert src_spec.group_sizes == [1, 0] + expected_sizes = [0] * (recurrent_groups + 1) + expected_sizes[full_group_idx] = 1 + assert src_spec.group_sizes == expected_sizes + assert src_spec.block_indices == [0] * len(expected_sizes) def test_partial_lookup_returns_exact_boundary_and_group_load_keys(): diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index f1a1cb16ce37..29a3980c84a9 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -154,6 +154,25 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): _build_config(kv_connector=kv_connector) +@pytest.mark.parametrize("spec_name", [None, "CPUOffloadingSpec"]) +def test_cpu_offload_allows_expandable_segments(monkeypatch, spec_name): + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + _build_config( + kv_connector="OffloadingConnector", + kv_connector_extra_config={} if spec_name is None else {"spec_name": spec_name}, + ) + + +@pytest.mark.parametrize("spec_name", ["TieringOffloadingSpec", "ExternalSpec"]) +def test_nonlocal_offload_rejects_expandable_segments(monkeypatch, spec_name): + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + with pytest.raises(ValueError, match="expandable_segments"): + _build_config( + kv_connector="OffloadingConnector", + kv_connector_extra_config={"spec_name": spec_name}, + ) + + def test_lmcache_mp_engine_driven_allows_expandable_segments(monkeypatch): monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") _build_config( diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 10faa8b671a0..90039cdc7208 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -196,6 +196,117 @@ def __getitem__(self, _key): model._build_context_kv_buffers(layers_attn, has_bias=False) +def test_dflash_bf16_context_projection_preserves_weights(): + from torch import nn + + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + torch.manual_seed(7) + layers_attn = [ + SimpleNamespace( + q_size=4, + qkv_proj=SimpleNamespace( + quant_method=UnquantizedLinearMethod(), + weight=torch.randn(8, 32, dtype=torch.bfloat16), + ), + k_norm=SimpleNamespace(weight=torch.ones(2, dtype=torch.bfloat16)), + ) + for _ in range(2) + ] + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.hidden_norm = SimpleNamespace(weight=torch.ones(32, dtype=torch.bfloat16)) + model.layers = [SimpleNamespace(self_attn=attn) for attn in layers_attn] + model._build_context_kv_buffers(layers_attn, has_bias=False) + expected = torch.cat([attn.qkv_proj.weight[4:] for attn in layers_attn]) + torch.testing.assert_close(model._fused_kv_weight, expected, rtol=0, atol=0) + assert model._fused_kv_weight_scale is None + model.process_weights_after_loading() + torch.testing.assert_close(model._fused_kv_weight, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_dflash_serialized_mxfp8_context_has_independent_kernel(monkeypatch): + from torch import nn + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.quantization.modelopt import ( + build_linear_method, + ) + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + from vllm.utils.torch_utils import set_default_torch_dtype + + monkeypatch.setenv( + "VLLM_DISABLED_KERNELS", + "B12xMxfp8LinearKernel,FlashInferCutedslMxfp8LinearKernel,FlashInferCutlassMxfp8LinearKernel", + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) + torch.manual_seed(9) + layers_attn = [] + expected_weights = [] + with ( + set_current_vllm_config(VllmConfig()), + set_default_torch_dtype(torch.bfloat16), + torch.device("cuda"), + ): + for _ in range(2): + method = build_linear_method(None, "MXFP8", "") + linear = nn.Module() + linear.quant_method = method + method.create_weights(linear, 512, [512], 512, 512, torch.bfloat16) + linear.weight.data.copy_(torch.randn(512, 512).to(torch.float8_e4m3fn)) + linear.weight_scale.data.fill_(127) + expected_weights.append(linear.weight[256:].to(torch.bfloat16)) + layers_attn.append( + SimpleNamespace( + q_size=256, + qkv_proj=linear, + k_norm=SimpleNamespace(weight=torch.ones(128)), + ) + ) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.hidden_norm = SimpleNamespace(weight=torch.ones(512)) + model.layers = [SimpleNamespace(self_attn=attn) for attn in layers_attn] + model._fused_kv_linear = nn.Module() + model._build_context_kv_buffers(layers_attn, has_bias=False) + for attn in layers_attn: + attn.qkv_proj.quant_method.process_weights_after_loading(attn.qkv_proj) + source_kernel = layers_attn[0].qkv_proj.quant_method.kernel + source_weight = layers_attn[0].qkv_proj.weight.clone() + model.process_weights_after_loading() + assert model._fused_kv_quant_method.kernel is not source_kernel + torch.testing.assert_close( + layers_attn[0].qkv_proj.weight, source_weight, rtol=0, atol=0 + ) + expected_weight = torch.cat(expected_weights) + for rows in (1, 7, 256): + x = torch.randn(rows, 512) + + def run(x=x): + return model._fused_kv_quant_method.apply(model._fused_kv_linear, x) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + for _ in range(2): + x.normal_() + graph.replay() + expected = torch.nn.functional.linear(x, expected_weight) + assert torch.isfinite(actual).all() and actual.abs().max() > 0 + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.float().flatten(), dim=0 + ) + assert cosine > 0.9999 + + @pytest.mark.parametrize("mxfp8", [False, True]) @pytest.mark.parametrize("has_bias", [False, True]) def test_fused_context_projection_owns_its_linear_method(monkeypatch, mxfp8, has_bias): diff --git a/tests/v1/spec_decode/test_dflash_dcp.py b/tests/v1/spec_decode/test_dflash_dcp.py index 5c326b4cdf3c..51cbca33b382 100644 --- a/tests/v1/spec_decode/test_dflash_dcp.py +++ b/tests/v1/spec_decode/test_dflash_dcp.py @@ -2,16 +2,62 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +import numpy as np import pytest import torch import vllm.v1.worker.gpu.spec_decode.dflash.speculator as dflash_speculator +from vllm.config import ( + AttentionConfig, + CacheConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.layers.attention import Attention from vllm.model_executor.models.qwen3_dflash import DFlashAttention from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.attention.selector import _cached_get_attn_backend from vllm.v1.kv_cache_interface import SlidingWindowSpec +def test_dflash_selects_local_triton_without_relaxing_target_dcp(): + """A replicated draft selects local attention; a sharded target cannot.""" + config = VllmConfig( + attention_config=AttentionConfig(backend=AttentionBackendEnum.TRITON_ATTN), + cache_config=CacheConfig(block_size=16), + parallel_config=ParallelConfig( + tensor_parallel_size=16, + decode_context_parallel_size=16, + distributed_executor_backend="mp", + ), + ) + _cached_get_attn_backend.cache_clear() + with set_current_vllm_config(config): + draft = DFlashAttention( + 2, + 128, + 128**-0.5, + num_kv_heads=1, + cache_config=config.cache_config, + prefix="draft", + ) + assert draft.get_attn_backend().get_name() == "TRITON_ATTN" + assert draft.get_kv_cache_spec(config).dcp_replicated + with pytest.raises(ValueError, match="DCP not supported"): + Attention( + 2, + 128, + 128**-0.5, + num_kv_heads=1, + cache_config=config.cache_config, + prefix="target", + ) + assert config.parallel_config.decode_context_parallel_size == 16 + + def test_dflash_sliding_window_cache_uses_aligned_block_size(): attention = SimpleNamespace( sliding_window=2048, @@ -133,10 +179,11 @@ def capture_prepare_args(*args): ), dp_size=1, dp_rank=0, + pcp_manager=None, _group_causal=False, kv_cache_config=SimpleNamespace(), draft_tokens=torch.zeros((1, 1), dtype=torch.int64), - _build_draft_attn_metadata=lambda **_kwargs: {}, + _build_uniform_attn_metadata=lambda **_kwargs: {}, _prepare_eplb_forward=lambda *_args: None, _generate_draft=lambda *_args, **_kwargs: query_runs.append("eager"), ) @@ -181,7 +228,7 @@ def capture_metadata(_self, *_args, **kwargs): monkeypatch.setattr( dflash_speculator.DraftModelSpeculator, - "_build_draft_attn_metadata", + "_build_attn_metadata", capture_metadata, ) speculator = object.__new__(dflash_speculator.DFlashSpeculator) @@ -198,10 +245,10 @@ def capture_metadata(_self, *_args, **kwargs): ) object.__setattr__(speculator, "num_query_per_req", 4) - result = speculator._build_draft_attn_metadata( + result = speculator._build_attn_metadata( num_reqs=2, - num_reqs_padded=2, - num_tokens_padded=8, + batch_desc=SimpleNamespace(num_reqs=2, num_tokens=8), + query_start_loc_np=np.array([0, 4, 8], dtype=np.int32), seq_lens_cpu_upper_bound=seq_lens, step=0, ) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 673c0d711c7c..2da455d4bff0 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -15,6 +15,7 @@ import vllm.v1.hisparse.binding as attn_utils_module from tests.v1.attention.utils import dense_kv_cache_views +from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadataBuilder from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -221,6 +222,9 @@ def __init__(self, spec, layer_names, config, device): super().__init__(AttentionCGSupport.ALWAYS) self.config = config + def set_kernel_block_size(self, block_size): + self.kernel_block_size = block_size + class Backend: @staticmethod def full_cls_name(): @@ -235,7 +239,12 @@ def get_builder_cls(): return Builder configs = [ - SimpleNamespace(cache_config=SimpleNamespace(kv_cache_layout=None)) + SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_layout=None, kv_sharing_fast_prefill=False + ), + parallel_config=SimpleNamespace(use_ubatching=False), + ) for _ in range(2) ] layers = { @@ -270,6 +279,100 @@ def get_builder_cls(): assert group.get_metadata_builder(0).config is expected_config +@pytest.mark.parametrize("num_ubatches", [1, 2]) +@pytest.mark.parametrize("num_target_groups,num_draft_groups", [(4, 1), (8, 2)]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_mla_prefill_scratch_shares_groups_but_isolates_execution_lanes( + monkeypatch, num_ubatches, num_target_groups, num_draft_groups, device +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + + class Builder(MLACommonMetadataBuilder): + def __init__(self, spec, layer_names, config, device): + self.chunked_prefill_workspace = torch.empty( + (128, spec.head_size), dtype=spec.dtype, device=device + ) + + def get_cudagraph_support(self, *_args): + return AttentionCGSupport.ALWAYS + + class Backend: + @staticmethod + def full_cls_name(): + return (__name__, "MLABackend") + + @staticmethod + def get_supported_kernel_block_sizes(): + return [16] + + @staticmethod + def get_builder_cls(): + return Builder + + configs = [ + SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_layout=None, kv_sharing_fast_prefill=False + ), + parallel_config=SimpleNamespace( + use_ubatching=num_ubatches > 1, num_ubatches=num_ubatches + ), + ) + for _ in range(2) + ] + names = [f"target.{i}" for i in range(num_target_groups)] + [ + f"draft.{i}" for i in range(num_draft_groups) + ] + layers = { + name: SimpleNamespace(get_attn_backend=lambda: Backend, num_heads=6) + for name in names + } + monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda _: {}) + monkeypatch.setattr( + attn_utils, + "get_layers_from_vllm_config", + lambda config, layer_type, names: {name: layers[name] for name in names}, + ) + spec = MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=576, dtype=torch.bfloat16 + ) + cache = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec([name], spec) for name in names], + ) + groups, _, _ = init_attn_backend( + cache, + configs[0], + torch.device(device), + layer_vllm_configs={ + name: configs[1] for name in names if name.startswith("draft") + }, + ) + target = [group[0] for group in groups[:num_target_groups]] + draft = [group[0] for group in groups[num_target_groups:]] + pointers = set() + for lane, lane_groups in enumerate((target, draft)): + for ubatch in range(num_ubatches): + scratch = ( + lane_groups[0].get_metadata_builder(ubatch).chunked_prefill_workspace + ) + pointers.add(scratch.data_ptr()) + scratch.fill_(1 + lane * num_ubatches + ubatch) + assert all( + group.get_metadata_builder(ubatch).chunked_prefill_workspace is scratch + for group in lane_groups + ) + assert len(pointers) == 2 * num_ubatches + for lane, lane_groups in enumerate((target, draft)): + for ubatch in range(num_ubatches): + scratch = ( + lane_groups[0].get_metadata_builder(ubatch).chunked_prefill_workspace + ) + assert torch.all(scratch == 1 + lane * num_ubatches + ubatch) + + def test_attention_checks_preserve_global_and_target_scoped_support(): spec = FullAttentionSpec( block_size=16, diff --git a/tests/v1/worker/test_dspark_greedy_markov.py b/tests/v1/worker/test_dspark_greedy_markov.py index 348838251234..2dbc9d8921e2 100644 --- a/tests/v1/worker/test_dspark_greedy_markov.py +++ b/tests/v1/worker/test_dspark_greedy_markov.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Exact full-vocabulary DSpark greedy tokens, including graph replay.""" +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch @@ -11,6 +14,56 @@ ) +@pytest.mark.parametrize("needs_distribution", [False, True]) +def test_sequential_sharded_sampling_preserves_each_token(needs_distribution): + from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + shared = torch.zeros(2, dtype=torch.int64) + seen_anchors = [] + + def embed(ids): + seen_anchors.append(ids.clone()) + return ids[:, None].float() + + def sample(base, bias): + return shared.copy_((base + bias).argmax(-1)) + + # Distinct position maxima expose overwrites from a reused argmax output. + base = torch.tensor([[3.0, 0.0], [0.0, 3.0], [3.0, 0.0], [0.0, 3.0]]) + model = SimpleNamespace( + supports_local_draft_argmax=lambda: True, + compute_local_draft_logits=Mock(return_value=base), + markov_embed=embed, + compute_local_markov_bias=lambda x: x.expand(-1, 2), + sample_local_draft_logits=Mock(side_effect=sample), + gather_local_draft_logits=Mock(side_effect=lambda b, m: b + m), + ) + spec = SimpleNamespace( + _draft_topk=None, + num_speculative_steps=2, + model=model, + sample_indices=torch.arange(4), + sample_idx_mapping=torch.tensor([0, 0, 1, 1]), + sample_pos=torch.arange(4), + use_confidence_head=False, + input_buffers=SimpleNamespace(input_ids=torch.tensor([5, 9])), + _anchor_idx=torch.tensor([0, 1]), + draft_tokens=torch.full((2, 2), -1, dtype=torch.int64), + draft_logits=object() if needs_distribution else None, + acceptance_estimator=None, + _sample_logits=Mock(side_effect=lambda logits, *args: logits.argmax(-1)), + ) + DSparkSpeculator._sample_sequential(spec, 2, torch.zeros(4, 8)) + assert spec.draft_tokens.tolist() == [[0, 1], [0, 1]] + assert [ids.tolist() for ids in seen_anchors] == [[5, 9], [0, 0]] + assert model.sample_local_draft_logits.call_count == ( + 0 if needs_distribution else 2 + ) + assert model.gather_local_draft_logits.call_count == ( + 2 if needs_distribution else 0 + ) + + @pytest.mark.parametrize("base_vocab,bias_vocab", [(2049, 2048), (2048, 2049)]) def test_greedy_markov_rejects_mismatched_vocabulary_before_dispatch( base_vocab, bias_vocab diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py index 94655f192c89..2d3a073cd2e2 100644 --- a/tests/v1/worker/test_dspark_utils.py +++ b/tests/v1/worker/test_dspark_utils.py @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass, field +from types import SimpleNamespace + +import pytest +import torch.nn as nn from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config @@ -49,3 +53,67 @@ def test_dspark_parallel_config_disables_eplb_atomically(): assert draft_config.eplb_config.num_redundant_experts == 0 assert not draft_config.enable_elastic_ep assert draft_config.eplb_config is not target_config.eplb_config + + +@pytest.mark.parametrize("draft_format", [None, "auto", "safetensors"]) +def test_dspark_selects_explicit_draft_loader(monkeypatch, draft_format): + """A CPU-backed draft loader must not inherit target GPU staging buffers.""" + from vllm.config import AttentionConfig, CacheConfig, LoadConfig + from vllm.model_executor import model_loader + from vllm.model_executor.models import utils as model_utils + from vllm.v1.worker.gpu.spec_decode import utils as spec_utils + from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model + from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils + + @dataclass + class Config: + speculative_config: object + parallel_config: object + kernel_config: object + attention_config: object + cache_config: object + load_config: object + model_config: object + + target_load = LoadConfig(load_format="instanttensor") + draft_load = LoadConfig(load_format=draft_format) if draft_format else None + model_config = SimpleNamespace( + hf_config=SimpleNamespace(model_type="qwen3", num_hidden_layers=1), + get_vocab_size=lambda: 32, + ) + config = Config( + speculative_config=SimpleNamespace( + draft_model_config=model_config, + draft_parallel_config=SimpleNamespace(tensor_parallel_size=4), + attention_backend=None, + moe_backend=None, + kv_cache_dtype=None, + draft_load_config=draft_load, + ), + parallel_config=_FakeParallelConfig(), + kernel_config=None, + attention_config=AttentionConfig(), + cache_config=CacheConfig(), + load_config=target_load, + model_config=model_config, + ) + target, draft = nn.Module(), nn.Module() + target.model, draft.model = nn.Module(), nn.Module() + selected_loaders = [] + + def select_loader(load_config): + selected_loaders.append(load_config) + return SimpleNamespace(load_model=lambda **kwargs: draft) + + monkeypatch.setattr(model_loader, "get_model_loader", select_loader) + monkeypatch.setattr(model_utils, "get_draft_quant_config", lambda config: None) + monkeypatch.setattr( + spec_utils, "get_pp_group", lambda: SimpleNamespace(world_size=1) + ) + monkeypatch.setattr( + eagle_utils, "get_pp_group", lambda: SimpleNamespace(world_size=1) + ) + assert load_dspark_model(target, config) is draft + assert selected_loaders == [draft_load or target_load] + assert config.load_config is target_load + assert config.load_config.load_format == "instanttensor" diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index dff19a289439..5b14fd0c21da 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -254,12 +254,26 @@ def test_reinterpret_u64_as_i64_preserves_pointer_bits(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_aligned_state_indices_graph_replay_masks_padding_and_refreshes_blocks(): +@pytest.mark.parametrize("num_groups", [2, 69]) +@pytest.mark.parametrize("num_state_slots", [1, 8]) +@pytest.mark.parametrize("num_reqs", [1, 4, 33]) +def test_aligned_state_indices_graph_replay_masks_padding_and_refreshes_blocks( + num_groups, num_state_slots, num_reqs +): + max_reqs, table_width, sentinel = num_reqs + 2, 16, -77 tables = [ - torch.arange(16, dtype=torch.int32, device="cuda").view(4, 4) + group * 100 - for group in range(2) + torch.arange(max_reqs * table_width, dtype=torch.int32, device="cuda") + .view(max_reqs, table_width) + .add(group * 10000) + for group in range(num_groups) ] - indices = torch.empty((2, 4, 1), dtype=torch.int32, device="cuda") + storage = torch.full( + (num_groups, max_reqs, num_state_slots * 2), + sentinel, + dtype=torch.int32, + device="cuda", + ) + indices = storage[:, :, ::2] ctx = SimpleNamespace( is_initialized=True, aligned_state_indices=indices, @@ -268,25 +282,55 @@ def test_aligned_state_indices_graph_replay_masks_padding_and_refreshes_blocks() dtype=torch.int64, device="cuda", ), - block_table_stride_req=4, + block_table_stride_req=table_width, block_size=16, - num_groups=2, + num_groups=num_groups, ) - seq_lens = torch.tensor([16, 17, 0, 0], dtype=torch.int32, device="cuda") + seq_lens = torch.zeros(max_reqs * 2, dtype=torch.int32, device="cuda")[::2] + + def set_lengths(values): + lengths = torch.tensor( + [values[row % len(values)] for row in range(max_reqs)], dtype=torch.int32 + ) + seq_lens.copy_(lengths) + return lengths + + lengths = set_lengths([16, 17, 0]) def compute(): - MambaSpecDecodeGPUContext.compute_aligned_state_indices(ctx, seq_lens, 4) + MambaSpecDecodeGPUContext.compute_aligned_state_indices(ctx, seq_lens, num_reqs) + + def check(lengths, table_delta=0): + expected = torch.full((num_groups, max_reqs, num_state_slots), sentinel) + for group in range(num_groups): + for row in range(num_reqs): + first_slot = max((int(lengths[row]) - 1) // 16, 0) + expected[group, row] = ( + group * 10000 + + row * table_width + + first_slot + + torch.arange(num_state_slots) + + table_delta + if lengths[row] > 0 + else -1 + ) + torch.testing.assert_close( + indices.cpu(), expected.to(torch.int32), rtol=0, atol=0 + ) + assert torch.all(storage[:, :, 1::2] == sentinel) compute() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): compute() graph.replay() - assert indices[:, :, 0].tolist() == [[0, 5, -1, -1], [100, 105, -1, -1]] + check(lengths) - seq_lens.copy_(torch.tensor([17, 0, 33, 0], dtype=torch.int32, device="cuda")) + lengths = set_lengths([33, 0, 17]) + for table in tables: + table.add_(3) graph.replay() - assert indices[:, :, 0].tolist() == [[1, -1, 10, -1], [101, -1, 110, -1]] + check(lengths, table_delta=3) def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index a310bee570a5..3c89e303699f 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -157,10 +157,75 @@ def test_workspace_reservation_covers_every_execution_slot(monkeypatch) -> None: assert view.data_ptr() == reserved.data_ptr() manager.lock() - with pytest.raises(AssertionError, match="reserve_all"): + with pytest.raises(AssertionError, match="Workspace is locked"): manager.reserve_all(((1024,), torch.uint8)) +def test_workspace_reservation_preserves_independent_lane_capacities(monkeypatch): + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + manager.get_simultaneous(((1024,), torch.uint8)) + active_ubatch[0] = 1 + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((256,), torch.uint8)) + manager.reserve_by_lane() + assert [buffer.numel() for buffer in manager._current_workspaces] == [ + 1024, + 256, + 1024, + 256, + ] + pointers = [buffer.data_ptr() for buffer in manager._current_workspaces] + assert len(set(pointers)) == 4 + manager.lock() + manager.reserve_by_lane() + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((256,), torch.uint8)) + with pytest.raises(AssertionError, match="Workspace is locked"): + manager.get_simultaneous(((512,), torch.uint8)) + assert pointers == [buffer.data_ptr() for buffer in manager._current_workspaces] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_lane_reservation_replays_graphs_with_independent_changed_inputs(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager(torch.device("cuda"), num_lanes=2) + manager.get_simultaneous(((1024,), torch.float32)) + with workspace.use_workspace_lane(1): + manager.get_simultaneous(((64,), torch.float32)) + manager.reserve_by_lane() + manager.lock() + assert [buffer.numel() for buffer in manager._current_workspaces] == [4096, 256] + source = torch.ones(64, device="cuda") + output = torch.empty_like(source) + + def run(): + (target,) = manager.get_simultaneous(((1024,), torch.float32)) + target[:64].copy_(source) + with workspace.use_workspace_lane(1): + (draft,) = manager.get_simultaneous(((64,), torch.float32)) + torch.mul(source, 3, out=draft) + torch.add(target[:64], draft, out=output) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + pointers = [buffer.data_ptr() for buffer in manager._current_workspaces] + for _ in range(3): + source.normal_() + graph.replay() + torch.testing.assert_close(output, source + source * 3, rtol=0, atol=0) + assert pointers == [buffer.data_ptr() for buffer in manager._current_workspaces] + + def test_workspace_lane_validation(monkeypatch) -> None: monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) diff --git a/vllm/config/model.py b/vllm/config/model.py index 707404a79805..4a0fbc01d9aa 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1278,6 +1278,7 @@ def _verify_quantization(self) -> None: "awq_marlin", "inc", "moe_wna16", + "exl3", "modelopt", "modelopt_fp4", "modelopt_mxfp8", @@ -1957,8 +1958,12 @@ def use_mla(self) -> bool: # Bidirectional DeepSeek variants (is_causal=False, used by some # embedding models) must use the non-MLA attention path, since the - # MLA kernels only support causal attention. - if not getattr(self.hf_text_config, "is_causal", True): + # ordinary MLA path only supports causal attention. Kimi-K3 DFlash2 + # supplies a separate non-causal block-query implementation. + if ( + not getattr(self.hf_text_config, "is_causal", True) + and "DFlash2KimiK3Model" not in self.architectures + ): return False return self.is_deepseek_mla diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 1b827e7cf5d7..127537fc96fb 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -15,6 +15,7 @@ from vllm.config.kernel import MoEBackend from vllm.config.model import HfOverrides, ModelConfig from vllm.config.parallel import ParallelConfig +from vllm.config.quantization import QuantizationConfigArgs, resolve_quantization_config from vllm.config.utils import config from vllm.logger import init_logger from vllm.transformers_utils.config import get_hf_text_config @@ -408,6 +409,8 @@ class SpeculativeConfig: """Quantization method that was used to quantize the draft model weights. If `None`, we assume the model weights are not quantized. Note that it only takes effect when using the draft model-based speculative method.""" + quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None + """Per-layer online weight quantization settings for the draft model only.""" moe_backend: MoEBackend | None = None """MoE backend to use for the draft model. When `None`, the draft model inherits the target model's `--moe-backend` setting. Useful when the @@ -1419,6 +1422,9 @@ def __post_init__(self): max_model_len=self.max_model_len, # type: ignore[arg-type] spec_target_max_model_len=self.target_model_config.max_model_len, quantization=self.quantization, + quantization_config=resolve_quantization_config( + self.quantization, self.quantization_config + ), enforce_eager=self.target_model_config.enforce_eager, max_logprobs=self.target_model_config.max_logprobs, hf_overrides=draft_hf_overrides, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a7695fac3e6a..459d81f68dbc 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1240,9 +1240,8 @@ def _verify_kv_transfer_compat(self) -> None: # registrations pointing at stale physical pages after any remap, # producing RDMA failures like IBV_WC_REM_ACCESS_ERR / # NIXL_ERR_REMOTE_DISCONNECT at the first inter-node KV transfer. - # We can't enumerate every in-tree and out-of-tree connector that - # pins memory, so we conservatively reject the combination whenever - # any KV connector is configured. + # Unknown transports remain rejected unless the connector explicitly + # declares a configuration that does not retain GPU page registrations. # # CuMem allocator is exempt: CuMemAllocator.use_memory_pool toggles # expandable_segments off around its pool (see #40812), so the KV @@ -1254,6 +1253,8 @@ def _verify_kv_transfer_compat(self) -> None: return if self.model_config is not None and (self.model_config.enable_cumem_allocator): return + if self._connector_supports_vmm_safe_transfers(): + return # Engine-driven LMCache MP transport does not pin or register KV # addresses; GPU gather/scatter remains in the vLLM worker. @@ -1277,9 +1278,26 @@ def _verify_kv_transfer_compat(self) -> None: "unset expandable_segments:True or enable the cumem allocator " "(sleep mode does this automatically and also " "routes KV allocations through CuMemAllocator's pool, where " - "expandable_segments is automatically disabled)." + "expandable_segments is automatically disabled), or use a connector " + "that declares SupportsVmmSafeTransfers." + ) + + def _connector_supports_vmm_safe_transfers(self) -> bool: + from vllm.distributed.kv_transfer.kv_connector.factory import ( + KVConnectorFactory, + ) + from vllm.distributed.kv_transfer.kv_connector.v1 import ( + supports_vmm_safe_transfers, ) + try: + connector_cls = KVConnectorFactory.get_connector_class( + self.kv_transfer_config + ) + except (AttributeError, ImportError, TypeError, ValueError): + return False + return supports_vmm_safe_transfers(connector_cls, self.kv_transfer_config) + def _verify_sampling_replay_config(self) -> None: model_config = self.model_config if model_config is None or not model_config.return_sampling_mask: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py index 47329207f4b7..96ab863d0260 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py @@ -4,7 +4,9 @@ KVConnectorBase_V1, KVConnectorRole, SupportsHMA, + SupportsVmmSafeTransfers, supports_hma, + supports_vmm_safe_transfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import ( # noqa: E501 DecodeBenchConnector, @@ -15,5 +17,7 @@ "KVConnectorBase_V1", "supports_hma", "SupportsHMA", + "SupportsVmmSafeTransfers", + "supports_vmm_safe_transfers", "DecodeBenchConnector", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index a7b6c2aa14f8..4cadacdf4a1f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -54,7 +54,7 @@ from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: - from vllm.config import VllmConfig + from vllm.config import KVTransferConfig, VllmConfig from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( KVConnectorPromMetrics, @@ -137,6 +137,32 @@ def supports_hma(connector: Any) -> bool: return isinstance(connector, SupportsHMA) +class SupportsVmmSafeTransfers(ABC): + """Declare configurations that do not retain GPU physical-page registrations. + + A qualifying transport accesses GPU buffers through virtual addresses at + transfer time, without RDMA, GPUDirect or GDS page registrations. + """ + + @classmethod + @abstractmethod + def supports_vmm_safe_transfer_config( + cls, kv_transfer_config: "KVTransferConfig" + ) -> bool: + """Return whether this exact connector configuration is VMM-safe.""" + raise NotImplementedError + + +def supports_vmm_safe_transfers( + connector: Any, kv_transfer_config: "KVTransferConfig" +) -> bool: + """Evaluate a connector's explicit CUDA VMM transfer contract.""" + connector_cls = connector if isinstance(connector, type) else type(connector) + return issubclass( + connector_cls, SupportsVmmSafeTransfers + ) and connector_cls.supports_vmm_safe_transfer_config(kv_transfer_config) + + class KVConnectorRole(enum.Enum): # Connector running in the scheduler process SCHEDULER = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index ff31c12da537..f04633d8f5bf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1494,7 +1494,10 @@ def _reachable_store_block_mask( use_eagle=group_config.is_eagle_group, retention_interval=self.config.retention_interval, reachable_boundaries=reachable_boundaries, - dcp_world_size=self.config.dcp_world_size, + # The offload plan already resolves each group's logical span. + # Replicated draft pages retain their physical token width even + # when target attention pages are sharded across DCP ranks. + dcp_world_size=(group_config.tokens_per_block // kv_cache_spec.block_size), final_segment_end_block=( final_segment_end_chunk_idx * blocks_per_chunk if final_segment_end_chunk_idx is not None @@ -1567,9 +1570,6 @@ def _build_store_jobs( group_config, group_state, num_offloadable_tokens ) - if group_config.requires_cow_source: - continue - start_chunk_idx = group_state.next_stored_chunk_idx prompt_horizon_chunks = ( prompt_offloadable_tokens // group_config.tokens_per_chunk @@ -1606,6 +1606,8 @@ def _build_store_jobs( partial_segment_start_block // blocks_per_chunk ) start_chunk_idx = min(start_chunk_idx, partial_segment_start_chunk) + # Transfer descriptors retain one slot per cache group, even + # when recurrent state is stored only by boundary handoffs. group_store_ranges.append((start_chunk_idx, num_chunks)) if group_config.requires_cow_source: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 73570881753d..9a5fe61ee7a5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -6,12 +6,13 @@ import torch -from vllm.config import VllmConfig +from vllm.config import KVTransferConfig, VllmConfig from vllm.distributed.kv_events import KVCacheEvent from vllm.distributed.kv_transfer.kv_connector.v1 import ( KVConnectorBase_V1, KVConnectorRole, SupportsHMA, + SupportsVmmSafeTransfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -49,7 +50,16 @@ from vllm.v1.request import Request -class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): +class OffloadingConnector(KVConnectorBase_V1, SupportsHMA, SupportsVmmSafeTransfers): + @classmethod + def supports_vmm_safe_transfer_config( + cls, kv_transfer_config: KVTransferConfig + ) -> bool: + # CPU copies pin host storage, not GPU cache pages. Tiered or external + # storage implementations may retain GPU page registrations. + extra = kv_transfer_config.kv_connector_extra_config or {} + return extra.get("spec_name", "CPUOffloadingSpec") == "CPUOffloadingSpec" + @cached_property def _bounding_group_ids(self) -> tuple[int, ...]: """Prefix-cacheable groups this connector does not offload. diff --git a/vllm/model_executor/kernels/linear/bf16_staging.py b/vllm/model_executor/kernels/linear/bf16_staging.py new file mode 100644 index 000000000000..11093113ad46 --- /dev/null +++ b/vllm/model_executor/kernels/linear/bf16_staging.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Retained BF16 input storage for sequential auxiliary-state producers.""" + +from __future__ import annotations + +import torch + + +class Bf16InputAccumulator: + """Copy feature slices without quantization, then invoke one local linear. + + The caller supplies a BF16-input quantization method. Output gathering is + owned by the caller: invoking ``quant_method.apply`` avoids running a + column-parallel layer's collective twice. No source slice remains live + after its copy completes on the caller's stream. + """ + + def __init__(self, layer, reference: torch.Tensor, slice_width: int): + self.input_width = int(layer.input_size) + if reference.ndim != 2 or reference.dtype != torch.bfloat16: + raise ValueError("BF16 staging requires a BF16 [capacity,N] reference") + if slice_width <= 0 or self.input_width % slice_width: + raise ValueError("BF16 staging requires complete feature slices") + self.max_tokens = int(reference.shape[0]) + if self.max_tokens <= 0: + raise ValueError("BF16 staging capacity must be positive") + self.slice_width = slice_width + self.layer = layer + self.values = torch.empty( + (self.max_tokens, self.input_width), + dtype=reference.dtype, + device=reference.device, + ) + self._tokens = self._column = 0 + + def begin(self, tokens: int) -> None: + if not 0 < tokens <= self.max_tokens: + raise ValueError("BF16 staging row count is outside reserved capacity") + self._tokens, self._column = tokens, 0 + + def append(self, source: torch.Tensor) -> None: + rows, column, width = self._tokens, self._column, self.slice_width + if ( + rows == 0 + or source.shape != (rows, width) + or source.dtype != self.values.dtype + or source.device != self.values.device + or column + width > self.input_width + ): + raise ValueError("BF16 staging slice has incompatible geometry or dtype") + self.values[:rows, column : column + width].copy_(source) + self._column += width + + def finish(self) -> torch.Tensor: + if self._tokens == 0 or self._column != self.input_width: + raise RuntimeError("BF16 staging input is incomplete") + rows = self._tokens + self._tokens = 0 + return self.layer.quant_method.apply( + self.layer, self.values[:rows], getattr(self.layer, "bias", None) + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/staged.py b/vllm/model_executor/kernels/linear/mxfp8/staged.py new file mode 100644 index 000000000000..bde0034bd957 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/staged.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepared MXFP8 staging for sequential auxiliary-state producers.""" + +from __future__ import annotations + +import torch + +from vllm.utils.b12x import B12xPreparationUnit + + +def _scale_view(storage: torch.Tensor, rows: int, columns: int) -> torch.Tensor: + """Map caller-owned UE8M0 bytes to the documented F8_128x4 layout.""" + mt, kt = (rows + 127) // 128, columns // 128 + return storage.as_strided( + (32, 4, mt, 4, kt, 1), (16, 4, kt * 512, 1, 512, mt * kt * 512) + ) + + +class B12xMxfp8InputAccumulator: + """Assemble quantized feature slices, then perform one complete GEMM. + + Each slice ends at a K128 boundary. Quantization therefore preserves the + same K32 groups as concatenation followed by quantization. All buffers are + allocated before serving; each invocation passes fresh views to prepared + kernels. No partial GEMM summation or target-model quantization is added. + """ + + def __init__(self, layer, output: torch.Tensor, slice_width: int): + from b12x.gemm import blockscaled + from b12x.quantization import mxfp8 + + packed = getattr(layer, "b12x_mxfp8_packed_weight", None) + if ( + packed is None + or getattr(layer, "b12x_activation_mode", None) != "quantized" + ): + raise ValueError("staging requires a quantized-activation B12X MXFP8 FC") + if output.ndim != 2 or output.dtype != torch.bfloat16: + raise ValueError("staging output must be a BF16 [capacity,N] tensor") + if ( + output.shape[1] != packed.out_features + or output.device != packed.weight.values.device + ): + raise ValueError( + "staging output differs from the FC weight geometry/device" + ) + if ( + packed.in_features != packed.padded_in_features + or slice_width <= 0 + or slice_width % 128 + or packed.in_features % slice_width + ): + raise ValueError("staging requires complete K128 slices without K padding") + self.input_width = int(packed.in_features) + self.max_tokens = int(output.shape[0]) + self.slice_width = slice_width + self.output = output + self.packed = packed + self._tokens = self._column = 0 + + def buffers(width): + return ( + torch.empty( + (self.max_tokens, width), + dtype=torch.float8_e4m3fn, + device=output.device, + ), + torch.empty( + (self.max_tokens, width // 32), + dtype=torch.uint8, + device=output.device, + ), + torch.empty( + ((self.max_tokens + 127) // 128 * (width // 128) * 512,), + dtype=torch.uint8, + device=output.device, + ), + ) + + self.values, self.scales, self.mma = buffers(self.input_width) + self.slice_values, self.slice_scales, self.slice_mma = buffers(slice_width) + self.workspace = torch.empty(0, dtype=torch.float32, device=output.device) + self.quant_plan = mxfp8.plan( + mxfp8.Mxfp8Query( + rows=self.max_tokens, + columns=slice_width, + dtype="bfloat16", + value_order="linear", + scale_rows_layout="row_major", + scale_rows_storage="uint8", + scale_mma_layout="linear_storage", + scale_mma_storage="uint8", + ) + ) + lhs, rhs, out = self._operands(self.max_tokens) + self.gemm_plan = blockscaled.plan( + blockscaled.query_from_call( + lhs, + rhs, + out=out, + workspace=self.workspace, + ab_dtype="float8_e4m3fn", + sf_dtype="float8_e8m0fnu", + sf_vec_size=32, + expected_m=self.max_tokens, + ) + ) + + def _operands(self, rows): + return ( + ( + self.values[:rows].unsqueeze(-1), + _scale_view(self.mma, rows, self.input_width).view( + torch.float8_e8m0fnu + ), + ), + (self.packed.weight.values.unsqueeze(-1), self.packed.weight.scale_mma), + self.output[:rows].unsqueeze(-1), + ) + + def preparation_unit(self, name): + from b12x.preparation import PreparedCall + + def quant_call(state): + source = torch.full( + (self.max_tokens, self.slice_width), + 0.125, + dtype=torch.bfloat16, + device=self.output.device, + ) + return PreparedCall( + run=lambda: state.run( + source, + self.slice_values, + self.slice_scales, + self.slice_mma, + ) + ) + + def gemm_call(state): + lhs, rhs, out = self._operands(self.max_tokens) + return PreparedCall( + produce=lambda: (self.values.fill_(0.125), self.mma.fill_(127)), + run=lambda: state.run( + *lhs, *rhs, None, out, self.workspace, self.output.dtype, None + ), + owners=(self.values, self.mma, self.packed, self.output), + ) + + return B12xPreparationUnit( + name="MXFP8_STAGED_INPUT", + key=( + self.max_tokens, + self.input_width, + self.slice_width, + self.output.shape[1], + ), + requests=( + self.quant_plan.request( + name=f"{name}.quant", + prepare_call=quant_call, + benchmark_call=quant_call, + ), + self.gemm_plan.request( + name=f"{name}.gemm", + prepare_call=gemm_call, + benchmark_call=gemm_call, + ), + ), + stage="weights", + ) + + def begin(self, tokens): + if not 0 < tokens <= self.max_tokens: + raise ValueError("staged input token count is outside prepared capacity") + self._tokens, self._column = tokens, 0 + + def append(self, source): + from b12x.quantization import mxfp8 + + rows, column, width = self._tokens, self._column, self.slice_width + if ( + rows == 0 + or source.shape != (rows, width) + or column + width > self.input_width + ): + raise ValueError( + "staged input slice is out of order or has incompatible geometry" + ) + mxfp8.quantize_rows( + source.contiguous(), + self.slice_values[:rows], + self.slice_scales[:rows], + self.slice_mma, + plan=self.quant_plan, + ) + self.values[:rows, column : column + width].copy_(self.slice_values[:rows]) + self.scales[:rows, column // 32 : (column + width) // 32].copy_( + self.slice_scales[:rows] + ) + _scale_view(self.mma, rows, self.input_width)[ + :, :, :, :, column // 128 : (column + width) // 128 + ].copy_(_scale_view(self.slice_mma, rows, width)) + self._column += width + + def finish(self): + from b12x.gemm import blockscaled + + if self._tokens == 0 or self._column != self.input_width: + raise RuntimeError("staged MXFP8 input is incomplete") + lhs, rhs, out = self._operands(self._tokens) + blockscaled.mm( + lhs, + rhs, + plan=self.gemm_plan, + out=out, + workspace=self.workspace, + ) + self._tokens = self._column = 0 + return out.squeeze(-1) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 56841cba17c3..7733d8ab7f47 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -405,6 +405,8 @@ def apply_mxfp8_marlin_linear( size_k: int, bias: torch.Tensor | None = None, use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, + *, + output: torch.Tensor | None = None, ) -> torch.Tensor: reshaped_x = input.reshape(-1, input.shape[-1]) out_shape = input.shape[:-1] + (size_n,) @@ -412,6 +414,18 @@ def apply_mxfp8_marlin_linear( padded_n, padded_k = marlin_repacked_nk(weight, num_bits=8) reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k) + if output is not None and ( + output.shape != (reshaped_x.shape[0], padded_n) + or output.dtype != input.dtype + or output.device != input.device + or not output.is_contiguous() + or output.untyped_storage().data_ptr() == input.untyped_storage().data_ptr() + ): + raise ValueError( + "Marlin output must be contiguous, disjoint from input, and match " + "the padded output shape, activation dtype and device" + ) + use_atomic_add = should_use_atomic_add_reduce( m=reshaped_x.size(0), n=padded_n, @@ -422,7 +436,7 @@ def apply_mxfp8_marlin_linear( output = ops.marlin_gemm( a=reshaped_x, - c=None, + c=output, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, diff --git a/vllm/model_executor/layers/rotary_embedding/compact.py b/vllm/model_executor/layers/rotary_embedding/compact.py new file mode 100644 index 000000000000..e5c2d69aaf76 --- /dev/null +++ b/vllm/model_executor/layers/rotary_embedding/compact.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bounded position tables for unscaled rotary embeddings.""" + +import torch + +from vllm import _custom_ops as ops + +from .base import RotaryEmbedding + + +class CompactRotaryEmbedding(RotaryEmbedding): + """Materialize only positions consumed by a forward in retained storage. + + Frequencies and trigonometric functions use the full-table implementation's + FP32 operations, followed by the same cache-dtype conversion. Capacity is a + maximum row count, not a maximum global position. Buffers are instance-owned + and reused only by serial forwards, including CUDA Graph replay. + """ + + def __init__(self, *args, capacity: int, **kwargs): + if capacity <= 0: + raise ValueError("compact rotary capacity must be positive") + super().__init__(*args, **kwargs, init_cache=False) + self.capacity = capacity + inv_freq = self._compute_inv_freq(self.base) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer( + "freqs_workspace", + torch.empty( + capacity, + self.rotary_dim // 2, + dtype=torch.float32, + device=inv_freq.device, + ), + persistent=False, + ) + self.register_buffer( + "cos_sin_cache", + torch.empty( + capacity, self.rotary_dim, dtype=self.dtype, device=inv_freq.device + ), + persistent=False, + ) + self.register_buffer( + "local_positions", + torch.arange(capacity, dtype=torch.int64, device=inv_freq.device), + persistent=False, + ) + + def materialize(self, positions): + if positions.ndim != 1 or positions.numel() > self.capacity: + raise ValueError("rotary position rows exceed prepared capacity") + rows = positions.numel() + freqs = self.freqs_workspace[:rows] + cache = self.cos_sin_cache[:rows] + half = self.rotary_dim // 2 + torch.mul(positions[:, None], self.inv_freq[None, :], out=freqs) + torch.cos(freqs, out=cache[:, :half]) + torch.sin(freqs, out=cache[:, half:]) + return self.local_positions[:rows], cache + + def forward_cuda(self, positions, query, key=None): + if query.dtype != self.dtype: + raise ValueError("compact rotary cache and query dtypes must match") + local_positions, cache = self.materialize(positions) + ops.rotary_embedding( + local_positions, query, key, self.head_size, cache, self.is_neox_style + ) + return query, key + + def forward_native(self, positions, query, key=None): + local_positions, cache = self.materialize(positions) + return self.forward_static( + local_positions, + query, + key, + self.head_size, + self.rotary_dim, + cache, + self.is_neox_style, + ) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index e6797dee9a21..c5336afced11 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -483,6 +483,73 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: ): quant_config["quant_method"] = "mxfp4" + @staticmethod + def update_model_config_for_parallelism( + model_config: "ModelConfig", parallel_config: "ParallelConfig" + ) -> None: + if model_config.quantization != "exl3": + return + text = model_config.hf_text_config + tp = parallel_config.tensor_parallel_size + if not 2 <= tp <= 24 or parallel_config.enable_expert_parallel: + raise ValueError( + "EXL3 experts require tensor parallel size 2..24 without EP" + ) + original = getattr( + text, "original_num_attention_heads", text.num_attention_heads + ) + text.original_num_attention_heads = original + text.num_attention_heads = round_up(original, tp) + if text.linear_attn_config is not None: + kda = dict(text.linear_attn_config) + original = kda.get("original_num_heads", kda["num_heads"]) + kda["original_num_heads"] = original + kda["num_heads"] = round_up(original, tp) + text.linear_attn_config = kda + model_config.model_arch_config = model_config.get_model_arch_config() + + +class Qwen3DSparkConfig(VerifyAndUpdateConfig): + @staticmethod + def update_model_config_for_parallelism( + model_config: "ModelConfig", parallel_config: "ParallelConfig" + ) -> None: + text = model_config.hf_text_config + heads = getattr(text, "original_num_attention_heads", text.num_attention_heads) + kv_heads = getattr( + text, "original_num_key_value_heads", text.num_key_value_heads + ) + if heads % kv_heads: + raise ValueError("DSpark query heads must form complete KV groups") + tp = parallel_config.tensor_parallel_size + # Preserve the checkpoint's GQA grouping, including supported KV replication. + compatible = heads % tp == 0 and (kv_heads % tp == 0 or tp % kv_heads == 0) + text.original_num_attention_heads = heads + text.original_num_key_value_heads = kv_heads + if getattr(text, "head_dim", None) is None: + text.head_dim = text.hidden_size // heads + text.num_key_value_heads = kv_heads if compatible else round_up(kv_heads, tp) + text.num_attention_heads = ( + heads if compatible else (text.num_key_value_heads * (heads // kv_heads)) + ) + model_config.model_arch_config = model_config.get_model_arch_config() + + +class KimiK3MLADraftConfig(VerifyAndUpdateConfig): + @staticmethod + def update_model_config_for_parallelism( + model_config: "ModelConfig", parallel_config: "ParallelConfig" + ) -> None: + text = model_config.hf_text_config + original = getattr( + text, "original_num_attention_heads", text.num_attention_heads + ) + text.original_num_attention_heads = original + text.num_attention_heads = round_up( + original, parallel_config.tensor_parallel_size + ) + model_config.model_arch_config = model_config.get_model_arch_config() + class GptOssForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod @@ -1131,6 +1198,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig, "DeepseekV41ForCausalLM": DeepseekV41ForCausalLMConfig, "DSparkV41DraftModel": DeepseekV41ForCausalLMConfig, + "DFlash2KimiK3Model": KimiK3MLADraftConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 @@ -1151,6 +1219,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "JinaVLForRanking": JinaVLForSequenceClassificationConfig, "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig, "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig, + "K3DSparkModel": KimiK3MLADraftConfig, "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, "LlamaBidirectionalModel": LlamaBidirectionalConfig, "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig, @@ -1163,6 +1232,8 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "NomicBertModel": NomicBertModelConfig, "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig, "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, + "Qwen3DSparkModel": Qwen3DSparkConfig, + "Qwen3OmniDSparkModel": Qwen3DSparkConfig, "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig, "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForCausalLM": Qwen3_5ForCausalLMConfig, diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3b5b01f5afb2..5f09a411b0b3 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -4,6 +4,7 @@ import dataclasses import io from collections.abc import Iterable +from typing import Any, ClassVar import torch import torch.nn.functional as F @@ -11,33 +12,40 @@ from transformers import Qwen3Config from vllm import _custom_ops as ops +from vllm import envs from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, ) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.compact import CompactRotaryEmbedding from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.parameter import copy_tensor_parallel_shard from vllm.model_executor.weight_transfer import flush_weight_transfers from vllm.multimodal.inputs import NestedTensors from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes +from vllm.utils.math_utils import round_up from vllm.v1.attention.backend import AttentionType from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -63,6 +71,30 @@ _SLIDING_ATTENTION = "sliding_attention" +_STREAMED_AUX_MIN_TOKENS = 1024 + + +def _enable_dflash_tp_padding(layer: nn.Module) -> None: + """Zero checkpoint-absent tails before optional online MXFP8 conversion.""" + from vllm.model_executor.layers.quantization.online.mxfp8 import ( + Mxfp8OnlineLinearMethod, + ) + + if not isinstance( + layer.quant_method, (UnquantizedLinearMethod, Mxfp8OnlineLinearMethod) + ): + raise ValueError( + "DFlash/DSpark TP padding supports BF16/FP16 checkpoints with " + "unquantized or online MXFP8 linear weights" + ) + for parameter in layer.parameters(recurse=False): + parameter.allow_tp_padding = True + + +def _dflash_padded_width(width: int, tp: int) -> int: + # Whole 128-channel tiles preserve MXFP8 block boundaries and avoid extra + # activation-padding copies in the W8A16 projection kernels. + return width if width % tp == 0 else round_up(width, tp * 128) def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: @@ -162,6 +194,8 @@ def _resolve_layer_attention( class DFlashAttention(Attention): """Attention whose small draft KV is replicated across DCP ranks.""" + dcp_replicated: ClassVar[bool] = True + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: dcp_replicated = vllm_config.parallel_config.decode_context_parallel_size > 1 if self.sliding_window is not None: @@ -214,6 +248,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", attn_type: str = AttentionType.DECODER, + pad_heads: bool = False, ) -> None: super().__init__() self.layer_name = prefix @@ -249,13 +284,33 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.o_proj", ) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - is_neox_style=is_neox_style, - rope_parameters=rope_parameters, - ) + if pad_heads: + _enable_dflash_tp_padding(self.qkv_proj) + _enable_dflash_tp_padding(self.o_proj) + + if envs.VLLM_DFLASH_COMPACT_ROPE: + params = rope_parameters or {} + if params.get("rope_type", "default") != "default" or set(params) - { + "rope_type", + "rope_theta", + }: + raise ValueError("DFlash compact RoPE requires unscaled full-head RoPE") + self.rotary_emb = CompactRotaryEmbedding( + self.head_dim, + self.head_dim, + max_position, + params.get("rope_theta", 10000), + is_neox_style, + torch.get_default_dtype(), + capacity=get_current_vllm_config().scheduler_config.max_num_batched_tokens, + ) + else: + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position, + is_neox_style=is_neox_style, + rope_parameters=rope_parameters, + ) self.attention_sink_bias = ( torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) @@ -360,14 +415,24 @@ def __init__( rope_parameters=config.rope_parameters, prefix=f"{prefix}.self_attn", attn_type=attn_type, + pad_heads=config.num_attention_heads + != getattr( + config, "original_num_attention_heads", config.num_attention_heads + ), + ) + intermediate_size = _dflash_padded_width( + config.intermediate_size, get_tensor_model_parallel_world_size() ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, + intermediate_size=intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=f"{prefix}.mlp", ) + if intermediate_size != config.intermediate_size: + _enable_dflash_tp_padding(self.mlp.gate_up_proj) + _enable_dflash_tp_padding(self.mlp.down_proj) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -476,17 +541,87 @@ def __init__( ] ) if self.use_aux_hidden_state: - self.fc = ReplicatedLinear( + if ( + envs.VLLM_DFLASH_AUX_MXFP8_STREAMING + and envs.VLLM_DFLASH_AUX_BF16_STAGING + ): + raise ValueError("Select either MXFP8 or BF16 auxiliary input staging") + fc_cls = ReplicatedLinear + fc_kwargs = {} + fc_output_size = self.config.hidden_size + if envs.VLLM_DFLASH_SHARD_AUX_PROJECTION: + fc_cls = ColumnParallelLinear + fc_kwargs["gather_output"] = True + fc_output_size = _dflash_padded_width( + fc_output_size, get_tensor_model_parallel_world_size() + ) + self.fc = fc_cls( input_size=_get_dflash_fc_input_size( vllm_config, ), - output_size=self.config.hidden_size, + output_size=fc_output_size, bias=False, params_dtype=vllm_config.model_config.dtype, quant_config=self.quant_config, prefix=maybe_prefix(prefix, "fc"), return_bias=False, + **fc_kwargs, ) + if fc_output_size != self.config.hidden_size: + _enable_dflash_tp_padding(self.fc) + if envs.VLLM_DFLASH_AUX_MXFP8_STREAMING: + from vllm.model_executor.kernels.linear.mxfp8.b12x import ( + B12xMxfp8LinearKernel, + Mxfp8LinearLayerConfig, + ) + from vllm.model_executor.layers.quantization.online.mxfp8 import ( + Mxfp8OnlineLinearMethod, + ) + + method = self.fc.quant_method + if not isinstance(method, Mxfp8OnlineLinearMethod): + raise ValueError( + "DFlash MXFP8 staging requires online MXFP8 FC weights" + ) + method.kernel = B12xMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + method.use_a16 = False + # The wrapper prepares small FC calls separately from staged + # prefill, so their activation scratch is not reserved twice. + self.fc.b12x_preparation_suppressed = True + elif envs.VLLM_DFLASH_AUX_BF16_STAGING: + from vllm.model_executor.kernels.linear.mxfp8.marlin import ( + MarlinMxfp8LinearKernel, + Mxfp8LinearLayerConfig, + ) + from vllm.model_executor.layers.quantization.online.mxfp8 import ( + Mxfp8OnlineLinearMethod, + ) + + method = self.fc.quant_method + if not isinstance(method, Mxfp8OnlineLinearMethod): + raise ValueError( + "DFlash BF16 staging requires online MXFP8 FC weights" + ) + # Preserve BF16 activations through the same W8A16 linear path + # used without staging; weight sharding does not quantize inputs. + method.kernel = MarlinMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + self._aux_projection_tp_size = ( + get_tensor_model_parallel_world_size() + if envs.VLLM_DFLASH_SHARD_AUX_PROJECTION + else 1 + ) + self._streamed_aux_layer_ids = tuple( + get_eagle3_aux_layers_from_config(vllm_config.speculative_config) or () + ) + self._target_hidden_size = int( + getattr(self.config, "target_hidden_size", None) or self.config.hidden_size + ) + self._streamed_aux_accumulator = None + self._streamed_aux_scratch = None + self._streamed_aux_tokens = self._streamed_aux_index = 0 + self._streamed_aux_generation = self._completed_stream_generation = 0 + self._consumed_stream_generation = 0 + self._completed_stream_result = None self.hidden_norm = RMSNorm( self.config.hidden_size, eps=self.config.rms_norm_eps, @@ -512,6 +647,140 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds) return embeds + def bind_auxiliary_stream( + self, + accumulator: Any, + scratch: torch.Tensor, + ) -> None: + """Bind retained storage for large target auxiliary-state projections.""" + if scratch.ndim != 2 or int(scratch.shape[1]) != int(self.config.hidden_size): + raise ValueError( + "DFlash auxiliary output scratch must have shape " + f"[max_tokens,{self.config.hidden_size}], got {tuple(scratch.shape)}" + ) + if int(scratch.shape[1]) < self._target_hidden_size: + raise ValueError( + "DFlash auxiliary output scratch also stages target residual " + f"sums and must be at least {self._target_hidden_size} wide, " + f"got {scratch.shape[1]}" + ) + self._streamed_aux_accumulator = accumulator + self._streamed_aux_scratch = scratch + + def can_stream_auxiliary_states( + self, + layer_ids: tuple[int, ...], + hidden_states: torch.Tensor, + ) -> bool: + """Return whether a target forward can release auxiliary states early.""" + accumulator = self._streamed_aux_accumulator + scratch = self._streamed_aux_scratch + if accumulator is None or scratch is None or not self.use_aux_hidden_state: + return False + if hidden_states.ndim != 2 or hidden_states.shape[0] < _STREAMED_AUX_MIN_TOKENS: + return False + if tuple(layer_ids) != self._streamed_aux_layer_ids: + return False + if hidden_states.shape[1] * len(layer_ids) != accumulator.input_width: + return False + if hidden_states.shape[0] > accumulator.max_tokens: + return False + if ( + hidden_states.dtype != scratch.dtype + or hidden_states.device != scratch.device + ): + return False + return not (hidden_states.is_cuda and torch.cuda.is_current_stream_capturing()) + + def begin_auxiliary_stream(self, hidden_states: torch.Tensor) -> None: + """Start one ordered sequence of target auxiliary-state slices.""" + if not self.can_stream_auxiliary_states( + self._streamed_aux_layer_ids, hidden_states + ): + raise RuntimeError( + "DFlash auxiliary streaming was started for an unsupported geometry" + ) + self._streamed_aux_tokens = int(hidden_states.shape[0]) + self._streamed_aux_index = 0 + self._streamed_aux_generation += 1 + self._completed_stream_result = None + self._streamed_aux_accumulator.begin(self._streamed_aux_tokens) + logger.info_once( + "DFlash staged auxiliary projection is active: " + "tokens=%d target_width=%d slices=%d.", + self._streamed_aux_tokens, + self._target_hidden_size, + len(self._streamed_aux_layer_ids), + ) + + def accumulate_auxiliary_state( + self, + primary: torch.Tensor, + residual: torch.Tensor | None, + ) -> None: + """Store one complete target state in the retained linear input.""" + if self._streamed_aux_index >= len(self._streamed_aux_layer_ids): + raise RuntimeError("DFlash received too many auxiliary states") + expected_shape = (self._streamed_aux_tokens, self._target_hidden_size) + if tuple(primary.shape) != expected_shape: + raise ValueError( + "DFlash auxiliary state shape changed during streaming: " + f"got={tuple(primary.shape)}, expected={expected_shape}" + ) + source = primary + if residual is not None: + if tuple(residual.shape) != expected_shape: + raise ValueError( + "DFlash auxiliary residual shape changed during streaming: " + f"got={tuple(residual.shape)}, expected={expected_shape}" + ) + assert self._streamed_aux_scratch is not None + source = self._streamed_aux_scratch[ + : self._streamed_aux_tokens, : self._target_hidden_size + ] + torch.add(primary, residual, out=source) + self._streamed_aux_accumulator.append(source) + self._streamed_aux_index += 1 + + def finish_auxiliary_stream(self) -> torch.Tensor: + """Project the assembled input with the draft's FC weights.""" + expected = len(self._streamed_aux_layer_ids) + if self._streamed_aux_index != expected: + raise RuntimeError( + "DFlash auxiliary stream ended before every configured state " + f"was received: received={self._streamed_aux_index}, " + f"expected={expected}" + ) + output = self._gather_auxiliary_projection( + self._streamed_aux_accumulator.finish() + ) + self._completed_stream_generation = self._streamed_aux_generation + self._completed_stream_result = output + return output + + def _gather_auxiliary_projection(self, output: torch.Tensor) -> torch.Tensor: + if getattr(self, "_aux_projection_tp_size", 1) > 1: + output = tensor_model_parallel_all_gather(output, dim=-1) + return output[..., : self.config.hidden_size] + + def is_streamed_context_states(self, states: list[torch.Tensor]) -> bool: + """Claim one completed streamed projection exactly once.""" + if len(states) != 1 or self._completed_stream_result is None: + return False + candidate = states[0] + expected = self._completed_stream_result + matches = ( + self._completed_stream_generation > self._consumed_stream_generation + and candidate is expected + and candidate.shape == expected.shape + and candidate.dtype == expected.dtype + and candidate.device == expected.device + and candidate.data_ptr() == expected.data_ptr() + ) + if matches: + self._consumed_stream_generation = self._completed_stream_generation + return matches + def _build_context_kv_buffers( self, layers_attn: list[nn.Module], @@ -733,8 +1002,12 @@ def precompute_and_store_context_kv( # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy). # In-place RoPE: pass K as the "query" arg with key=None. all_k_flat = all_k_normed.view(L * num_ctx, kv) - positions_repeated = context_positions.repeat(L) - cos_sin_cache = self._rope_cos_sin_cache + rotary = self.layers[0].self_attn.rotary_emb + if isinstance(rotary, CompactRotaryEmbedding): + rope_positions, cos_sin_cache = rotary.materialize(context_positions) + else: + rope_positions, cos_sin_cache = context_positions, self._rope_cos_sin_cache + positions_repeated = rope_positions.repeat(L) if cos_sin_cache.dtype != all_k_flat.dtype: cos_sin_cache = cos_sin_cache.to(dtype=all_k_flat.dtype) ops.rotary_embedding( @@ -751,7 +1024,7 @@ def precompute_and_store_context_kv( # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) - per_layer = isinstance(context_slot_mapping, (list, tuple)) + per_layer = isinstance(context_slot_mapping, list | tuple) for i in range(L): slot_mapping = ( context_slot_mapping[i] if per_layer else context_slot_mapping @@ -798,10 +1071,17 @@ def _preprocess( if "attention_sink_bias" in name: # Sink bias is per-head; shard it across TP ranks like the # attention heads themselves. - heads_per_rank = loaded_weight.shape[0] // tp_size - loaded_weight = loaded_weight.narrow( - 0, tp_rank * heads_per_rank, heads_per_rank + heads_per_rank = self.config.num_attention_heads // tp_size + shard = loaded_weight.new_empty(heads_per_rank) + copy_tensor_parallel_shard( + shard, + loaded_weight, + 0, + tp_rank * heads_per_rank, + heads_per_rank, + allow_padding=True, ) + loaded_weight = shard yield name, loaded_weight def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -918,11 +1198,101 @@ def combine_hidden_states( "means the draft model's target_layer_ids reference layers that " "do not exist in the target model (incompatible draft/target pair)." ) - result = self.model.fc(hidden_states) + accumulator = self.model._streamed_aux_accumulator + if ( + accumulator is not None + and hidden_states.shape[0] >= _STREAMED_AUX_MIN_TOKENS + ): + # Large non-streamed calls, including graph capture, use the same + # staged buffers. Only smaller calls need a separate linear plan. + accumulator.begin(hidden_states.shape[0]) + for source in hidden_states.split(accumulator.slice_width, dim=-1): + accumulator.append(source) + result = self.model._gather_auxiliary_projection(accumulator.finish()) + else: + result = self.model.fc(hidden_states) + result = result[..., : self.model.config.hidden_size] if needs_squeeze: result = result.squeeze(0) return result + def bind_target_auxiliary_stream(self, target_model, scratch) -> None: + """Stage large Kimi auxiliary states with the selected input precision.""" + bf16_staging = envs.VLLM_DFLASH_AUX_BF16_STAGING + if not envs.VLLM_DFLASH_AUX_MXFP8_STREAMING and not bf16_staging: + return + from vllm.model_executor.kernels.linear.mxfp8.staged import ( + B12xMxfp8InputAccumulator, + ) + from vllm.utils.b12x import set_b12x_preparation_provider + + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + target = getattr(language_model, "model", language_model) + setter = getattr(target, "set_aux_hidden_state_projector", None) + if not callable(setter) or not self.model.use_aux_hidden_state: + raise ValueError( + "DFlash auxiliary staging requires a target auxiliary-state hook" + ) + fc = self.model.fc + if bf16_staging: + from vllm.model_executor.kernels.linear.bf16_staging import ( + Bf16InputAccumulator, + ) + + accumulator = Bf16InputAccumulator( + fc, scratch, self.model._target_hidden_size + ) + self.model.bind_auxiliary_stream(accumulator, scratch) + setter(self.model) + logger.info_once("DFlash auxiliary projection retains BF16 input slices.") + return + local_width = fc.b12x_mxfp8_packed_weight.out_features + output = scratch + if local_width != scratch.shape[1]: + output = torch.empty( + (scratch.shape[0], local_width), + dtype=scratch.dtype, + device=scratch.device, + ) + accumulator = B12xMxfp8InputAccumulator( + fc, output, self.model._target_hidden_size + ) + self.model.bind_auxiliary_stream(accumulator, scratch) + setter(self.model) + set_b12x_preparation_provider(self, self) + logger.info_once("DFlash auxiliary projection uses prepared MXFP8 staging.") + + def get_b12x_preparation_units(self, layer, workload): + accumulator = self.model._streamed_aux_accumulator + if workload.stage != "weights" or accumulator is None: + return () + if not callable(getattr(accumulator, "preparation_unit", None)): + # BF16 input storage uses the ordinary Marlin FC without B12X plans. + return () + capacity = min(workload.max_tokens, _STREAMED_AUX_MIN_TOKENS - 1) + small_workload = dataclasses.replace( + workload, + max_tokens=capacity, + token_counts=tuple( + sorted({capacity, *(n for n in workload.token_counts if n < capacity)}) + ), + fixed_token_counts=tuple( + n for n in workload.fixed_token_counts if n < capacity + ), + ) + linear = self.model.fc.b12x_linear + return ( + linear.unit(small_workload, name=f"linear.mxfp8.{linear.layer_name}"), + accumulator.preparation_unit("dflash.auxiliary"), + ) + + def is_streamed_context_states(self, states: list[torch.Tensor]) -> bool: + return self.model.is_streamed_context_states(states) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights = {} includes_draft_id_mapping = False diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 7318d225f8d2..ec1172a2b370 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -20,6 +20,7 @@ import torch import torch.nn as nn +from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.linear import ReplicatedLinear @@ -31,6 +32,7 @@ ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, + VocabParallelEmbedding, ) from vllm.platforms import current_platform @@ -53,9 +55,9 @@ class DSparkMarkovHead(nn.Module): (``draft_vocab_size``) added to the base draft logits. The two sizes coincide for full-vocab drafts. - Both weights are replicated because the head runs sequentially for every - draft position. Sharding them would add an all-reduce and a full-vocab - gather to each position. + Weights are replicated by default. An opt-in TP-sharded layout reduces + draft residency and supports rank-local greedy sampling. Keeping W1 + replicated avoids an embedding all-reduce at each draft position. """ def __init__( @@ -68,14 +70,26 @@ def __init__( retain_weight_for_gather: bool = False, ) -> None: super().__init__() - self.markov_w1 = nn.Embedding(vocab_size, markov_rank) + self.shard_across_tp = envs.VLLM_DSPARK_SHARD_MARKOV_HEAD + self.replicate_w1 = envs.VLLM_DSPARK_REPLICATE_MARKOV_W1 + if self.replicate_w1 and not self.shard_across_tp: + raise ValueError("Replicated Markov W1 requires a TP-sharded Markov head") + if self.shard_across_tp and retain_weight_for_gather: + raise ValueError("Gathered Markov top-k requires replicated weights") + self.markov_w1 = ( + VocabParallelEmbedding( + vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1") + ) + if self.shard_across_tp and not self.replicate_w1 + else nn.Embedding(vocab_size, markov_rank) + ) self.markov_w2 = ParallelLMHead( draft_vocab_size, markov_rank, bias=False, quant_config=quant_config, prefix=maybe_prefix(prefix, "markov_w2"), - disable_tp=True, + disable_tp=not self.shard_across_tp, ) self.markov_w2._retain_weight_for_gather = retain_weight_for_gather self.markov_w2.is_w4a16_nvfp4 = False @@ -92,6 +106,14 @@ def bias( """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).""" return logits_processor(self.markov_w2, markov_embed) + def local_bias( + self, markov_embed: torch.Tensor, logits_processor: LogitsProcessor + ) -> torch.Tensor: + """Return the unprocessed bias shard matching the target LM-head shard.""" + if not self.shard_across_tp: + raise RuntimeError("Local Markov bias requires a TP-sharded head") + return logits_processor._apply_head(self.markov_w2, markov_embed, None) + def apply_bias_gathered( self, markov_embed: torch.Tensor, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 0b3324b394c2..268592e9d242 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -643,6 +643,10 @@ "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"), + "DFlash2KimiK3Model": ( + "vllm.models.kimi_k3.nvidia.dflash2_mla", + "DFlash2KimiK3ForCausalLM", + ), # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash # implementation. EAGLEConfig rewrites a dflash draft's architecture to # DFlash{arch} unless it already starts or ends with "DFlash" (see diff --git a/vllm/models/kimi_k3/nvidia/dflash2_mla.py b/vllm/models/kimi_k3/nvidia/dflash2_mla.py new file mode 100644 index 000000000000..a9b518e72eb5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/dflash2_mla.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MLA DFlash2 draft with grouped convolutions and a candidate transition head.""" + +from collections.abc import Iterable + +import torch + +from vllm.compilation.backends import set_model_tag +from vllm.config import VllmConfig +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dflash import _resolve_layer_attention +from vllm.model_executor.models.qwen3_dflash2 import ( + CandidateSelector, + DFlashGroupedConv, +) +from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.models.kimi_k3.nvidia.dspark_mla import ( + K3DSparkDecoderLayer, + K3DSparkForCausalLM, + K3DSparkModel, +) + + +class DFlash2KimiK3DecoderLayer(K3DSparkDecoderLayer): + @staticmethod + def get_attention_window(config, layer_idx: int) -> int: + window, causal = _resolve_layer_attention(config, layer_idx) + if causal: + raise ValueError("Kimi-K3 MLA DFlash2 requires non-causal draft layers") + # Zero explicitly preserves full attention instead of inheriting the + # DSpark tail-window override. Sliding layers retain checkpoint windows. + return window or 0 + + def __init__(self, *, vllm_config, config, layer_idx, start_layer_id, prefix): + super().__init__( + vllm_config=vllm_config, + config=config, + layer_idx=layer_idx, + start_layer_id=start_layer_id, + prefix=prefix, + ) + draft = config.dflash_config + spec = vllm_config.speculative_config + assert spec is not None + conv_args = dict( + hidden_size=config.hidden_size, + taps=int(draft["conv_kernel_size"]), + group_size=int(draft["conv_group_size"]), + block_size=1 + spec.num_speculative_tokens, + params_dtype=vllm_config.model_config.dtype, + quant_config=None, + ) + self.attention_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "attention_conv") + ) + self.mlp_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "mlp_conv") + ) + + def forward(self, positions, hidden_states, residual, rope_cos_sin_cache=None): + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + hidden_states, coefficients = self.attention_conv.prepare(hidden_states) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + rope_cos_sin_cache=rope_cos_sin_cache, + ) + hidden_states = self.attention_conv.finish(hidden_states, coefficients) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + hidden_states, coefficients = self.mlp_conv.prepare(hidden_states) + hidden_states = self.mlp(hidden_states) + return self.mlp_conv.finish(hidden_states, coefficients), residual + + +class DFlash2KimiK3Model(K3DSparkModel): + decoder_layer_cls = DFlash2KimiK3DecoderLayer + has_markov_head = False + + def __init__(self, *, vllm_config, start_layer_id, prefix): + super().__init__( + vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix + ) + draft = self.config.dflash_config + self.input_embedding_scale = float(draft.get("input_embedding_scale", 1.0)) + with set_model_tag("dflash2_candidate_selector"): + self.candidate_selector = CandidateSelector( + hidden_size=self.config.hidden_size, + vocab_size=self.config.vocab_size, + rank=int(draft["selector_rank"]), + top_k=int(draft["selector_top_k"]), + params_dtype=vllm_config.model_config.dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "candidate_selector"), + ) + + def embed_input_ids(self, input_ids): + return super().embed_input_ids(input_ids) * self.input_embedding_scale + + +class DFlash2KimiK3ForCausalLM(K3DSparkForCausalLM): + model_cls = DFlash2KimiK3Model + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + spec = vllm_config.speculative_config + assert spec is not None + config = spec.draft_model_config.hf_config + draft = config.dflash_config + if draft.get("attention_mode") != "mla" or config.q_lora_rank is None: + raise ValueError("Kimi-K3 DFlash2 requires latent Q and KV projections") + config.target_layer_ids = list(draft["target_layer_ids"]) + if not config.target_layer_ids: + raise ValueError("Kimi-K3 DFlash2 requires target auxiliary layer IDs") + config.target_hidden_size = getattr( + config, "target_hidden_size", config.hidden_size + ) + config.num_target_layers = len(config.target_layer_ids) + config.draft_vocab_size = config.vocab_size + super().__init__(vllm_config=vllm_config, prefix=prefix) + softcap = float(draft.get("final_logit_softcapping") or 0.0) + self.candidate_logits_processor = LogitsProcessor( + vllm_config.model_config.get_vocab_size(), + scale=float(draft.get("output_multiplier", 1.0)), + soft_cap=softcap or None, + ) + + def compute_candidates(self, hidden_states): + return self.candidate_logits_processor.get_top_k_tokens( + self.lm_head, hidden_states, self.model.candidate_selector.top_k + ) + + def get_draft_attn_causal(self) -> list[bool]: + return [False] * len(self.model.layers) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def remap(): + names = { + "fc": "context_proj", + "hidden_norm": "context_norm", + "norm": "final_norm", + } + for name, tensor in weights: + first, separator, suffix = name.partition(".") + yield names.get(first, first) + separator + suffix, tensor + + return super().load_weights(remap()) diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 87f9e94d9b66..6a86ac6694e6 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -2,19 +2,30 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """K3 dense MLA draft model for DSpark speculative decoding.""" -from collections.abc import Iterable +from collections.abc import Iterable, Iterator +from contextlib import contextmanager import torch import torch.nn as nn +import torch.nn.functional as F import vllm._custom_ops as ops +from vllm import envs from vllm.config import VllmConfig +from vllm.distributed import ( + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce_in_place, +) +from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -25,14 +36,38 @@ from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention from vllm.models.kimi_k3.nvidia.model import KimiMLP +from vllm.utils.b12x import ( + B12xPreparationUnit, + B12xWorkload, + PreparationResourceUnavailableError, + set_b12x_preparation_provider, +) from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.worker.workspace import current_workspace_manager +logger = init_logger(__name__) +_COMPACT_ROPE_PROTECTED_IDS: set[int] = set() +_STREAMED_AUX_MIN_TOKENS = 1024 + _GROUPED_KV_CACHE_DTYPES = frozenset( {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"} ) +@contextmanager +def protect_k3_compact_rope_sources(model: nn.Module) -> Iterator[None]: + """Prevent a draft from releasing position tables owned by ``model``.""" + previous = _COMPACT_ROPE_PROTECTED_IDS.copy() + _COMPACT_ROPE_PROTECTED_IDS.update( + id(module) for module in model.modules() if hasattr(module, "cos_sin_cache") + ) + try: + yield + finally: + _COMPACT_ROPE_PROTECTED_IDS.clear() + _COMPACT_ROPE_PROTECTED_IDS.update(previous) + + def _duplicate_context_kv_weights( weights: Iterable[tuple[str, torch.Tensor]], num_layers: int ) -> Iterable[tuple[str, torch.Tensor]]: @@ -55,7 +90,37 @@ def _duplicate_context_kv_weights( yield f"context_kv_proj.{param_name}", fused_weight +def _fill_compact_rope_cache( + positions: torch.Tensor, + inv_freq: torch.Tensor, + freqs_workspace: torch.Tensor, + cache_workspace: torch.Tensor, + *, + mscale: float, +) -> torch.Tensor: + """Materialize RoPE rows consumed by one Kimi-K3 draft forward.""" + num_positions = int(positions.shape[0]) + if num_positions > int(freqs_workspace.shape[0]): + raise ValueError( + "Kimi-K3 DSpark compact RoPE workspace is too small: " + f"positions={num_positions}, capacity={freqs_workspace.shape[0]}." + ) + freqs = freqs_workspace[:num_positions] + cache = cache_workspace[:num_positions] + half_dim = int(inv_freq.shape[0]) + torch.mul(positions[:, None], inv_freq[None, :], out=freqs) + torch.cos(freqs, out=cache[:, :half_dim]) + torch.sin(freqs, out=cache[:, half_dim:]) + if mscale != 1.0: + cache.mul_(mscale) + return cache + + class K3DSparkDecoderLayer(nn.Module): + @staticmethod + def get_attention_window(config, layer_idx: int) -> int | None: + return None + def __init__( self, *, @@ -83,6 +148,7 @@ def __init__( ), use_rope=True, non_causal_multi_token_decode=True, + draft_kv_window=self.get_attention_window(config, layer_idx), ) # Both row-parallel outputs stay un-reduced; their all-reduces are fused # into the RMSNorm that follows via fused_allreduce_rms_norm. @@ -93,6 +159,7 @@ def __init__( hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, + pad_intermediate=True, prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -105,6 +172,7 @@ def forward( positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, + rope_cos_sin_cache: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if residual is None: # First layer: hidden_states is the (already reduced) embedding. @@ -118,6 +186,7 @@ def forward( hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, + rope_cos_sin_cache=rope_cos_sin_cache, ) hidden_states, residual = fused_allreduce_rms_norm( hidden_states, residual, self.post_attention_layernorm @@ -129,6 +198,9 @@ def forward( class K3DSparkModel(nn.Module): + decoder_layer_cls = K3DSparkDecoderLayer + has_markov_head = True + def __init__( self, *, @@ -144,21 +216,46 @@ def __init__( # The frozen target embedding is aliased after the draft checkpoint loads. self.embed_tokens: nn.Module | None = None - self.context_proj = ReplicatedLinear( - self.config.target_hidden_size * self.config.num_target_layers, - self.config.hidden_size, - bias=False, - return_bias=False, - quant_config=self.quant_config, - prefix=maybe_prefix(prefix, "context_proj"), + parallel_config = getattr(vllm_config, "parallel_config", None) + tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1)) + self.context_proj_sharded = ( + tp_size > 1 and self.config.hidden_size % tp_size == 0 ) + context_input_size = ( + self.config.target_hidden_size * self.config.num_target_layers + ) + context_prefix = maybe_prefix(prefix, "context_proj") + if self.context_proj_sharded: + # Target auxiliary states are identical across TP ranks. Each rank + # retains and evaluates only its output rows; the gathered result + # preserves the draft hidden-state layout exactly. Keep this one + # projection in BF16: large-prefill streaming reads contiguous + # input-column slices before the complete target state exists. + self.context_proj = ColumnParallelLinear( + context_input_size, + self.config.hidden_size, + bias=False, + gather_output=True, + return_bias=False, + quant_config=None, + prefix=context_prefix, + ) + else: + self.context_proj = ReplicatedLinear( + context_input_size, + self.config.hidden_size, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=context_prefix, + ) self.context_norm = RMSNorm( self.config.hidden_size, eps=self.config.rms_norm_eps ) self.layers = nn.ModuleList( [ - K3DSparkDecoderLayer( + self.decoder_layer_cls( vllm_config=vllm_config, config=self.config, layer_idx=layer_idx, @@ -182,16 +279,167 @@ def __init__( disable_tp=True, ) self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) - self.markov_head = DSparkMarkovHead( - self.config.vocab_size, - self.config.draft_vocab_size, - self.config.markov_rank, - prefix=maybe_prefix(prefix, "markov_head"), + self.markov_head = ( + DSparkMarkovHead( + self.config.vocab_size, + self.config.draft_vocab_size, + self.config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + if self.has_markov_head + else None ) self._max_num_context_tokens = ( vllm_config.scheduler_config.max_num_batched_tokens ) + self._streamed_aux_layer_ids = tuple( + int(layer_id) + 1 + for layer_id in getattr(self.config, "target_layer_ids", ()) + ) + self._streamed_aux_scratch: torch.Tensor | None = None + self._streamed_aux_tokens = 0 + self._streamed_aux_index = 0 + self._streamed_aux_generation = 0 + self._completed_stream_generation = 0 + self._consumed_stream_generation = 0 + self._completed_stream_result: torch.Tensor | None = None + self._context_local_width = ( + int(self.config.hidden_size) // tp_size + if self.context_proj_sharded + else int(self.config.hidden_size) + ) + self._context_local_start = ( + int(getattr(self.context_proj, "tp_rank", 0)) * self._context_local_width + ) + context_proj_weight = self.context_proj.weight + if self.context_proj_sharded: + self.register_buffer( + "_streamed_context_states", + torch.empty( + self._max_num_context_tokens, + self._context_local_width, + dtype=context_proj_weight.dtype, + device=context_proj_weight.device, + ), + persistent=False, + ) + else: + self.register_buffer( + "_streamed_context_states", + torch.empty( + 0, + dtype=context_proj_weight.dtype, + device=context_proj_weight.device, + ), + persistent=False, + ) + self._compact_rope_enabled = bool(envs.VLLM_DSPARK_COMPACT_ROPE) + if self._compact_rope_enabled: + self._init_compact_rope() + + def _init_compact_rope(self) -> None: + rotary_modules = [layer.self_attn.rotary_emb for layer in self.layers] + if not rotary_modules or any(rotary is None for rotary in rotary_modules): + raise RuntimeError( + "Kimi-K3 DSpark compact RoPE requires rotary draft layers." + ) + rotary = rotary_modules[0] + assert rotary is not None + if any(candidate is not rotary for candidate in rotary_modules[1:]): + raise RuntimeError( + "Kimi-K3 DSpark compact RoPE requires every draft layer to " + "share one immutable rotary embedding." + ) + if not hasattr(rotary, "scaling_factor"): + raise TypeError( + "Kimi-K3 DSpark compact RoPE requires a YaRN rotary embedding, " + f"got {type(rotary).__name__}." + ) + + full_cache = rotary.cos_sin_cache + if full_cache.dtype != torch.float32 or full_cache.ndim != 2: + raise TypeError( + "Kimi-K3 DSpark compact RoPE requires a 2-D fp32 source cache, " + f"got shape={tuple(full_cache.shape)}, dtype={full_cache.dtype}." + ) + with torch.device(full_cache.device): + inv_freq = rotary._compute_inv_freq(rotary.scaling_factor) # noqa: SLF001 + inv_freq = inv_freq.to(dtype=torch.float32) + half_dim = int(inv_freq.shape[0]) + if int(full_cache.shape[1]) != 2 * half_dim: + raise ValueError( + "Kimi-K3 DSpark compact RoPE frequency width does not match " + f"its source table: inv_freq={half_dim}, " + f"cache={full_cache.shape[1]}." + ) + + self.register_buffer("_compact_rope_inv_freq", inv_freq, persistent=False) + self.register_buffer( + "_compact_rope_freqs", + torch.empty( + (self._max_num_context_tokens, half_dim), + dtype=torch.float32, + device=full_cache.device, + ), + persistent=False, + ) + self.register_buffer( + "_compact_rope_cache", + torch.empty( + (self._max_num_context_tokens, 2 * half_dim), + dtype=torch.float32, + device=full_cache.device, + ), + persistent=False, + ) + self.register_buffer( + "_compact_rope_positions", + torch.arange( + self._max_num_context_tokens, + dtype=torch.int64, + device=full_cache.device, + ), + persistent=False, + ) + self._compact_rope_mscale = float(rotary.mscale) + + if id(rotary) in _COMPACT_ROPE_PROTECTED_IDS: + logger.warning_once( + "Kimi-K3 DSpark compact RoPE retains a position table shared " + "with the target model." + ) + return + + released_bytes = full_cache.numel() * full_cache.element_size() + rotary.cos_sin_cache = torch.empty( + (0, 2 * half_dim), + dtype=torch.float32, + device=full_cache.device, + ) + logger.info_once( + "Kimi-K3 DSpark compact RoPE released %.2f MiB/rank and " + "materializes at most %d fp32 rows per forward.", + released_bytes / (1024**2), + self._max_num_context_tokens, + ) + + def _get_rope_inputs( + self, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + rotary = self.layers[0].self_attn.rotary_emb + assert rotary is not None + if not self._compact_rope_enabled: + return positions, rotary.cos_sin_cache + cache = _fill_compact_rope_cache( + positions, + self._compact_rope_inv_freq, + self._compact_rope_freqs, + self._compact_rope_cache, + mscale=self._compact_rope_mscale, + ) + return self._compact_rope_positions[: positions.shape[0]], cache + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: assert self.embed_tokens is not None return self.embed_tokens(input_ids) @@ -199,6 +447,160 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.context_norm(self.context_proj(hidden_states)) + def bind_auxiliary_stream_scratch(self, scratch: torch.Tensor) -> bool: + """Bind caller-owned storage used to form one target auxiliary state.""" + expected_width = int(self.config.target_hidden_size) + if scratch.ndim != 2 or scratch.shape[1] != expected_width: + logger.warning_once( + "Kimi-K3 DSpark auxiliary streaming is disabled because its " + "caller-owned scratch has shape %s instead of [tokens, %d].", + tuple(scratch.shape), + expected_width, + ) + self._streamed_aux_scratch = None + return False + if scratch.shape[0] < self._max_num_context_tokens: + raise ValueError( + "Kimi-K3 DSpark auxiliary scratch has insufficient token capacity: " + f"capacity={scratch.shape[0]}, required={self._max_num_context_tokens}." + ) + self._streamed_aux_scratch = scratch + return True + + def can_stream_auxiliary_states( + self, + layer_ids: tuple[int, ...], + hidden_states: torch.Tensor, + ) -> bool: + """Return whether a large target forward can use streamed projection.""" + if not self.context_proj_sharded or self._streamed_aux_scratch is None: + return False + if hidden_states.ndim != 2 or hidden_states.shape[0] < _STREAMED_AUX_MIN_TOKENS: + return False + if tuple(layer_ids) != self._streamed_aux_layer_ids: + return False + if hidden_states.shape[1] != self.config.target_hidden_size: + return False + if hidden_states.dtype != self.context_proj.weight.dtype: + return False + if hidden_states.device != self.context_proj.weight.device: + return False + if self._streamed_aux_scratch.dtype != hidden_states.dtype: + return False + if self._streamed_aux_scratch.device != hidden_states.device: + return False + if hidden_states.is_cuda and torch.cuda.is_current_stream_capturing(): + return False + context_kv_weight = getattr(self.context_kv_proj, "weight", None) + return context_kv_weight is not None and context_kv_weight.ndim == 2 + + def begin_auxiliary_stream(self, hidden_states: torch.Tensor) -> None: + """Start one ordered sequence of target auxiliary-state projections.""" + if not self.can_stream_auxiliary_states( + self._streamed_aux_layer_ids, hidden_states + ): + raise RuntimeError( + "Kimi-K3 DSpark auxiliary streaming was started for an " + "unsupported tensor geometry." + ) + self._streamed_aux_tokens = int(hidden_states.shape[0]) + self._streamed_aux_index = 0 + self._streamed_aux_generation += 1 + self._completed_stream_result = None + + def accumulate_auxiliary_state( + self, + primary: torch.Tensor, + residual: torch.Tensor | None, + ) -> None: + """Project one target state into its TP-local draft-hidden rows.""" + index = self._streamed_aux_index + if index >= len(self._streamed_aux_layer_ids): + raise RuntimeError("Kimi-K3 DSpark received too many auxiliary states.") + num_tokens = self._streamed_aux_tokens + if tuple(primary.shape) != (num_tokens, self.config.target_hidden_size): + raise ValueError( + "Kimi-K3 DSpark auxiliary state shape changed during projection: " + f"got={tuple(primary.shape)}, tokens={num_tokens}, " + f"width={self.config.target_hidden_size}." + ) + assert self._streamed_aux_scratch is not None + scratch = self._streamed_aux_scratch[:num_tokens] + if residual is None: + scratch.copy_(primary) + else: + torch.add(primary, residual, out=scratch) + + input_width = int(self.config.target_hidden_size) + weight = self.context_proj.weight[ + :, index * input_width : (index + 1) * input_width + ] + output = self._streamed_context_states[:num_tokens] + if index == 0: + torch.mm(scratch, weight.t(), out=output) + else: + torch.addmm( + output, + scratch, + weight.t(), + beta=1.0, + alpha=1.0, + out=output, + ) + self._streamed_aux_index = index + 1 + + def finish_auxiliary_stream(self) -> torch.Tensor: + """Normalize and return the TP-local projected target context.""" + expected = len(self._streamed_aux_layer_ids) + if self._streamed_aux_index != expected: + raise RuntimeError( + "Kimi-K3 DSpark auxiliary stream ended with an incomplete " + f"projection: received={self._streamed_aux_index}, expected={expected}." + ) + output = self._streamed_context_states[: self._streamed_aux_tokens] + squared_norm = torch.linalg.vector_norm( + output, + ord=2, + dim=-1, + keepdim=True, + dtype=torch.float32, + ).square_() + tensor_model_parallel_all_reduce_in_place(squared_norm) + squared_norm.div_(self.config.hidden_size).add_( + self.context_norm.variance_epsilon + ) + squared_norm.rsqrt_() + output.mul_(squared_norm.to(dtype=output.dtype)) + + output.mul_( + self.context_norm.weight[ + self._context_local_start : self._context_local_start + + self._context_local_width + ] + ) + self._completed_stream_generation = self._streamed_aux_generation + self._completed_stream_result = output + return output + + def is_streamed_context_states(self, states: list[torch.Tensor]) -> bool: + """Claim the completed TP-local stream output exactly once.""" + if not self.context_proj_sharded or len(states) != 1: + return False + candidate = states[0] + expected = self._streamed_context_states[: self._streamed_aux_tokens] + matches = ( + candidate is self._completed_stream_result + and self._completed_stream_generation > self._consumed_stream_generation + and candidate.shape == expected.shape + and candidate.dtype == expected.dtype + and candidate.device == expected.device + and candidate.data_ptr() == expected.data_ptr() + ) + if matches: + self._consumed_stream_generation = self._completed_stream_generation + self._completed_stream_result = None + return matches + @torch.inference_mode() def precompute_and_store_context_kv( self, @@ -251,6 +653,15 @@ def _precompute_fused_context_kv( context_positions: torch.Tensor, context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, ) -> None: + if ( + self.context_proj_sharded + and context_states.shape[-1] == self._context_local_width + ): + self._precompute_streamed_context_kv( + context_states, context_positions, context_slot_mapping + ) + return + num_ctx = context_states.shape[0] num_layers = self._num_context_layers @@ -278,7 +689,8 @@ def _precompute_fused_context_kv( ((num_layers * self._max_num_context_tokens,), torch.int64), ) repeated_positions = repeated_positions[: num_layers * num_ctx] - repeated_positions.view(num_layers, num_ctx).copy_(context_positions) + rope_positions, rope_cos_sin_cache = self._get_rope_inputs(context_positions) + repeated_positions.view(num_layers, num_ctx).copy_(rope_positions) # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. rotary_emb = self.layers[0].self_attn.rotary_emb @@ -288,7 +700,7 @@ def _precompute_fused_context_kv( all_k_pe_flat, None, rotary_emb.head_size, - rotary_emb.cos_sin_cache, + rope_cos_sin_cache, rotary_emb.is_neox_style, ) all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) @@ -310,7 +722,7 @@ def _precompute_fused_context_kv( # Grouped context KV insert assumes that all layers share the same # cache dtype and block layout. - if isinstance(context_slot_mapping, (list, tuple)): + if isinstance(context_slot_mapping, list | tuple): per_layer_slot_mappings = [ s for s in context_slot_mapping if s is not None ] @@ -342,7 +754,7 @@ def _precompute_fused_context_kv( for layer_idx, layer in enumerate(self.layers): slot_mapping = ( context_slot_mapping[layer_idx] - if isinstance(context_slot_mapping, (list, tuple)) + if isinstance(context_slot_mapping, list | tuple) else context_slot_mapping ) if slot_mapping is None: @@ -357,6 +769,77 @@ def _precompute_fused_context_kv( attn._k_scale, ) + def _precompute_streamed_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + """Project TP-local context one draft layer at a time.""" + num_ctx = int(context_states.shape[0]) + full_weight = self.context_kv_proj.weight + if full_weight.ndim != 2: + raise RuntimeError( + "Streamed Kimi-K3 DSpark context KV requires an unquantized " + "two-dimensional fused projection weight." + ) + + rope_positions, rope_cos_sin_cache = self._get_rope_inputs(context_positions) + rotary_emb = self.layers[0].self_attn.rotary_emb + assert rotary_emb is not None + + for layer_idx, layer in enumerate(self.layers): + row_start = layer_idx * self._context_kv_width + weight = full_weight[ + row_start : row_start + self._context_kv_width, + self._context_local_start : self._context_local_start + + self._context_local_width, + ] + layer_kv = F.linear(context_states, weight) + tensor_model_parallel_all_reduce_in_place(layer_kv) + + kv_c = layer_kv[:, : self._context_kv_lora_rank] + normalized_kv_c = torch.empty( + (num_ctx, self._context_kv_lora_rank), + dtype=kv_c.dtype, + device=kv_c.device, + ) + ops.rms_norm( + normalized_kv_c, + kv_c, + self._context_kv_norm_weights[layer_idx], + self._context_rms_norm_eps, + ) + k_pe = layer_kv[:, self._context_kv_lora_rank :].contiguous() + k_pe = k_pe.view(num_ctx, 1, self._context_rope_dim) + ops.rotary_embedding( + rope_positions, + k_pe, + None, + rotary_emb.head_size, + rope_cos_sin_cache, + rotary_emb.is_neox_style, + ) + + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, list | tuple) + else context_slot_mapping + ) + if slot_mapping is None: + del layer_kv, kv_c, normalized_kv_c, k_pe + continue + attn = layer.self_attn + attn.impl.do_kv_cache_update( + normalized_kv_c, + k_pe, + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + del layer_kv, kv_c, normalized_kv_c, k_pe + def _has_uniform_block_layout( self, cache_layers: list[MultiHeadLatentAttention], @@ -398,11 +881,13 @@ def forward( hidden_states = inputs_embeds residual = None + rope_positions, rope_cos_sin_cache = self._get_rope_inputs(positions) for layer in self.layers: hidden_states, residual = layer( - positions=positions, + positions=rope_positions, hidden_states=hidden_states, residual=residual, + rope_cos_sin_cache=rope_cos_sin_cache, ) hidden_states, _ = fused_allreduce_rms_norm( hidden_states, residual, self.final_norm @@ -411,6 +896,7 @@ def forward( class K3DSparkForCausalLM(nn.Module): + model_cls = K3DSparkModel has_own_embed_tokens = False has_own_lm_head = False draft_id_to_target_id = None @@ -437,7 +923,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() - self.model = K3DSparkModel( + self.model = self.model_cls( vllm_config=vllm_config, start_layer_id=target_layer_num, prefix=maybe_prefix(prefix, "model"), @@ -450,6 +936,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) + self._argmax_capacity = min(vllm_config.scheduler_config.max_num_seqs, 8) + self._argmax_enabled = envs.VLLM_KIMI_K3_B12X_DSPARK_ARGMAX + self._argmax_runtime = None + self._argmax_plan = None + self._argmax_output = None + set_b12x_preparation_provider(self, self) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -457,6 +949,32 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.model.combine_hidden_states(hidden_states) + def bind_target_auxiliary_stream( + self, + target_model: nn.Module, + scratch: torch.Tensor, + ) -> None: + """Connect the target's large-prefill states to the draft projector.""" + get_language_model = getattr(target_model, "get_language_model", None) + target_language_model = ( + get_language_model() if callable(get_language_model) else target_model + ) + target_inner = getattr(target_language_model, "model", None) + setter = getattr(target_inner, "set_aux_hidden_state_projector", None) + if not callable(setter): + logger.warning_once( + "Kimi-K3 DSpark auxiliary streaming is disabled because the " + "target model does not expose an auxiliary-state projector hook." + ) + return + if not self.model.bind_auxiliary_stream_scratch(scratch): + return + setter(self.model) + + def is_streamed_context_states(self, states: list[torch.Tensor]) -> bool: + """Return whether target auxiliary states are already TP-local.""" + return self.model.is_streamed_context_states(states) + def get_draft_kv_cache_layer_names(self) -> list[str]: return [layer.self_attn.layer_name for layer in self.model.layers] @@ -485,6 +1003,150 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.compute_logits(hidden_states) + def supports_local_draft_argmax(self) -> bool: + """Return whether rank-local target and Markov logits can be combined.""" + markov_head = self.model.markov_head + if markov_head is None or not markov_head.shard_across_tp: + return False + if not isinstance(self.lm_head, VocabParallelEmbedding): + return False + markov_w2 = markov_head.markov_w2 + if not isinstance(markov_w2, VocabParallelEmbedding): + return False + layout_fields = ( + "tp_size", + "tp_rank", + "num_embeddings", + "org_vocab_size", + "num_embeddings_padded", + "num_embeddings_per_partition", + ) + if any( + getattr(self.lm_head, field) != getattr(markov_w2, field) + for field in layout_fields + ): + return False + if self.lm_head.shard_indices != markov_w2.shard_indices: + return False + # Rank-major base-vocabulary padding preserves B12X's global token + # mapping. Added-vocabulary shards interleave padding and require a + # different local-index mapping. + if self.lm_head.num_embeddings != self.lm_head.org_vocab_size: + return False + if self.logits_processor.soft_cap is not None: + return False + return self.logits_processor.scale == 1.0 + + def _mask_local_draft_padding(self, logits: torch.Tensor) -> None: + """Exclude this rank's padded vocabulary tail from token selection.""" + assert isinstance(self.lm_head, VocabParallelEmbedding) + valid_rows = self.lm_head.shard_indices.num_org_elements + if valid_rows < logits.shape[-1]: + logits[..., valid_rows:].fill_(float("-inf")) + + def compute_local_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project target hidden states into this rank's vocabulary shard.""" + assert isinstance(self.lm_head, VocabParallelEmbedding) + return self.logits_processor._apply_head(self.lm_head, hidden_states, None) + + def compute_local_markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + """Project a Markov embedding into the matching vocabulary shard.""" + return self.model.markov_head.local_bias(markov_embed, self.logits_processor) + + def gather_local_draft_logits( + self, + base_logits: torch.Tensor, + markov_bias: torch.Tensor, + ) -> torch.Tensor: + """Add matching shards and gather the exact complete distribution.""" + logits = tensor_model_parallel_all_gather(base_logits + markov_bias, dim=-1) + return logits[..., : self.logits_processor.org_vocab_size] + + def get_b12x_preparation_units(self, layer, workload: B12xWorkload): + if ( + workload.stage != "weights" + or not self._argmax_enabled + or not self.supports_local_draft_argmax() + or self.lm_head.tp_size not in (8, 12, 16) + ): + return () + from b12x.comm.pcie import _owner_preparation + from b12x.comm.pcie.pcie_vocab_argmax import PCIeVocabParallelArgmax + + if self._argmax_runtime is None: + self._argmax_runtime = PCIeVocabParallelArgmax.from_exchange_group( + exchange_group=get_tp_group().cpu_group, + device=self.lm_head.weight.device, + local_vocab_size=self.lm_head.num_embeddings_per_partition, + max_batch_size=self._argmax_capacity, + ) + self._argmax_output = torch.empty( + self._argmax_capacity, + dtype=torch.int64, + device=self.lm_head.weight.device, + ) + runtime = self._argmax_runtime + if self._argmax_plan is None: + query = _owner_preparation.query_from_runtime(runtime) + self._argmax_plan = _owner_preparation.plan(query, runtime=runtime) + + def prepare(state): + base = torch.full( + (runtime.max_batch_size, runtime.local_vocab_size), + 0.5, + dtype=torch.bfloat16, + device=runtime.device, + ) + bias = torch.full_like(base, 0.25) + return _owner_preparation.prepared_call( + state, base=base, bias=bias, out=self._argmax_output + ) + + name = "kimi_dspark.vocab_argmax" + return ( + B12xPreparationUnit( + name="DSPARK_VOCAB_ARGMAX", + key=( + runtime.world_size, + runtime.local_vocab_size, + self._argmax_capacity, + ), + requests=( + self._argmax_plan.request( + name=name, + prepare_call=prepare, + collective=_owner_preparation.collective(runtime, key=name), + ), + ), + stage="weights", + ), + ) + + def sample_local_draft_logits( + self, base_logits: torch.Tensor, markov_bias: torch.Tensor + ) -> torch.Tensor: + """Sample exact greedy tokens without gathering full vocabulary shards.""" + self._mask_local_draft_padding(base_logits) + self._mask_local_draft_padding(markov_bias) + batch = base_logits.shape[0] + if ( + self._argmax_enabled + and self.lm_head.tp_size in (8, 12, 16) + and base_logits.dtype == markov_bias.dtype == torch.bfloat16 + and 0 < batch <= self._argmax_capacity + ): + if self._argmax_plan is None: + raise PreparationResourceUnavailableError( + "Kimi DSpark vocabulary argmax must be prepared before sampling" + ) + return self._argmax_runtime.fused_add_argmax( + base_logits, + markov_bias, + out=self._argmax_output[:batch], + plan=self._argmax_plan, + ) + return self.gather_local_draft_logits(base_logits, markov_bias).argmax(dim=-1) + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: return draft_ids diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 52025f35a75e..5307fd268a88 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -309,6 +309,8 @@ def get_quantization_config(self): return quant_cfg def is_deepseek_mla(self) -> bool: + if "DFlash2KimiK3Model" in self.get_architectures(): + return getattr(self.hf_text_config, "kv_lora_rank", None) is not None if not hasattr(self.hf_text_config, "model_type"): return False elif self.hf_text_config.model_type in ( @@ -809,6 +811,19 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: return max_position_embeddings, "language_config.max_position_embeddings" +class Qwen3ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__(self, hf_config, hf_text_config, revision=None): + # MLA and GQA DFlash2 checkpoints share the serialized Qwen3 type. + # Resolve the implementation before registry inspection and KV sizing. + if ( + getattr(hf_config, "architectures", None) == ["DFlash2DraftModel"] + and (getattr(hf_config, "dflash_config", None) or {}).get("attention_mode") + == "mla" + ): + hf_config.architectures = ["DFlash2KimiK3Model"] + super().__init__(hf_config, hf_text_config, revision) + + # hf_config.model_type -> convertor class MODEL_ARCH_CONFIG_CONVERTORS = { "bailing_hybrid_mtp": BailingHybridMTPModelArchConfigConvertor, @@ -840,6 +855,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "nemotron-nas": NemotronNasModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, + "qwen3": Qwen3ModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, "qwen4_exp_mtp": Qwen4ExpMTPModelArchConfigConvertor, diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index b7e2a41d00e5..a778d18857f6 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -272,6 +272,9 @@ def update_draft_decode_metadata(self, _metadata: TritonAttentionMetadata) -> No class TritonAttentionBackend(AttentionBackend): + # Replicated draft KV uses complete local sequence metadata, not DCP shards. + supports_dcp_replicated: ClassVar[bool] = True + @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack inline fp32 scales after each head's diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 078033fc55cb..758af5ecc6b8 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1548,6 +1548,7 @@ def _get_kv_cache_groups_uniform_page_size( kv_cache_spec: dict[str, KVCacheSpec], *, log_padding: bool = True, + group_size_override: int | None = None, ) -> list[KVCacheGroupSpec]: """ Generates the KV cache groups for hybrid models with multiple @@ -1608,6 +1609,8 @@ def _get_kv_cache_groups_uniform_page_size( Args: kv_cache_spec: The KVCacheSpec of each attention layer in the model + log_padding: Whether to report padding introduced by grouping. + group_size_override: Maximum layers per group, or the default heuristic. Returns: The generated KVCacheGroupSpecs """ @@ -1637,6 +1640,10 @@ def _get_kv_cache_groups_uniform_page_size( # layers while accommodating speculative decoding drafters that add # extra layers to one attention type. group_size = max_num_layers + if group_size_override is not None: + if group_size_override <= 0: + raise ValueError("KV cache group size override must be positive") + group_size = group_size_override group_counts = [] for layers in layer_buckets: num_padding_layers = group_size - len(layers) % group_size @@ -2503,7 +2510,7 @@ def _warn_if_unannotated_eagle_mamba( kv_cache_groups: Groups as they will be handed to consumers. """ spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): + if spec_config is None or not spec_config.use_eagle_block_drop(): return if any(group.is_eagle_group for group in kv_cache_groups): return @@ -2626,6 +2633,51 @@ def get_kv_cache_groups( # attention free models. return [] + group_size = envs.VLLM_K3_KV_GROUP_SIZE + model_type = getattr( + getattr(getattr(vllm_config, "model_config", None), "hf_config", None), + "model_type", + None, + ) + if ( + group_size > 0 + and model_type == "kimi_k3" + and any(isinstance(spec, MambaSpec) for spec in kv_cache_spec.values()) + and ( + any( + isinstance(spec, SlidingWindowMLASpec) + and spec.non_causal_multi_token_decode + for spec in kv_cache_spec.values() + ) + or ( + getattr(vllm_config.speculative_config, "method", None) == "dflash" + and any( + isinstance(spec, SlidingWindowSpec) and spec.dcp_replicated + for spec in kv_cache_spec.values() + ) + ) + ) + ): + # All groups draw from one pool sized by its widest group. Bound the + # target and draft widths together so a short draft tail does not + # consume blocks sized for all target attention layers. + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + aligned_specs = ( + kv_cache_spec + if layout.is_block_outermost + else unify_kv_cache_spec_page_size(kv_cache_spec) + ) + groups = _get_kv_cache_groups_uniform_page_size( + aligned_specs, group_size_override=group_size + ) + _annotate_eagle_groups(vllm_config, aligned_specs, groups) + logger.info_once( + "Kimi-K3 uses hybrid KV group size %d (%d groups).", + group_size, + len(groups), + ) + return groups + target_specs, draft_specs = _partition_parallel_draft_specs( vllm_config, kv_cache_spec ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 66ed299afbb9..7d1351da6f5a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -961,6 +961,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" cache_dtype_str: str | None = None + non_causal_multi_token_decode: bool = False # DeepseekV4-only: see MLAAttentionSpec.model_version. alignment: int | None = None # Default to None for no padding. model_version: str | None = None @@ -999,6 +1000,7 @@ def merge(cls, specs: list[Self]) -> Self: sliding_window_set = set(spec.sliding_window for spec in specs) extra_retained_set = set(spec.extra_retained_tokens for spec in specs) dcp_replicated_set = {spec.dcp_replicated for spec in specs} + non_causal_set = {spec.non_causal_multi_token_decode for spec in specs} prefix_policy_set = { (spec.prefix_cache_enabled, spec.prefill_replay_window) for spec in specs } @@ -1009,12 +1011,13 @@ def merge(cls, specs: list[Self]) -> Self: and len(sliding_window_set) == 1 and len(extra_retained_set) == 1 and len(dcp_replicated_set) == 1 + and len(non_causal_set) == 1 and len(prefix_policy_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, sliding " "window size, retained token count, DCP replication mode, " - "and prefix policy." + "non-causal mode, and prefix policy." ) return cls( block_size=specs[0].block_size, @@ -1030,6 +1033,7 @@ def merge(cls, specs: list[Self]) -> Self: tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), dcp_replicated=dcp_replicated_set.pop(), + non_causal_multi_token_decode=non_causal_set.pop(), prefix_cache_enabled=specs[0].prefix_cache_enabled, prefill_replay_window=specs[0].prefill_replay_window, ) @@ -1041,6 +1045,7 @@ def is_uniform_with_collection( isinstance(spec, SlidingWindowMLASpec) and spec.sliding_window == self.sliding_window and spec.dcp_replicated == self.dcp_replicated + and spec.non_causal_multi_token_decode == self.non_causal_multi_token_decode and spec.prefix_cacheable == self.prefix_cacheable and spec.prefill_replay_tokens == self.prefill_replay_tokens for spec in kv_cache_specs.values() diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index f955bf8c37c5..b8d40ac19a88 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -14,6 +14,7 @@ ) from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadataBuilder from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.v1.attention.backend import ( @@ -283,13 +284,17 @@ def init_attn_backend( # Phase 3: create metadata builders and determine cudagraph support. attn_backend_workspace: torch.Tensor | None = None + mla_prefill_workspaces: dict[ + tuple[int, int, torch.Size, torch.dtype], torch.Tensor + ] = {} for kv_cache_group_id, groups in enumerate(attn_groups): kernel_block_size = None if kv_cache_group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[kv_cache_group_id] for group in groups: + layer_config = layer_vllm_configs.get(group.layer_names[0], vllm_config) group.create_metadata_builders( - vllm_config=layer_vllm_configs.get(group.layer_names[0], vllm_config), + vllm_config=layer_config, device=device, kernel_block_size=kernel_block_size, # Microbatches build attention metadata concurrently, and some @@ -300,7 +305,21 @@ def init_attn_backend( # The microbatches' builders share the workspace: they all issue # attention on the one compute stream the threads hand off, so the # buffer is written serially, as it already is across steps. - for builder in group.metadata_builders: + for ubatch_id, builder in enumerate(group.metadata_builders): + if isinstance(builder, MLACommonMetadataBuilder): + # Context gathering writes this scratch only during layer + # execution. Sequential cache groups can reuse it, but + # concurrent target/draft lanes and microbatches cannot. + workspace = builder.chunked_prefill_workspace + key = ( + id(layer_config), + ubatch_id, + workspace.shape, + workspace.dtype, + ) + builder.chunked_prefill_workspace = ( + mla_prefill_workspaces.setdefault(key, workspace) + ) if attn_backend_workspace is None: if hasattr(builder, "_get_workspace_buffer"): attn_backend_workspace = builder._get_workspace_buffer() diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index b527c9f21bdf..13f1ed24b6db 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -15,7 +15,10 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): return ExtractHiddenStatesSpeculator(vllm_config, device) elif speculative_config.method == "dflash": - if "DFlash2DraftModel" in speculative_config.draft_model_config.architectures: + if any( + arch in ("DFlash2DraftModel", "DFlash2KimiK3Model") + for arch in speculative_config.draft_model_config.architectures + ): from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( DFlash2Speculator, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 8ae1ec9014b0..0c6568fa49d2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -183,7 +183,11 @@ def load_draft_model( target_model: nn.Module, target_attn_layer_names: set[str], ) -> nn.Module: - return load_dflash_model(target_model, self.vllm_config) + model = load_dflash_model(target_model, self.vllm_config) + bind_auxiliary_stream = getattr(model, "bind_target_auxiliary_stream", None) + if callable(bind_auxiliary_stream): + bind_auxiliary_stream(target_model, self.hidden_states) + return model def set_attn( self, @@ -428,11 +432,21 @@ def propose( # number of rejected tokens, we maintain the size of input_ids and # hidden_states the same as the target model's. This means, we pad each # request's query length to include any rejected positions. + is_streamed_context = getattr(self.model, "is_streamed_context_states", None) + context_states_are_streamed = bool( + aux_hidden_states + and not context_kv_is_restored + and ced_indices is None + and callable(is_streamed_context) + and is_streamed_context(aux_hidden_states) + ) + context_states = self.hidden_states[:num_context_tokens] use_context_graph = ( not dummy_run and ced_indices is None and not is_profile and not context_kv_is_restored + and not context_states_are_streamed and bool(aux_hidden_states) and self._context_preparer is not None and self._context_preparer.can_run(num_target_tokens) @@ -447,22 +461,23 @@ def propose( ] else: last_hidden_states = gather_rows(last_hidden_states, ced_indices) - if aux_hidden_states: + if context_states_are_streamed: + context_states = aux_hidden_states[0] + elif aux_hidden_states: hidden_states = self.model.combine_hidden_states( torch.cat(aux_hidden_states, dim=-1) ) else: hidden_states = last_hidden_states - self.hidden_states[:num_context_tokens].copy_( - hidden_states[:num_context_tokens] - ) + if not context_states_are_streamed: + context_states.copy_(hidden_states[:num_context_tokens]) if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. # Since DFlash needs to build its own attention metadata, we must skip the # preparation in this path and run a minimal forward pass. self.model.precompute_and_store_context_kv( - self.hidden_states[:num_context_tokens], + context_states, self.context_positions[:num_context_tokens], ) # DFlash processes all speculative tokens in one forward pass, @@ -555,7 +570,7 @@ def propose( self._context_preparer.run(aux_hidden_states, num_target_tokens) elif not context_kv_is_restored: self.model.precompute_and_store_context_kv( - self.hidden_states[:num_context_tokens], + context_states, context_positions, context_slots, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index bc861577bb2d..930ac82779c2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -121,6 +121,9 @@ def load_draft_model( target_attn_layer_names: set[str], ) -> torch.nn.Module: model = load_dspark_model(target_model, self.vllm_config) + bind_auxiliary_stream = getattr(model, "bind_target_auxiliary_stream", None) + if callable(bind_auxiliary_stream): + bind_auxiliary_stream(target_model, self.hidden_states) # Reduced draft vocab: probabilistic rejection sampling indexes draft # logits by target id, so precompute the draft->target column map and a # scratch buffer to scatter logits into target vocab before sampling. @@ -139,7 +142,7 @@ def load_draft_model( ) self.use_confidence_head = ( self.enable_adaptive_verification - and model.model.confidence_head is not None + and getattr(model.model, "confidence_head", None) is not None ) if self.use_confidence_head: # The acceptance estimator is not needed when a trained confidence head @@ -201,7 +204,12 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: # Per-(req, position) head hidden, ordered (req, step). sample_hidden = head_hidden[self.sample_indices[:num_sample]] # Draft-vocab logits; sampled ids are remapped to target vocab below. - base_logits = self.model.compute_draft_logits(sample_hidden) + local_head = getattr(self.model, "supports_local_draft_argmax", lambda: False)() + base_logits = ( + self.model.compute_local_draft_logits(sample_hidden) + if local_head + else self.model.compute_draft_logits(sample_hidden) + ) vocab_size = base_logits.shape[-1] base_logits = base_logits.view(num_reqs, n_spec, vocab_size) @@ -218,6 +226,22 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: markov_embed = self.model.markov_embed(prev) if self.use_confidence_head: confidence_markov_embeds.append(markov_embed) + if local_head: + bias = self.model.compute_local_markov_bias(markov_embed) + if self.draft_logits is None and self.acceptance_estimator is None: + draft_sampled_i = self.model.sample_local_draft_logits( + base_logits[:, i], bias + ) + else: + logits_i = self.model.gather_local_draft_logits( + base_logits[:, i], bias + ) + draft_sampled_i = self._sample_logits( + logits_i, idx_map[:, i], sample_pos[:, i], i + ) + self.draft_tokens[:num_reqs, i] = draft_sampled_i + prev = draft_sampled_i + continue bias = self.model.markov_bias(markov_embed) if ( self.draft_logits is None diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 36fd92af48f4..1b7dff03d7d5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext + import torch.nn as nn from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace @@ -108,9 +110,19 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo # config is retained for DSpark's target-layer metadata, so we must override it. draft_vllm_config.quant_config = get_draft_quant_config(vllm_config) - with set_model_tag("dspark_head"): + if draft_model_config.hf_config.model_type == "k3_dspark": + from vllm.models.kimi_k3.nvidia.dspark_mla import ( + protect_k3_compact_rope_sources, + ) + + rope_ownership = protect_k3_compact_rope_sources(target_model) + else: + rope_ownership = nullcontext() + with rope_ownership, set_model_tag("dspark_head"): draft_model = get_model( - vllm_config=draft_vllm_config, model_config=draft_model_config + vllm_config=draft_vllm_config, + model_config=draft_model_config, + load_config=speculative_config.draft_load_config, ) target_language_model = ( diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 6b05709a033c..0f3a36ecb453 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -66,8 +66,7 @@ def get_aligned_state_indices_multi_group_kernel( ) first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) - # load multiple block table for each group - groups = tl.arange(0, BLOCK_GROUPS) + groups = tl.program_id(1) * BLOCK_GROUPS + tl.arange(0, BLOCK_GROUPS) valid_group = groups < NUM_GROUPS group_base_addrs = tl.load( block_table_ptrs_ptr + groups, @@ -1295,8 +1294,21 @@ def compute_aligned_state_indices( return self.aligned_state_indices[:, :0] num_state_slots = self.aligned_state_indices.shape[2] - block_rows = 32 - grid = (triton.cdiv(num_reqs, block_rows),) + block_state_slots = triton.next_power_of_2(num_state_slots) + # Bound each program's gather tile: large group counts otherwise spill + # registers and reserve hundreds of MiB of device-local backing memory. + block_rows = min( + 32, triton.next_power_of_2(num_reqs), max(1, 1024 // block_state_slots) + ) + block_groups = min( + 8, + triton.next_power_of_2(self.num_groups), + max(1, 1024 // (block_rows * block_state_slots)), + ) + grid = ( + triton.cdiv(num_reqs, block_rows), + triton.cdiv(self.num_groups, block_groups), + ) get_aligned_state_indices_multi_group_kernel[grid]( self.block_table_ptrs, seq_lens, @@ -1309,9 +1321,9 @@ def compute_aligned_state_indices( num_reqs, CACHE_BLOCK_SIZE=self.block_size, NUM_GROUPS=self.num_groups, - BLOCK_GROUPS=triton.next_power_of_2(self.num_groups), + BLOCK_GROUPS=block_groups, NUM_STATE_SLOTS=num_state_slots, - BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + BLOCK_STATE_SLOTS=block_state_slots, BLOCK_ROWS=block_rows, num_warps=1, ) From 5ef1847281b8c933026a85c110e0e6d9b93f6dd7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 22 Sep 2026 02:41:57 +0000 Subject: [PATCH 175/365] Reuse identical native payloads and recover incomplete autotune caches Signed-off-by: Martin Vit --- .dockerignore | 7 +++ .../test_flashinfer_autotune_cache.py | 51 ++++++++++++++++++ tools/jovian_wheel_release/Dockerfile | 19 ++++++- tools/jovian_wheel_release/build_bundle.sh | 52 ++++++++++++++++++- .../jovian_wheel_release/build_vllm_wheel.sh | 6 ++- .../verify_native_payload.py | 41 +++++++++++++++ .../warmup/flashinfer_autotune_cache.py | 42 +++++++++++++++ vllm/model_executor/warmup/kernel_warmup.py | 9 ++-- 8 files changed, 218 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/test_flashinfer_autotune_cache.py create mode 100644 tools/jovian_wheel_release/verify_native_payload.py diff --git a/.dockerignore b/.dockerignore index e0284129eefa..9b72c9f5cbbb 100644 --- a/.dockerignore +++ b/.dockerignore @@ -5,8 +5,15 @@ /wheel-export/ dist vllm/*.so +vllm/**/*.so vllm/vllm-rs +# Generated FlashAttention packages come from the build or verified wheel. +vllm/vllm_flash_attn/* +!vllm/vllm_flash_attn/.gitkeep +!vllm/vllm_flash_attn/__init__.py +!vllm/vllm_flash_attn/flash_attn_interface.py + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/tests/model_executor/test_flashinfer_autotune_cache.py b/tests/model_executor/test_flashinfer_autotune_cache.py new file mode 100644 index 000000000000..debcbcc84d5b --- /dev/null +++ b/tests/model_executor/test_flashinfer_autotune_cache.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + read_flashinfer_autotune_cache, + save_flashinfer_autotune_cache, +) + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("contents", [None, b"", b'{"version":', b"\xff"]) +def test_interrupted_autotune_cache_is_a_miss(tmp_path, contents): + path = tmp_path / "autotune_configs.json" + if contents is not None: + path.write_bytes(contents) + assert read_flashinfer_autotune_cache(path) is None + + +def test_autotune_cache_preserves_serialized_bytes(tmp_path): + path = tmp_path / "autotune_configs.json" + contents = b'{"version": 1, "configs": []}\n' + tuner = Mock() + tuner.save_configs.side_effect = lambda output: Path(output).write_bytes(contents) + save_flashinfer_autotune_cache(path, tuner) + assert read_flashinfer_autotune_cache(path) == contents + assert list(tmp_path.iterdir()) == [path] + + +@pytest.mark.parametrize("serialization_failure", [False, True]) +def test_failed_autotune_save_preserves_existing_cache(tmp_path, serialization_failure): + path = tmp_path / "autotune_configs.json" + contents = b'{"version": 1, "configs": []}\n' + path.write_bytes(contents) + + def save(output): + Path(output).write_bytes(b'{"version":') + if serialization_failure: + raise RuntimeError("Interrupted tuner serialization") + + tuner = Mock() + tuner.save_configs.side_effect = save + with pytest.raises(RuntimeError if serialization_failure else ValueError): + save_flashinfer_autotune_cache(path, tuner) + assert read_flashinfer_autotune_cache(path) == contents + assert list(tmp_path.iterdir()) == [path] diff --git a/tools/jovian_wheel_release/Dockerfile b/tools/jovian_wheel_release/Dockerfile index 2a617ba8942a..40b986f3f3b3 100644 --- a/tools/jovian_wheel_release/Dockerfile +++ b/tools/jovian_wheel_release/Dockerfile @@ -7,7 +7,7 @@ ARG BUILDER_IMAGE FROM ${UV_IMAGE} AS uv FROM ${RUST_IMAGE} AS rust -FROM ${BUILDER_IMAGE} AS wheels +FROM ${BUILDER_IMAGE} AS wheel-source # DeepGEMM's DeepJIT exception formatter includes the elfutils DWARF API. RUN --mount=type=cache,id=lil-wheel-apt-packages,target=/var/cache/apt,sharing=locked \ @@ -59,6 +59,7 @@ assert torch._C._GLIBCXX_USE_CXX11_ABI assert importlib.metadata.version("nvidia-cutlass-dsl") == os.environ["CUTLASS_DSL_VERSION"] PY +FROM wheel-source AS wheels RUN --mount=type=cache,id=lil-vllm-native-cu134-torch214-sm120-${DEPENDENCY_RECIPE},target=/src/vllm/build,sharing=locked \ --mount=type=cache,id=lil-vllm-fetchcontent-cu134-${DEPENDENCY_RECIPE},target=/tmp/vllm-fetchcontent,sharing=locked \ --mount=type=cache,id=lil-vllm-generated-cu134-torch214-sm120-${DEPENDENCY_RECIPE},target=/tmp/vllm-generated,sharing=locked \ @@ -73,3 +74,19 @@ RUN --mount=type=cache,id=lil-vllm-native-cu134-torch214-sm120-${DEPENDENCY_RECI FROM scratch AS export COPY --from=wheels /wheelhouse/ /wheels/ + +FROM wheel-source AS precompiled-wheels +RUN --mount=type=bind,from=native-bundle,source=/,target=/native-bundle,ro \ + --mount=type=cache,id=lil-wheel-uv,target=/root/.cache/uv,sharing=locked \ + test "$(find /native-bundle/wheels -name 'vllm-*.whl' | wc -l)" -eq 1 && \ + SOURCE_DATE_EPOCH="${SOURCE_DATE_EPOCH}" \ + BUILD_JOBS="${BUILD_JOBS}" \ + VLLM_USE_PRECOMPILED=1 \ + VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1 \ + VLLM_PRECOMPILED_WHEEL_LOCATION="$(find /native-bundle/wheels -name 'vllm-*.whl')" \ + tools/jovian_wheel_release/build_vllm_wheel.sh && \ + /build-venv/bin/python tools/jovian_wheel_release/verify_native_payload.py \ + /native-bundle/wheels/vllm-*.whl /wheelhouse/vllm-*.whl + +FROM scratch AS export-precompiled +COPY --from=precompiled-wheels /wheelhouse/ /wheels/ diff --git a/tools/jovian_wheel_release/build_bundle.sh b/tools/jovian_wheel_release/build_bundle.sh index 728ddd4b9db9..abc7100d08e6 100755 --- a/tools/jovian_wheel_release/build_bundle.sh +++ b/tools/jovian_wheel_release/build_bundle.sh @@ -16,6 +16,12 @@ lock_value() { source_commit=$(git -C "${repo_root}" rev-parse HEAD) source_tree=$(git -C "${repo_root}" rev-parse 'HEAD^{tree}') +# A worktree's .git file refers outside Docker's source context. A standalone +# snapshot may supply the context, but must contain exactly the reviewed tree. +source_context=${VLLM_BUILD_CONTEXT:-${repo_root}} +test "$(git -C "${source_context}" rev-parse HEAD)" = "${source_commit}" +test "$(git -C "${source_context}" rev-parse 'HEAD^{tree}')" = "${source_tree}" +test -z "$(git -C "${source_context}" status --porcelain)" # FetchContent applies tracked patches inside dependency checkouts. Different # recipes cannot safely share those mutable checkouts or their native objects. dependency_recipe=$(git -C "${repo_root}" rev-parse HEAD:cmake/external_projects) @@ -24,6 +30,42 @@ repository=${GITHUB_REPOSITORY:-local-inference-lab/vllm} release_tag=${VLLM_RELEASE_TAG:-"vllm-jovian-cu134-beta-${source_commit}"} builder=$(lock_value buildx.builder) test -z "$(git -C "${repo_root}" status --porcelain)" +build_target='export' +native_args=() +native_source_commit= +native_wheel_sha256= +if [[ -n ${VLLM_PRECOMPILED_BUNDLE:-} ]]; then + native_bundle=$(realpath "${VLLM_PRECOMPILED_BUNDLE}") + (cd "${native_bundle}" && sha256sum --check SHA256SUMS) + native_manifest=${native_bundle}/manifest.json + jq -e '.schema == "local-inference-vllm-wheel-release/v2"' \ + "${native_manifest}" >/dev/null + native_source_commit=$(jq -er .source.commit "${native_manifest}") + native_wheel_sha256=$(jq -er '.packages[] | select(.name == "vllm") | .sha256' \ + "${native_manifest}") + for field in builder_image python cuda pytorch pytorch_commit cuda_arch_list cutlass_dsl; do + case "$field" in + builder_image) key=builder.image ;; + python) key=python.version ;; + cuda) key=cuda.version ;; + pytorch) key=pytorch.version ;; + pytorch_commit) key=pytorch.commit ;; + cuda_arch_list) key=cuda.arch-list ;; + cutlass_dsl) key=cutlass-dsl.version ;; + esac + test "$(jq -er --arg key "$field" '.runtime[$key]' "${native_manifest}")" \ + = "$(lock_value "$key")" + done + # Python-only releases may reuse binaries only with identical native inputs. + git -C "${repo_root}" diff --exit-code "${native_source_commit}" HEAD -- \ + csrc cmake rust requirements CMakeLists.txt setup.py pyproject.toml \ + vllm/vllm_flash_attn vllm/third_party \ + tools/jovian_wheel_release/runtime.lock \ + tools/jovian_wheel_release/build-requirements.lock \ + tools/jovian_wheel_release/normalize_wheel.py + build_target=export-precompiled + native_args=(--build-context "native-bundle=${native_bundle}") +fi mkdir -p "$(dirname "${output_dir}")" if ! mkdir "${output_dir}"; then @@ -44,9 +86,10 @@ docker buildx build \ --build-arg "DEPENDENCY_RECIPE=${dependency_recipe}" \ --build-arg "CUTLASS_DSL_VERSION=$(lock_value cutlass-dsl.version)" \ --build-arg "VLLM_BUILD_CUTLASS_SCALED_MM_C2X=$(lock_value build.cutlass-scaled-mm-c2x)" \ - --target export \ + --target "${build_target}" \ + "${native_args[@]}" \ --output "type=local,dest=${output_dir}/raw" \ - "${repo_root}" + "${source_context}" cp -a "${output_dir}/raw/wheels/." "${output_dir}/bundle/wheels/" wheel=$(find "${output_dir}/bundle/wheels" -maxdepth 1 -name 'vllm-*.whl' -print -quit) @@ -108,9 +151,14 @@ jq -n \ --arg cuda_arch_list "$(lock_value cuda.arch-list)" \ --arg cutlass_scaled_mm_c2x "$(lock_value build.cutlass-scaled-mm-c2x)" \ --arg cutlass_dsl "$(lock_value cutlass-dsl.version)" \ + --arg native_source_commit "${native_source_commit}" \ + --arg native_wheel_sha256 "${native_wheel_sha256}" \ '{schema: "local-inference-vllm-wheel-release/v2", status: $status, scope: "vLLM native and Python runtime for Qwen3.8 SM120 serving", source: {repository: $repository, commit: $commit, tree: $tree}, + native_reuse: (if $native_source_commit == "" then null else + {source_commit: $native_source_commit, wheel_sha256: $native_wheel_sha256, + contract: "identical native source, dependency recipe and runtime ABI"} end), package_version: $package_version, release_tag: $release_tag, runtime: {builder_image: $builder_image, rust_image: $rust_image, uv_image: $uv_image, python: $python, cuda: $cuda, pytorch: $pytorch, diff --git a/tools/jovian_wheel_release/build_vllm_wheel.sh b/tools/jovian_wheel_release/build_vllm_wheel.sh index 4a49bf3d0e78..ba40e543f291 100755 --- a/tools/jovian_wheel_release/build_vllm_wheel.sh +++ b/tools/jovian_wheel_release/build_vllm_wheel.sh @@ -25,8 +25,9 @@ save_generated_sources() { | xargs -0 -r cp --preserve=timestamps --target-directory="${moe_cache}" } -restore_generated_sources -trap save_generated_sources EXIT +if [[ ${VLLM_USE_PRECOMPILED:-0} != 1 ]]; then + restore_generated_sources + trap save_generated_sources EXIT # CMake records generator hashes in its persistent cache. A clean source tree # lacks generated files, so discard only its configuration when the generated @@ -37,6 +38,7 @@ if ! test -f "${dense_source}/kernel_selector.h" \ rm -f -- "${cmake_cache}" done < <(find "${source_root}/build" -type f -name CMakeCache.txt -print0) fi +fi env -u PYTHONPATH \ VLLM_TARGET_DEVICE=cuda \ diff --git a/tools/jovian_wheel_release/verify_native_payload.py b/tools/jovian_wheel_release/verify_native_payload.py new file mode 100644 index 000000000000..80575e56006d --- /dev/null +++ b/tools/jovian_wheel_release/verify_native_payload.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Reject missing, added or changed native payloads in a Python-only wheel.""" + +import argparse +import hashlib +import zipfile +from pathlib import Path + + +def native_payload(path: Path) -> dict[str, str]: + with zipfile.ZipFile(path) as archive: + return { + name: hashlib.sha256(archive.read(name)).hexdigest() + for name in archive.namelist() + if name.startswith("vllm/") + and (name.endswith(".so") or name == "vllm/vllm-rs") + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reference", type=Path) + parser.add_argument("candidate", type=Path) + args = parser.parse_args() + reference = native_payload(args.reference) + candidate = native_payload(args.candidate) + if not reference or candidate != reference: + changed = sorted( + name + for name in reference.keys() | candidate.keys() + if reference.get(name) != candidate.get(name) + ) + raise RuntimeError(f"Native wheel payload differs: {changed}") + print(f"Verified identical native payloads: {len(reference)} files") + + +if __name__ == "__main__": + main() diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index afc3ef3ac014..ff20c0e072f2 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -3,6 +3,7 @@ """FlashInfer autotune cache helpers.""" import hashlib +import json import os import tempfile from contextlib import suppress @@ -10,11 +11,45 @@ from typing import TYPE_CHECKING import vllm.envs as envs +from vllm.logger import init_logger if TYPE_CHECKING: + from flashinfer.autotuner import AutoTuner + from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +logger = init_logger(__name__) + + +def read_flashinfer_autotune_cache(cache_path: Path) -> bytes | None: + """Treat interrupted JSON writes as cache misses, not model-load failures.""" + try: + contents = cache_path.read_bytes() + except FileNotFoundError: + return None + try: + json.loads(contents) + except (json.JSONDecodeError, UnicodeDecodeError): + logger.warning( + "Ignoring invalid FlashInfer autotune JSON at %s; " + "autotuning will regenerate it.", + cache_path, + ) + return None + return contents + + +def save_flashinfer_autotune_cache(cache_path: Path, tuner: "AutoTuner") -> None: + """Publish only complete serialized tuner state to the shared cache.""" + cache_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=cache_path.parent) as directory: + temporary_path = Path(directory) / cache_path.name + tuner.save_configs(str(temporary_path)) + contents = temporary_path.read_bytes() + json.loads(contents) + write_flashinfer_autotune_cache(cache_path, contents) + def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: config_hash = runner.vllm_config.compute_hash(include_version=False) @@ -87,7 +122,14 @@ def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None: try: with os.fdopen(fd, "wb") as f: f.write(contents) + f.flush() + os.fsync(f.fileno()) os.replace(tmp_path, cache_path) + directory_fd = os.open(cache_path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) except BaseException: with suppress(OSError): os.unlink(tmp_path) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c765e12b7aea..12a0ff3cbaea 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -17,7 +17,9 @@ from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + read_flashinfer_autotune_cache, resolve_flashinfer_autotune_file, + save_flashinfer_autotune_cache, write_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( @@ -405,9 +407,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: # during synchronized autotuning. # Read cached autotune results and broadcast to all ranks. cached_results: bytes | None = None - if is_leader and cache_path.exists(): - with open(cache_path, "rb") as f: - cached_results = f.read() + if is_leader: + cached_results = read_flashinfer_autotune_cache(cache_path) cached_results = world.broadcast_object(cached_results, src=0) if cached_results is not None: write_flashinfer_autotune_cache(cache_path, cached_results) @@ -437,4 +438,4 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: if world.world_size > 1: world.barrier() if is_leader: - tuner.save_configs(str(cache_path)) + save_flashinfer_autotune_cache(cache_path, tuner) From a02f585d3f2418cb9f9e8d70c9229bb179b5f86c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 11:50:51 +0000 Subject: [PATCH 176/365] fix(recoverssm): keep boundary-ending speculative state in its block In align mode, a speculative window whose accepted tokens end exactly on a recurrent block boundary selected the following block-table column for the committed state. That column can still be unallocated: the commit kernel then skipped the whole write, while the worker bookkeeping recorded the state in the new column. The next step read an uninitialized page and skipped the migration, corrupting the recurrent state; prefix-cache hits at the request boundary then reused it. GLM-5.3 MTP uses this path since compact KDA recovery. Select the block that holds the last committed token in both the commit plan and the postprocess bookkeeping, matching the (n - 1) // block convention used elsewhere in align mode. Root cause reported by astra on the Local Inference Lab Discord. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .../kimi_k3/test_recoverssm_boundary.py | 129 ++++++++++++++++++ vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 6 +- vllm/v1/worker/gpu/model_states/recoverssm.py | 4 +- 3 files changed, 137 insertions(+), 2 deletions(-) create mode 100644 tests/models/kimi_k3/test_recoverssm_boundary.py diff --git a/tests/models/kimi_k3/test_recoverssm_boundary.py b/tests/models/kimi_k3/test_recoverssm_boundary.py new file mode 100644 index 000000000000..8023e4424366 --- /dev/null +++ b/tests/models/kimi_k3/test_recoverssm_boundary.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Align-mode RecoverSSM state columns at recurrent block boundaries. + +A speculative window whose accepted tokens end exactly on a recurrent block +boundary must keep its state in the block holding the last committed token. +The following block can still be unallocated, so selecting it would drop the +commit while the bookkeeping recorded the state as advanced. +""" + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia.ops.recoverssm import _prepare_commit_plan_kernel +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.worker.gpu.model_states.recoverssm import ( + _postprocess_recoverssm_align_kernel, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="The RecoverSSM kernels require a CUDA-alike device.", +) + +BLOCK_SIZE = 2048 +SPEC_QUERY_LEN = 4 +BLOCKS = [100, 101, 102, 103, 104, 105] + +# (computed before the step, committed tokens, allocated block columns, +# final column, boundary column or None) +CASES = [ + # 8,188 + 4 ends exactly on the 8,192 boundary; column 4 is unallocated. + (8188, 4, 4, 3, 3), + # The window crosses the boundary: state moves on, checkpoint stays behind. + (8190, 4, 5, 4, 3), + # The window stays inside column 3. + (8180, 4, 4, 3, None), + # The first token after a boundary belongs to the new column. + (8192, 1, 5, 4, None), + # Single-token commit ending on the first boundary. + (2047, 1, 1, 0, 0), +] + + +@pytest.mark.parametrize( + "num_computed,commit_len,allocated,final_col,boundary_col", CASES +) +def test_commit_plan_keeps_boundary_state_in_last_computed_block( + num_computed, commit_len, allocated, final_col, boundary_col +): + device = current_platform.device_type + columns = [ + block if col < allocated else NULL_BLOCK_ID for col, block in enumerate(BLOCKS) + ] + block_table = torch.tensor([columns], dtype=torch.int32, device=device) + num_accepted = torch.tensor([commit_len], dtype=torch.int32, device=device) + state_indices = torch.tensor( + [BLOCKS[(num_computed - 1) // BLOCK_SIZE]], dtype=torch.int32, device=device + ) + query_start_loc = torch.tensor( + [0, SPEC_QUERY_LEN], dtype=torch.int32, device=device + ) + computed = torch.tensor([num_computed], dtype=torch.int32, device=device) + commit_lens, final_indices, boundary_indices, recovery_lens = ( + torch.full((1,), -7, dtype=torch.int32, device=device) for _ in range(4) + ) + + _prepare_commit_plan_kernel[(1,)]( + num_accepted, + None, + state_indices, + query_start_loc, + block_table, + computed, + commit_lens, + final_indices, + boundary_indices, + recovery_lens, + NULL_BLOCK_ID, + BLOCK_SIZE, + block_table.shape[1], + num_accepted.stride(0), + 0, + state_indices.stride(0), + query_start_loc.stride(0), + block_table.stride(0), + block_table.stride(1), + computed.stride(0), + SPEC_QUERY_LEN=SPEC_QUERY_LEN, + num_warps=1, + ) + + assert commit_lens.item() == commit_len + assert final_indices.item() == BLOCKS[final_col] + expected_boundary = NULL_BLOCK_ID if boundary_col is None else BLOCKS[boundary_col] + assert boundary_indices.item() == expected_boundary + if boundary_col is not None: + next_boundary = (num_computed // BLOCK_SIZE + 1) * BLOCK_SIZE + assert recovery_lens.item() == next_boundary - num_computed + + +@pytest.mark.parametrize( + "num_computed,commit_len,allocated,final_col,boundary_col", CASES +) +def test_postprocess_records_the_committed_state_column( + num_computed, commit_len, allocated, final_col, boundary_col +): + device = current_platform.device_type + idx_mapping = torch.tensor([2], dtype=torch.int32, device=device) + num_sampled = torch.tensor([commit_len], dtype=torch.int32, device=device) + computed = torch.tensor([num_computed], dtype=torch.int32, device=device) + state_idx = torch.full((4,), -7, dtype=torch.int32, device=device) + num_accepted = torch.full((4,), -7, dtype=torch.int32, device=device) + + _postprocess_recoverssm_align_kernel[(1,)]( + idx_mapping, + num_sampled, + None, + computed, + state_idx, + num_accepted, + MAMBA_BLOCK_SIZE=BLOCK_SIZE, + BLOCK_TABLE_WIDTH=len(BLOCKS), + ) + + assert state_idx[2].item() == final_col + assert num_accepted[2].item() == 1 + assert state_idx[[0, 1, 3]].tolist() == [-7, -7, -7] diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index 2c0a79279fc3..59b9b3f9a161 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -283,8 +283,12 @@ def _prepare_commit_plan_kernel( tl.int32 ) final_num_computed = num_computed + commit_len + # The committed state belongs to the block holding the last committed + # token. A window ending exactly on a boundary stays in that block; the + # next block may still be unallocated. final_state_col = tl.minimum( - final_num_computed // mamba_block_size, block_table_width - 1 + tl.maximum(final_num_computed - 1, 0) // mamba_block_size, + block_table_width - 1, ) final_state_idx = tl.load( block_table_ptr diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index 50933d8bc65c..3be4242f59cc 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -95,10 +95,12 @@ def _postprocess_recoverssm_align_kernel( return num_sampled = tl.load(num_sampled_ptr + batch_idx) num_computed = tl.load(num_computed_ptr + batch_idx) + # Match the commit plan: the running state lives in the block holding the + # last computed token, so the next step migrates it across a boundary. tl.store( state_idx_ptr + req_state_idx, tl.minimum( - (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE, BLOCK_TABLE_WIDTH - 1, ), ) From 56571486426df0e5d2f44cab5e35548a75624393 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 12:13:47 +0000 Subject: [PATCH 177/365] docs(release): describe the RecoverSSM boundary fix Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-860.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-860.json diff --git a/.lil/changes/vllm-860.json b/.lil/changes/vllm-860.json new file mode 100644 index 000000000000..5af84788e1ed --- /dev/null +++ b/.lil/changes/vllm-860.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-860", + "category": "fix", + "summary": "Stop speculative decoding from corrupting GLM-5.3 recurrent state at cache block boundaries", + "models": ["GLM-5.3-Flash", "Kimi-K3"], + "compatibility": "No user action required.", + "details": [ + "With MTP or DFlash and prefix caching, a speculative step that ended exactly on a recurrent block boundary could skip the state commit, producing off-topic text, loops or multilingual garbage from that point.", + "Later requests that reused such a conversation from the prefix cache no longer start from a damaged state." + ], + "pull_requests": [860], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/860"] +} From 9323ef586c339bd812a058ffc2997184a08a6148 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 12:14:47 +0000 Subject: [PATCH 178/365] docs(release): describe the 2026-09-23 canonical sync for Karmic beta Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-813-withdrawal.json | 13 +++++++++++++ .../vllm-karmic-canonical-sync-20260923.json | 16 ++++++++++++++++ 2 files changed, 29 insertions(+) create mode 100644 .lil/changes/vllm-813-withdrawal.json create mode 100644 .lil/changes/vllm-karmic-canonical-sync-20260923.json diff --git a/.lil/changes/vllm-813-withdrawal.json b/.lil/changes/vllm-813-withdrawal.json new file mode 100644 index 000000000000..d5a14e11628a --- /dev/null +++ b/.lil/changes/vllm-813-withdrawal.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-813-withdrawal", + "category": "internal", + "summary": "DeepSeek V4 routed-MoE tuning again times candidates without the concurrent shared expert", + "models": ["DeepSeek-V4-Flash"], + "compatibility": "No user action required.", + "details": [ + "The current routed MoE kernel selects the same serving grid either way; an isolated comparison measured no C1, C8 or prefill difference." + ], + "pull_requests": [813], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/813#issuecomment-5794090484"] +} diff --git a/.lil/changes/vllm-karmic-canonical-sync-20260923.json b/.lil/changes/vllm-karmic-canonical-sync-20260923.json new file mode 100644 index 000000000000..4996edbcdc1d --- /dev/null +++ b/.lil/changes/vllm-karmic-canonical-sync-20260923.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-karmic-canonical-sync-20260923", + "category": "compatibility", + "summary": "Carry the current canonical Karmic runtime into beta", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "Weights load into ordinary PyTorch allocations; the loader keeps only its scoped checkpoint-copy transport.", + "Block-FP8 linears and DeepSeek attention-output projections are prepared with a capacity plan, so intermediate live sizes reuse prepared kernels while serving.", + "Qwen QSA derives page alignment from the speculative ring and accepts more than four MTP draft tokens.", + "Adaptive speculative verification warms each captured shape before timing and alternates sample rounds on every worker." + ], + "authors": ["lukealonso"], + "requires": ["b12x-master-sync-20260923"] +} From bb44ab501b3bba6a6ff4012cdd58cf087303c270 Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 20 Sep 2026 20:58:29 -0400 Subject: [PATCH 179/365] feat(qwen): support atomic external QSA checkpoints (cherry picked from commit 2ea103e6f081c462d9c07d53401e079579142a42) --- tests/models/qwen4_exp/test_b12x_qsa.py | 4 ++ .../test_partial_prefix_cache_primitives.py | 39 +++++++++++++++++++ tests/v1/kv_connector/unit/test_config.py | 26 +++++++++++++ vllm/config/vllm.py | 32 ++++++++++++++- vllm/models/qwen4_exp/nvidia/b12x_qsa.py | 2 +- 5 files changed, 101 insertions(+), 2 deletions(-) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 39da98b5a105..b099de9c69db 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/tests/models/qwen4_exp/test_b12x_qsa.py @@ -183,6 +183,10 @@ def test_qsa_backend_selects_the_manager_block_without_dense_page_limits( assert Qwen4ExpQSABackend.get_preferred_block_size(70) == 72 +def test_qsa_backend_supports_atomic_kv_connectors() -> None: + assert Qwen4ExpQSABackend.supports_kv_connector() + + def test_qsa_dcp_packs_qkv_into_one_collective() -> None: layer = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention) torch.nn.Module.__init__(layer) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py index 2038f173d7c6..4e37ba25dab4 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py @@ -182,6 +182,45 @@ def test_glm_boundary_adapter_requires_gpu_resident_atomic_state( assert adapter_enabled is not external_cache +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_qwen_b12x_boundary_adapter_supports_dcp(monkeypatch, dcp): + monkeypatch.setattr( + "vllm.platforms.current_platform", + SimpleNamespace(is_cuda=lambda: True), + ) + config = SimpleNamespace( + cache_config=SimpleNamespace( + recurrent_checkpoint_policy="request_boundaries", + enable_prefix_caching=True, + mamba_cache_mode="align", + kv_cache_layout=None, + kv_offloading_size=None, + ), + model_config=SimpleNamespace( + enable_sleep_mode=False, + enable_return_routed_experts=False, + hf_text_config=SimpleNamespace(model_type="qwen3_8_flash_next_text"), + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + data_parallel_size=1, + decode_context_parallel_size=dcp, + prefill_context_parallel_size=1, + ), + speculative_config=SimpleNamespace( + method="mtp", + uses_dynamic_speculative_decoding=lambda: False, + ), + use_v2_model_runner=True, + lora_config=None, + kv_transfer_config=object(), + external_boundary_checkpoint_adapter_available=True, + kernel_config=SimpleNamespace(linear_backend="b12x", moe_backend="b12x"), + ) + + assert VllmConfig.use_request_boundary_checkpoints.fget(config) + + def test_request_boundary_branches_deduplicate_and_survive_sibling_eviction(): pool = BlockPool(12, True, 8) cache = BoundaryCheckpointCache(pool) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index f1a1cb16ce37..b00cca641329 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -104,6 +104,7 @@ def _build_config( kv_connector_extra_config: dict | None = None, enable_sleep_mode: bool = False, enable_cumem_allocator: bool = False, + model_type: str | None = None, ) -> VllmConfig: """Build a model-free config for KV-transfer compatibility checks. @@ -112,6 +113,7 @@ def _build_config( kv_connector_extra_config: Optional connector-specific configuration. enable_sleep_mode: Whether sleep mode is enabled. enable_cumem_allocator: Whether the cuMem allocator is enabled. + model_type: Optional Hugging Face model type. Returns: A VllmConfig suitable for exercising KV-transfer compatibility checks. @@ -135,11 +137,35 @@ def _build_config( cfg.model_config = SimpleNamespace( enable_sleep_mode=enable_sleep_mode, enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode), + hf_text_config=SimpleNamespace(model_type=model_type), ) cfg._verify_kv_transfer_compat() return cfg +@pytest.mark.parametrize("supports_atomic_checkpoints", [False, True]) +def test_qwen_qsa_requires_atomic_checkpoint_connector( + monkeypatch, supports_atomic_checkpoints +): + class _StubConnector: + @classmethod + def supports_request_boundary_checkpoints(cls, config): + del config + return supports_atomic_checkpoints + + monkeypatch.setattr( + KVConnectorFactory, + "get_connector_class", + lambda config: _StubConnector, + ) + + if supports_atomic_checkpoints: + _build_config(kv_connector="stub", model_type="qwen3_8_flash_next_text") + else: + with pytest.raises(ValueError, match="atomic request-boundary"): + _build_config(kv_connector="stub", model_type="qwen3_8_flash_next_text") + + @pytest.mark.parametrize( "kv_connector", ["NixlConnector", "MooncakeConnectorV1", "SomeOOTConnector"] ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a7695fac3e6a..8c61475b917d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -687,7 +687,15 @@ def use_request_boundary_checkpoints(self) -> bool: and parallel.data_parallel_size == 1 and ( parallel.decode_context_parallel_size == 1 - or model.hf_text_config.model_type in ("glm5_next_text", "glm5_next") + or model.hf_text_config.model_type + in ( + "qwen3_8_flash_next_text", + "qwen3_8_flash_next", + "qwen4_exp_text", + "qwen4_exp", + "glm5_next_text", + "glm5_next", + ) ) and parallel.prefill_context_parallel_size == 1 and ( @@ -1233,6 +1241,28 @@ def _verify_kv_transfer_compat(self) -> None: ): return + model_type = getattr( + getattr(self.model_config, "hf_text_config", None), "model_type", None + ) + if model_type in ( + "qwen3_8_flash_next_text", + "qwen3_8_flash_next", + "qwen4_exp_text", + "qwen4_exp", + ): + from vllm.distributed.kv_transfer.kv_connector.factory import ( + KVConnectorFactory, + ) + + connector = KVConnectorFactory.get_connector_class( + self.kv_transfer_config + ) + if not connector.supports_request_boundary_checkpoints(self): + raise ValueError( + "Qwen QSA KV transfer requires an atomic request-boundary " + "checkpoint connector" + ) + # PyTorch's expandable_segments allocator uses CUDA VMM, which can # remap a virtual address range to different physical pages over the # engine's lifetime. KV connectors that pin KV cache memory (e.g. diff --git a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py index cd958edab143..4c9a7e57addd 100644 --- a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py @@ -408,7 +408,7 @@ def is_sparse(cls) -> bool: @classmethod def supports_kv_connector(cls) -> bool: - return False + return True @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: From 1dba9bd9dbc4af12a20a5c78d104c804d865cf29 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 09:38:04 -0400 Subject: [PATCH 180/365] docs(release): describe Qwen QSA checkpoint transfer Assisted-by: OpenAI Codex --- .../vllm-qwen38-qsa-checkpoints-20260923.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json diff --git a/.lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json b/.lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json new file mode 100644 index 000000000000..b9d91d81af0b --- /dev/null +++ b/.lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json @@ -0,0 +1,12 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-qwen38-qsa-checkpoints-20260923", + "category": "feature", + "summary": "Allow Qwen QSA checkpoints to use compatible external KV storage", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Qwen KV transfer requires a connector that saves and restores an atomic request-boundary checkpoint." + ], + "authors": ["yatesdr"] +} From 42c6c141bc503f5018426751c242ba5c5c55976b Mon Sep 17 00:00:00 2001 From: derek Date: Tue, 22 Sep 2026 22:35:46 -0400 Subject: [PATCH 181/365] fix(qsa): preserve stable DCP top-k candidate order (cherry picked from commit 562649d8c9b5cca4cda2e7d39b5e95a3ec39294f) --- .../v1/attention/test_indexer_dcp_localize.py | 48 ++++++++++++++++--- .../attention/dsa/dcp_indexer_cutedsl.py | 44 +++++++++++++++-- 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index f46bd2dd89da..377fb2108710 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -634,14 +634,9 @@ def test_sparse_decode_dcp_persistent_topk_matches_non_dcp(): merged_global_topks = _merge_local_topks_global_with_fake_dcp( local_logits, local_topks, topk, world, interleave ) - # The radix top-K kernel selects a deterministic SET but writes it in - # nondeterministic (atomicAdd) order; the production path is permutation- - # invariant (compaction + softmax), so all ranks must agree on the set, not - # the array order. (The fp64 fallback happens to return sorted order.) ref_topk = merged_global_topks[0] for rank_topk in merged_global_topks[1:]: - for row in range(rank_topk.shape[0]): - assert set(rank_topk[row].tolist()) == set(ref_topk[row].tolist()) + torch.testing.assert_close(rank_topk, ref_topk, rtol=0, atol=0) local_outs = [] local_lses = [] @@ -737,10 +732,51 @@ def test_cutedsl_dcp_candidate_pack_and_select_matches_reference( topk, ) + repeated = torch.empty_like(actual) + stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=repeated) + torch.testing.assert_close(repeated, actual, rtol=0, atol=0) for row in range(rows): assert set(actual[row].cpu().tolist()) == set(expected[row].cpu().tolist()) +@pytest.mark.skipif( + not current_platform.is_cuda() or not has_cutedsl(), + reason="This test requires CUDA and CuteDSL", +) +def test_cutedsl_dcp_select_packs_valid_ids_before_deterministic_padding(): + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( + stable_topk_from_gathered_candidates_cutedsl, + ) + + rows = 8 + topk = 512 + candidates = topk * 4 + valid = 37 + gathered = torch.full( + (rows, candidates, 2), + -1, + dtype=torch.float32, + device="cuda", + ) + gathered[..., 0].fill_(float("-inf")) + scores = torch.arange(valid, dtype=torch.float32, device="cuda") + token_ids = torch.arange(valid, dtype=torch.float32, device="cuda") + gathered[:, :valid, 0] = scores + gathered[:, :valid, 1] = token_ids + + expected_valid = torch.arange( + valid, dtype=torch.int32, device="cuda" + ).expand(rows, -1) + expected_padding = torch.full( + (rows, topk - valid), -1, dtype=torch.int32, device="cuda" + ) + expected = torch.cat((expected_valid, expected_padding), dim=1) + + for _ in range(20): + actual = stable_topk_from_gathered_candidates_cutedsl(gathered, topk) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") def test_sparse_prefill_dcp_metadata_localizes_causal_bounds(): device = torch.device("cuda") diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py index 34ce32856ad0..6e608f771610 100644 --- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py +++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py @@ -290,6 +290,7 @@ def kernel(compile_key: CompileKey) -> Any: @cute.struct class SharedStorage: hist: cute.struct.MemRange[Int32, hist_bins] + selected: cute.struct.MemRange[Int32, compile_key.num_candidates] committed_count: cute.struct.MemRange[Int32, 1] running_count: cute.struct.MemRange[Int32, 1] threshold_bin: cute.struct.MemRange[Int32, 1] @@ -363,7 +364,6 @@ def prefix_matches( @cute.jit def radix_pass( keys: cute.Tensor, - output: cute.Tensor, storage, tid: Int32, step: Int32, @@ -371,6 +371,9 @@ def radix_pass( is_final_pass: bool, ): hist_smem = storage.hist.get_tensor(cute.make_layout((hist_bins,))) + selected_smem = storage.selected.get_tensor( + cute.make_layout((compile_key.num_candidates,)) + ) committed_count_smem = storage.committed_count.data_ptr() running_count_smem = storage.running_count.data_ptr() threshold_bin_smem = storage.threshold_bin.data_ptr() @@ -460,7 +463,8 @@ def radix_pass( scope="cta", ) if dst < Int32(topk): - output[dst] = recast_val(~Uint32(key), Int32) + candidate = tid + Int32(key_idx * tb_size) + selected_smem[candidate] = Int32(1) cute.arch.sync_threads() pass_finished = include_threshold_bin_smem.load() @@ -479,13 +483,21 @@ def device_kernel(input: cute.Tensor, out: cute.Tensor): smem = cutlass.utils.SmemAllocator() storage = smem.allocate(shared_storage, 8) + selected_smem = storage.selected.get_tensor( + cute.make_layout((compile_key.num_candidates,)) + ) committed_count_smem = storage.committed_count.data_ptr() + running_count_smem = storage.running_count.data_ptr() prefix_smem = storage.prefix_s.data_ptr() + warp_totals_smem = storage.warp_totals.get_tensor( + cute.make_layout((1, warps_per_block)) + ) for i in range(tid, topk, tb_size): output_row[i] = Int32(-1) for key_idx in cutlass.range_constexpr(keys_per_thread): col = tid + Int32(key_idx * tb_size) + selected_smem[col] = Int32(0) score = Float32(input_row[col, 0]) token_id = Int32(input_row[col, 1]) keys[key_idx] = stable_key(score, token_id) @@ -500,7 +512,6 @@ def device_kernel(input: cute.Tensor, out: cute.Tensor): while finished == Int32(0) and step < Int32(radix_passes - 1): finished = radix_pass( keys, - output_row, storage, tid, step, @@ -512,7 +523,6 @@ def device_kernel(input: cute.Tensor, out: cute.Tensor): if finished == Int32(0): radix_pass( keys, - output_row, storage, tid, Int32(radix_passes - 1), @@ -520,6 +530,32 @@ def device_kernel(input: cute.Tensor, out: cute.Tensor): True, ) + # Keep selected candidates in input order for stable attention sums. + if tid == Int32(0): + running_count_smem.store(Int32(0)) + cute.arch.sync_threads() + for key_idx in cutlass.range_constexpr(keys_per_thread): + col = tid + Int32(key_idx * tb_size) + is_selected = selected_smem[col] * ( + keys[key_idx] != Uint64(0) + ).to(Int32) + prefix = block_scan_inclusive_i32( + is_selected, + cute.arch.lane_idx(), + cute.arch.warp_idx(), + warp_totals_smem, + warps_per_block, + ) + base = running_count_smem.load() + if is_selected != Int32(0): + output_row[base + prefix - Int32(1)] = recast_val( + ~Uint32(keys[key_idx]), Int32 + ) + cute.arch.sync_threads() + if tid == Int32(tb_size - 1): + running_count_smem.store(base + prefix) + cute.arch.sync_threads() + @cute.jit def host_entrypoint( gathered: cute.Tensor, From c885f2e571c4be385c4a2b548abc70bc64e9cabd Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 09:38:21 -0400 Subject: [PATCH 182/365] docs(release): describe stable QSA DCP top-k selection Assisted-by: OpenAI Codex --- .../vllm-qwen38-qsa-stable-dcp-topk-20260923.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json diff --git a/.lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json b/.lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json new file mode 100644 index 000000000000..85d8c796ddf9 --- /dev/null +++ b/.lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json @@ -0,0 +1,12 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-qwen38-qsa-stable-dcp-topk-20260923", + "category": "fix", + "summary": "Keep Qwen QSA DCP top-k selection in a stable order", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "DCP candidate compaction preserves the selected order instead of depending on atomic arrival order." + ], + "authors": ["yatesdr"] +} From 310f77785e572ca572e34f9771e5f371d956a441 Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 20 Sep 2026 20:58:33 -0400 Subject: [PATCH 183/365] fix(b12x): skip missing multimodal stages during warmup (cherry picked from commit 4cb35f67ed0f10a125df1ca777470a5f950f04f1) --- tests/v1/worker/test_gpu_worker.py | 16 ++++++++++++++++ vllm/model_executor/warmup/b12x_prepare.py | 4 ++++ 2 files changed, 20 insertions(+) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index f8f6d6a5d117..8fbcce51b9d2 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -66,6 +66,22 @@ def get_encoder_budget(self): assert all(module.b12x_eager_only for module in connector.modules()) +def test_mark_b12x_eager_shapes_skips_missing_vision_stage() -> None: + from vllm.model_executor.models.utils import StageMissingLayer + from vllm.model_executor.warmup.b12x_prepare import mark_b12x_eager_shapes + + model = nn.Module() + model.visual = StageMissingLayer("vision_tower") + model.get_mm_lora_token_counts = lambda **kwargs: (1, 1) + model.get_mm_mapping = lambda: SimpleNamespace(connector=("visual.merger",)) + worker = SimpleNamespace( + get_model=lambda: model, + model_runner=SimpleNamespace(mm_registry=object()), + ) + + mark_b12x_eager_shapes(worker) + + def test_b12x_workload_covers_target_and_draft_profile_shapes() -> None: from vllm.model_executor.warmup.b12x_prepare import b12x_workload diff --git a/vllm/model_executor/warmup/b12x_prepare.py b/vllm/model_executor/warmup/b12x_prepare.py index 313589ad970a..6bcf4c99f0cd 100644 --- a/vllm/model_executor/warmup/b12x_prepare.py +++ b/vllm/model_executor/warmup/b12x_prepare.py @@ -286,8 +286,12 @@ def mark_b12x_eager_shapes(worker: Worker) -> None: mm_registry = getattr(worker.model_runner, "mm_registry", None) visual = getattr(model, "visual", None) get_token_counts = getattr(model, "get_mm_lora_token_counts", None) + from vllm.model_executor.models.utils import StageMissingLayer + if mm_registry is None or visual is None or not callable(get_token_counts): return + if isinstance(visual, StageMissingLayer): + return from vllm.multimodal.encoder_budget import MultiModalBudget encoder_budget = MultiModalBudget( From fd9cbde58f87cd2ed26719ccd164163153ad98ac Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 09:38:41 -0400 Subject: [PATCH 184/365] docs(release): describe Qwen missing-vision warmup handling Assisted-by: OpenAI Codex --- .../vllm-qwen38-missing-vision-warmup-20260923.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json diff --git a/.lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json b/.lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json new file mode 100644 index 000000000000..90095281604c --- /dev/null +++ b/.lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json @@ -0,0 +1,12 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-qwen38-missing-vision-warmup-20260923", + "category": "fix", + "summary": "Start Qwen without a loaded vision stage during B12X warmup", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Warmup skips a missing multimodal stage rather than treating its placeholder as a vision module." + ], + "authors": ["yatesdr"] +} From 1fbc36ab798d2777c02c5531a0dbdf6a40e0eef7 Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 20 Sep 2026 23:31:06 -0400 Subject: [PATCH 185/365] fix(scheduler): bound contended prefill compute steps (cherry picked from commit 6fb0752df88585e4f50c5fa0b604d079ce6b9987) --- .../test_prefill_compute_share_scheduler.py | 76 ++++++++++++++++++- tests/v1/core/utils.py | 2 + vllm/config/scheduler.py | 19 +++++ vllm/engine/arg_utils.py | 10 +++ vllm/v1/core/sched/scheduler.py | 31 ++++++++ 5 files changed, 135 insertions(+), 3 deletions(-) diff --git a/tests/v1/core/test_prefill_compute_share_scheduler.py b/tests/v1/core/test_prefill_compute_share_scheduler.py index aea462db33da..e49587574795 100644 --- a/tests/v1/core/test_prefill_compute_share_scheduler.py +++ b/tests/v1/core/test_prefill_compute_share_scheduler.py @@ -99,6 +99,31 @@ def test_prefill_compute_share_cli_accepts_auto(): assert engine_args.prefill_compute_share == "auto" +def test_prefill_compute_token_cap_cli_contract(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + + namespace = parser.parse_args( + [ + "--prefill-compute-share", + "auto", + "--max-num-prefill-tokens-per-step", + "3008", + ] + ) + engine_args = EngineArgs.from_cli_args(namespace) + + assert engine_args.max_num_prefill_tokens_per_step == 3008 + + +def test_prefill_compute_token_cap_requires_compute_share(): + with pytest.raises(ValueError, match="requires prefill_compute_share"): + SchedulerConfig( + max_model_len=128, + is_encoder_decoder=False, + max_num_prefill_tokens_per_step=4, + ) + + def test_prefill_interleave_cli_contract(): parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) @@ -445,6 +470,47 @@ def test_prefill_fcfs_and_capacity_bound_does_not_bypass_share(opt_model_path): assert second in scheduler.waiting +def test_fairness_prefill_token_cap_applies_only_under_contention( + opt_model_path, +): + scheduler = _create_fair_scheduler( + opt_model_path, + max_num_batched_tokens=16, + max_model_len=128, + max_num_prefill_tokens_per_step=4, + ) + _establish_decode(scheduler) + (prefill,) = create_requests( + num_requests=1, + num_tokens=32, + req_ids=["prefill"], + ) + scheduler.add_request(prefill) + + decode_output = scheduler.schedule() + scheduler.record_compute_time("decode", 0.01, contended=True) + _update(scheduler, decode_output) + prefill_output = scheduler.schedule() + + assert prefill_output.compute_service_class == "prefill" + assert prefill_output.num_scheduled_tokens == {prefill.request_id: 4} + + uncontended = _create_fair_scheduler( + opt_model_path, + max_num_batched_tokens=16, + max_model_len=128, + ) + (only_prefill,) = create_requests( + num_requests=1, + num_tokens=32, + req_ids=["only-prefill"], + ) + uncontended.add_request(only_prefill) + + output = uncontended.schedule() + assert output.num_scheduled_tokens == {only_prefill.request_id: 16} + + def test_running_prefills_retain_fcfs_order(opt_model_path): scheduler = _create_fair_scheduler( opt_model_path, @@ -509,11 +575,14 @@ def block_prefill(request, num_new_tokens, *_args): assert output.num_scheduled_tokens == {"decode": 1} -def test_decode_turn_falls_back_when_decode_is_no_longer_runnable(opt_model_path): +def test_decode_turn_falls_back_when_decode_is_no_longer_runnable( + opt_model_path, +): scheduler = _create_fair_scheduler( opt_model_path, max_num_batched_tokens=16, max_model_len=128, + max_num_prefill_tokens_per_step=4, ) decode = _establish_decode(scheduler) # Model the async guard that can make a coarsely eligible decode unable to @@ -531,7 +600,7 @@ def test_decode_turn_falls_back_when_decode_is_no_longer_runnable(opt_model_path output = scheduler.schedule() assert output.compute_service_class == "prefill" - assert output.num_scheduled_tokens == {prefill.request_id: 16} + assert output.num_scheduled_tokens == {prefill.request_id: 4} def test_decode_turn_falls_back_to_running_prefill(opt_model_path): @@ -540,6 +609,7 @@ def test_decode_turn_falls_back_to_running_prefill(opt_model_path): max_num_batched_tokens=16, max_model_len=128, long_prefill_token_threshold=8, + max_num_prefill_tokens_per_step=4, ) decode = _establish_decode(scheduler) (prefill,) = create_requests( @@ -562,7 +632,7 @@ def test_decode_turn_falls_back_to_running_prefill(opt_model_path): fallback_output = scheduler.schedule() assert fallback_output.compute_service_class == "prefill" - assert fallback_output.num_scheduled_tokens == {prefill.request_id: 8} + assert fallback_output.num_scheduled_tokens == {prefill.request_id: 4} def test_full_apc_hit_is_decode_not_prefill(opt_model_path): diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 731d987a0268..a0652f86920d 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -66,6 +66,7 @@ def create_scheduler( long_prefill_token_threshold: int = 0, prefill_compute_share: PrefillComputeShare | None = None, prefill_compute_half_life: PrefillComputeHalfLife | None = None, + max_num_prefill_tokens_per_step: int = 0, max_parallel_prefills: MaxParallelPrefills = 1, prefill_policy: PrefillPolicy = "round-robin", decode_refill_target: DecodeRefillTarget = "auto", @@ -127,6 +128,7 @@ def create_scheduler( long_prefill_token_threshold=long_prefill_token_threshold, prefill_compute_share=prefill_compute_share, prefill_compute_half_life=prefill_compute_half_life, + max_num_prefill_tokens_per_step=max_num_prefill_tokens_per_step, max_parallel_prefills=max_parallel_prefills, prefill_policy=prefill_policy, decode_refill_target=decode_refill_target, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 577a908bb583..ab3cea2cee5e 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -210,6 +210,13 @@ class SchedulerConfig: auto mode. """ + max_num_prefill_tokens_per_step: int = Field(default=0, ge=0) + """Maximum local prefill tokens in a contended compute-share step. + + Zero keeps the normal batched-token budget. Decode-only and uncontended + prefill steps are unchanged. + """ + max_parallel_prefills: MaxParallelPrefills = 1 """Maximum local prefills that may share one model step. @@ -337,6 +344,18 @@ def __post_init__(self, max_model_len: int, is_encoder_decoder: bool) -> None: "prefill_compute_share cannot be combined with " "prefill_schedule_interval greater than one" ) + if ( + self.max_num_prefill_tokens_per_step > 0 + and self.prefill_compute_share is None + ): + raise ValueError( + "max_num_prefill_tokens_per_step requires prefill_compute_share" + ) + if self.max_num_prefill_tokens_per_step > self.max_num_batched_tokens: + raise ValueError( + "max_num_prefill_tokens_per_step cannot exceed " + "max_num_batched_tokens" + ) if ( isinstance(self.max_parallel_prefills, int) and self.max_parallel_prefills > self.max_num_seqs diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 54c19d9ef2f0..eb3af384efa2 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -720,6 +720,9 @@ class EngineArgs: prefill_compute_half_life: PrefillComputeHalfLife | None = ( SchedulerConfig.prefill_compute_half_life ) + max_num_prefill_tokens_per_step: int = ( + SchedulerConfig.max_num_prefill_tokens_per_step + ) max_parallel_prefills: MaxParallelPrefills = SchedulerConfig.max_parallel_prefills prefill_policy: PrefillPolicy = SchedulerConfig.prefill_policy decode_refill_target: DecodeRefillTarget = SchedulerConfig.decode_refill_target @@ -1736,6 +1739,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: scheduler_group.add_argument( "--prefill-compute-half-life", **prefill_compute_half_life_kwargs ) + scheduler_group.add_argument( + "--max-num-prefill-tokens-per-step", + **scheduler_kwargs["max_num_prefill_tokens_per_step"], + ) max_parallel_prefills_kwargs = scheduler_kwargs["max_parallel_prefills"] max_parallel_prefills_kwargs.pop("choices", None) max_parallel_prefills_kwargs["type"] = positive_int_or_auto_type @@ -2583,6 +2590,9 @@ def create_engine_config( prefill_schedule_interval=self.prefill_schedule_interval, prefill_compute_share=self.prefill_compute_share, prefill_compute_half_life=self.prefill_compute_half_life, + max_num_prefill_tokens_per_step=( + self.max_num_prefill_tokens_per_step + ), max_parallel_prefills=self.max_parallel_prefills, prefill_policy=self.prefill_policy, decode_refill_target=self.decode_refill_target, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 785212c3ae6b..6e5ea6d15b33 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -389,6 +389,9 @@ def __init__( if prefill_compute_share is not None else None ) + self.prefill_fairness_max_tokens = ( + self.scheduler_config.max_num_prefill_tokens_per_step or None + ) self._decode_compute_seconds = 0.0 self._prefill_compute_seconds = 0.0 self.prefill_interleave_controller = PrefillInterleaveController() @@ -905,6 +908,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: compute_contention = has_eligible_decode and has_prefill_candidate compute_contention_started = compute_contention and not prior_contention + if ( + selected_compute_class == "prefill" + and compute_contention + and self.prefill_fairness_max_tokens is not None + ): + token_budget = min(token_budget, self.prefill_fairness_max_tokens) + input_budget = min( + input_budget, + self.prefill_fairness_max_tokens + draft_slots, + ) + legacy_defer_prefills = ( throttle_prefills and not self.prefill_capacity_bound ) and has_eligible_decode @@ -1510,6 +1524,14 @@ def schedule_running_requests( # its resource checks. Fall back immediately. adaptive_defer_prefills = False defer_prefills = legacy_defer_prefills + if self.prefill_fairness_max_tokens is not None: + token_budget = min( + token_budget, self.prefill_fairness_max_tokens + ) + input_budget = min( + input_budget, + self.prefill_fairness_max_tokens + draft_slots, + ) else: # DP prefill balancing: defer this step's local # prefill compute to a cadence-aligned step. @@ -1872,6 +1894,12 @@ def schedule_running_requests( ): adaptive_defer_prefills = False defer_prefills = False + if self.prefill_fairness_max_tokens is not None: + token_budget = min(token_budget, self.prefill_fairness_max_tokens) + input_budget = min( + input_budget, + self.prefill_fairness_max_tokens + draft_slots, + ) schedule_running_requests("prefill") # A prefill turn gives prefills first use of the model-step capacity. @@ -2150,6 +2178,9 @@ def get_prefill_fairness(self) -> dict[str, Any]: "prefill_compute_half_life": ( self.scheduler_config.prefill_compute_half_life ), + "max_num_prefill_tokens_per_step": ( + self.scheduler_config.max_num_prefill_tokens_per_step + ), "effective_prefill_compute_half_life_seconds": ( controller.effective_prefill_compute_half_life if controller is not None From ee5d2046c4e40782dc996e377c2e35d051b12905 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 09:39:02 -0400 Subject: [PATCH 186/365] docs(release): describe optional contended prefill quantum Assisted-by: OpenAI Codex --- .../vllm-contended-prefill-quantum-20260923.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .lil/changes/vllm-contended-prefill-quantum-20260923.json diff --git a/.lil/changes/vllm-contended-prefill-quantum-20260923.json b/.lil/changes/vllm-contended-prefill-quantum-20260923.json new file mode 100644 index 000000000000..8e8284dd41be --- /dev/null +++ b/.lil/changes/vllm-contended-prefill-quantum-20260923.json @@ -0,0 +1,12 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-contended-prefill-quantum-20260923", + "category": "performance", + "summary": "Optionally bound prefill work per compute-share step under decode load", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "A new step-token limit applies only when prefill and decode contend; the default leaves scheduling unchanged." + ], + "authors": ["yatesdr"] +} From 17eef739a3b806673c62457af98b3a4dd30bc547 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 10:25:04 -0400 Subject: [PATCH 187/365] docs(release): identify QSA ordering change as PR 861 Assisted-by: OpenAI Codex --- ...qwen38-qsa-stable-dcp-topk-20260923.json => vllm-861.json} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename .lil/changes/{vllm-qwen38-qsa-stable-dcp-topk-20260923.json => vllm-861.json} (82%) diff --git a/.lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json b/.lil/changes/vllm-861.json similarity index 82% rename from .lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json rename to .lil/changes/vllm-861.json index 85d8c796ddf9..ee6ab6aa6a00 100644 --- a/.lil/changes/vllm-qwen38-qsa-stable-dcp-topk-20260923.json +++ b/.lil/changes/vllm-861.json @@ -1,6 +1,6 @@ { "schema": "local-inference-release-change/v1", - "id": "vllm-qwen38-qsa-stable-dcp-topk-20260923", + "id": "vllm-861", "category": "fix", "summary": "Keep Qwen QSA DCP top-k selection in a stable order", "models": ["Qwen3.8-Flash-Next"], @@ -8,5 +8,5 @@ "details": [ "DCP candidate compaction preserves the selected order instead of depending on atomic arrival order." ], - "authors": ["yatesdr"] + "pull_requests": [861] } From f7f4ae77aeaaf6108afb6f181c218387859aa8cf Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 10:25:30 -0400 Subject: [PATCH 188/365] docs(release): identify Qwen checkpoint change as PR 864 Assisted-by: OpenAI Codex --- ...llm-qwen38-qsa-checkpoints-20260923.json => vllm-864.json} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename .lil/changes/{vllm-qwen38-qsa-checkpoints-20260923.json => vllm-864.json} (83%) diff --git a/.lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json b/.lil/changes/vllm-864.json similarity index 83% rename from .lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json rename to .lil/changes/vllm-864.json index b9d91d81af0b..15e7d9665ca6 100644 --- a/.lil/changes/vllm-qwen38-qsa-checkpoints-20260923.json +++ b/.lil/changes/vllm-864.json @@ -1,6 +1,6 @@ { "schema": "local-inference-release-change/v1", - "id": "vllm-qwen38-qsa-checkpoints-20260923", + "id": "vllm-864", "category": "feature", "summary": "Allow Qwen QSA checkpoints to use compatible external KV storage", "models": ["Qwen3.8-Flash-Next"], @@ -8,5 +8,5 @@ "details": [ "Qwen KV transfer requires a connector that saves and restores an atomic request-boundary checkpoint." ], - "authors": ["yatesdr"] + "pull_requests": [864] } From f56eb68c5fc3465f92b1e78bdff30082a9e01c56 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 10:25:53 -0400 Subject: [PATCH 189/365] docs(release): identify multimodal warmup fix as PR 863 Assisted-by: OpenAI Codex --- ...en38-missing-vision-warmup-20260923.json => vllm-863.json} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename .lil/changes/{vllm-qwen38-missing-vision-warmup-20260923.json => vllm-863.json} (81%) diff --git a/.lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json b/.lil/changes/vllm-863.json similarity index 81% rename from .lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json rename to .lil/changes/vllm-863.json index 90095281604c..6a43558b613c 100644 --- a/.lil/changes/vllm-qwen38-missing-vision-warmup-20260923.json +++ b/.lil/changes/vllm-863.json @@ -1,6 +1,6 @@ { "schema": "local-inference-release-change/v1", - "id": "vllm-qwen38-missing-vision-warmup-20260923", + "id": "vllm-863", "category": "fix", "summary": "Start Qwen without a loaded vision stage during B12X warmup", "models": ["Qwen3.8-Flash-Next"], @@ -8,5 +8,5 @@ "details": [ "Warmup skips a missing multimodal stage rather than treating its placeholder as a vision module." ], - "authors": ["yatesdr"] + "pull_requests": [863] } From cb49c008cb5b860186d1e349b5509f95604a1385 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 23 Sep 2026 10:26:16 -0400 Subject: [PATCH 190/365] docs(release): identify prefill quantum change as PR 862 Assisted-by: OpenAI Codex --- ...-contended-prefill-quantum-20260923.json => vllm-862.json} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename .lil/changes/{vllm-contended-prefill-quantum-20260923.json => vllm-862.json} (83%) diff --git a/.lil/changes/vllm-contended-prefill-quantum-20260923.json b/.lil/changes/vllm-862.json similarity index 83% rename from .lil/changes/vllm-contended-prefill-quantum-20260923.json rename to .lil/changes/vllm-862.json index 8e8284dd41be..9e1ae92cc27a 100644 --- a/.lil/changes/vllm-contended-prefill-quantum-20260923.json +++ b/.lil/changes/vllm-862.json @@ -1,6 +1,6 @@ { "schema": "local-inference-release-change/v1", - "id": "vllm-contended-prefill-quantum-20260923", + "id": "vllm-862", "category": "performance", "summary": "Optionally bound prefill work per compute-share step under decode load", "models": ["all"], @@ -8,5 +8,5 @@ "details": [ "A new step-token limit applies only when prefill and decode contend; the default leaves scheduling unchanged." ], - "authors": ["yatesdr"] + "pull_requests": [862] } From 9d80bb39b8c7643abba907c0b5beaab4c6a14656 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 14:31:46 +0000 Subject: [PATCH 191/365] fix(qwen): size QSA capture metadata for every replay Breakable PIECEWISE graphs record B12X QSA context selection and runtime staging. PIECEWISE capture builds metadata with for_capture=False, so QSA chose its context from the dummy batch's few tokens. A graph with more rows than max_num_seqs * (1 + draft tokens) therefore baked the smallest prefill context and its three-page compressed block table; replaying it with a long request read past the table, causing an illegal memory access or silent KV corruption (b12x#416). Mark every capture batch explicitly and pass max_model_len to the QSA metadata builder during capture, as FULL graphs already do. Only QSA metadata changes; other attention groups keep the dummy lengths that their eager capture paths require. Reported with root-cause analysis by @randomvariable. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- tests/models/qwen4_exp/test_b12x_qsa.py | 41 ++++++++++++++++++++ tests/v1/cudagraph/test_cudagraph_manager.py | 36 +++++++++++++++++ vllm/models/qwen4_exp/nvidia/b12x_qsa.py | 3 +- vllm/models/qwen4_exp/nvidia/model_state.py | 9 +++++ vllm/v1/worker/gpu/cudagraph_utils.py | 1 + vllm/v1/worker/gpu/input_batch.py | 5 +++ 6 files changed, 94 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 39da98b5a105..c3ea49f9abcb 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/tests/models/qwen4_exp/test_b12x_qsa.py @@ -789,6 +789,47 @@ def bind_prepare_and_exercise(cache, backing, *, previous_plans): assert owner.get_b12x_preparation_units(owner, workload) == () +def test_qsa_capture_metadata_covers_max_model_len() -> None: + """A PIECEWISE capture must not size QSA metadata for dummy lengths. + + Breakable PIECEWISE graphs record B12X QSA context selection and staging. + A context chosen for the dummy batch's few tokens bakes a too-narrow block + table, and replays at long positions read past it (b12x#416). + """ + builder = Qwen4ExpQSAMetadataBuilder.__new__(Qwen4ExpQSAMetadataBuilder) + builder._request_ids = torch.empty(8, dtype=torch.int32) + builder.max_speculative_tokens = 3 + builder._capture_state_slot_ids = torch.arange(2, dtype=torch.int32) + builder._capture_state_is_fresh = torch.ones(2, dtype=torch.bool) + builder._capture_num_accepted_tokens = torch.ones(2, dtype=torch.int32) + builder._capture_is_prefilling = torch.zeros(2, dtype=torch.bool) + common = CommonAttentionMetadata( + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), + seq_lens=torch.tensor([7, 7], dtype=torch.int32), + num_reqs=2, + num_actual_tokens=8, + max_query_len=4, + max_seq_len=7, + block_table_tensor=torch.zeros((2, 1), dtype=torch.int32), + slot_mapping=torch.full((8,), -1, dtype=torch.int64), + is_prefilling=torch.zeros(2, dtype=torch.bool), + ) + common._token_to_req_indices_cache = torch.tensor( + [0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32 + ) + + eager = builder.build(common_prefix_len=0, common_attn_metadata=common) + captured = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + qsa_max_seq_len=1_048_576, + ) + + assert eager.max_seq_len == 7 + assert captured.max_seq_len == 1_048_576 + + def test_qsa_registers_piecewise_splitting_op_once() -> None: compilation_config = SimpleNamespace( static_forward_context={}, diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index ca4f1fe33d1c..16764d043bca 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -236,6 +236,42 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings(): block_tables.get_dummy_slot_mappings.assert_not_called() +@pytest.mark.parametrize("full_cudagraph", [False, True]) +def test_capture_marks_dummy_batch_for_every_graph_mode(full_cudagraph): + """Operations recorded in PIECEWISE graphs can size metadata for replay.""" + num_reqs = 8 + num_tokens = 56 + input_buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu")) + block_tables = MagicMock() + block_tables.cp_size = 1 + block_tables.get_dummy_block_tables.return_value = () + block_tables.get_dummy_slot_mappings.return_value = torch.zeros( + 1, num_tokens, dtype=torch.int64 + ) + model_state = MagicMock() + model_state.prepare_attn.return_value = {} + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[], + ) + + gpu_cudagraph_utils.prepare_inputs_to_capture( + num_reqs, + num_tokens, + model_state, + input_buffers, + block_tables, + [], + kv_cache_config, + full_cudagraph=full_cudagraph, + ) + + input_batch = model_state.prepare_attn.call_args.args[0] + assert input_batch.cudagraph_capture is True + assert model_state.prepare_attn.call_args.kwargs["for_capture"] is full_cudagraph + + def test_capture_synchronizes_auxiliary_warmup_streams(monkeypatch): """CUDA graph capture starts only after warmup work has completed.""" lifecycle: list[str] = [] diff --git a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py index cd958edab143..76ae1dabd85c 100644 --- a/vllm/models/qwen4_exp/nvidia/b12x_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/b12x_qsa.py @@ -248,6 +248,7 @@ def build( qsa_state_is_fresh: torch.Tensor | None = None, qsa_num_accepted_tokens: torch.Tensor | None = None, qsa_is_prefilling: torch.Tensor | None = None, + qsa_max_seq_len: int | None = None, ) -> Qwen4ExpQSAMetadata: del common_prefix_len, fast_build cm = common_attn_metadata @@ -281,7 +282,7 @@ def build( num_actual_tokens=cm.num_actual_tokens, max_query_len=cm.max_query_len, query_start_loc=cm.query_start_loc, - max_seq_len=cm.max_seq_len, + max_seq_len=cm.max_seq_len if qsa_max_seq_len is None else qsa_max_seq_len, seq_lens=cm.seq_lens, block_table=cm.block_table_tensor, slot_mapping=cm.slot_mapping, diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index 82c9726f004e..183468978fea 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -37,6 +37,9 @@ class Qwen4ExpAttnMetadata(MambaHybridAttnMetadata): qsa_state_is_fresh: torch.Tensor | None = None qsa_num_accepted_tokens: torch.Tensor | None = None qsa_is_prefilling: torch.Tensor | None = None + # Set during CUDA graph capture: B12X QSA selects and stages its context + # inside PIECEWISE graphs, so the captured context must cover max_model_len. + qsa_max_seq_len: int | None = None def get_extra_attn_kwargs( self, @@ -55,6 +58,7 @@ def get_extra_attn_kwargs( qsa_state_is_fresh=self.qsa_state_is_fresh[:num_reqs], qsa_num_accepted_tokens=self.qsa_num_accepted_tokens[:num_reqs], qsa_is_prefilling=self.qsa_is_prefilling[:num_reqs], + qsa_max_seq_len=self.qsa_max_seq_len, ) return kwargs @@ -374,6 +378,11 @@ def prepare_attn( qsa_state_is_fresh=qsa_state_is_fresh, qsa_num_accepted_tokens=qsa_num_accepted_tokens, qsa_is_prefilling=qsa_is_prefilling, + qsa_max_seq_len=( + self.max_model_len + if for_capture or input_batch.cudagraph_capture + else None + ), ) attn_metadata = build_attn_metadata( attn_groups=attn_groups, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 35878124f529..ea5d19cddf19 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1087,6 +1087,7 @@ def prepare_inputs_to_capture( input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) input_batch.uniform_decode_graph = full_cudagraph and uniform_decode_graph + input_batch.cudagraph_capture = True block_table_provider = pcp_manager or block_tables input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 78e9f3095d69..81b24b8b5e9b 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -127,6 +127,11 @@ class InputBatch: # from dummy query lengths: a mixed graph may be captured with uniform rows. uniform_decode_graph: bool = False + # Dummy batch for any CUDA graph capture, FULL or PIECEWISE. Operations + # recorded inside a PIECEWISE graph must size their metadata for every + # replay, not for the dummy sequence lengths. + cudagraph_capture: bool = False + @classmethod def make_dummy( cls, From 26b42cf9eb715b118e55f91f53fa0ab412a79c09 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 14:49:57 +0000 Subject: [PATCH 192/365] chore(release): describe the Qwen QSA capture fix Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-865.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-865.json diff --git a/.lil/changes/vllm-865.json b/.lil/changes/vllm-865.json new file mode 100644 index 000000000000..952a5a498589 --- /dev/null +++ b/.lil/changes/vllm-865.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-865", + "category": "fix", + "summary": "Stop Qwen QSA from corrupting the KV cache when CUDA graphs are larger than the decode-row limit", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required. The shipped profile was not affected; deployments that lowered max-num-seqs without lowering the CUDA graph capture size were.", + "details": [ + "Mixed batches replayed in graphs with more rows than max-num-seqs times (1 + draft tokens) could write another request's compressed KV page, causing wrong output without an error, or an illegal memory access.", + "QSA now sizes captured graphs for the full model length." + ], + "pull_requests": [865], + "evidence": ["https://github.com/local-inference-lab/vllm/pull/865"] +} From 3f5e43720861012ff8afbe9d37527cc0bd4c04c4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 15:47:10 +0000 Subject: [PATCH 193/365] docs(release): list every DCP model that uses the shared top-k merge The merge kernel serves the Qwen QSA, GLM and DeepSeek sparse indexers, so DCP deployments of all four models see the new stable order. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-861.json | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.lil/changes/vllm-861.json b/.lil/changes/vllm-861.json index ee6ab6aa6a00..ed79fb201473 100644 --- a/.lil/changes/vllm-861.json +++ b/.lil/changes/vllm-861.json @@ -2,11 +2,12 @@ "schema": "local-inference-release-change/v1", "id": "vllm-861", "category": "fix", - "summary": "Keep Qwen QSA DCP top-k selection in a stable order", - "models": ["Qwen3.8-Flash-Next"], - "compatibility": "No user action required.", + "summary": "Keep DCP top-k attention selection in a stable order", + "models": ["Qwen3.8-Flash-Next", "GLM-5.3-Flash", "DeepSeek-V4-Flash", "DeepSeek-V4.1-Flash"], + "compatibility": "No user action required. Only decode context parallelism (DCP > 1) is affected.", "details": [ - "DCP candidate compaction preserves the selected order instead of depending on atomic arrival order." + "DCP candidate compaction preserves the selected order instead of depending on atomic arrival order.", + "The merge kernel is shared by the Qwen QSA, GLM and DeepSeek sparse indexers." ], "pull_requests": [861] } From e08582200d703c8123662c32a76ff84edfb7df5e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 19:02:38 +0000 Subject: [PATCH 194/365] feat(engine): report where a request or the engine loop is held The periodic stats line cannot tell a blocked engine core from an idle one, and a request that waits before reaching the engine leaves no trace. Name the stage and its duration instead: - The API server records when each request passed rendering, input processing, engine submission, first and final output, and the end of its stream. It warns when a request spends longer than VLLM_REQUEST_STALL_WARNING_S (default 60) outside engine work, or ends without any output after that long. VLLM_LOG_REQUEST_TIMELINE=1 logs the line for every request. - The engine core warns when an ADD request waited that long in its input queue, and a watchdog logs the loop thread's stack when one loop step or client request runs that long, then the step's total time. - SIGUSR1 to the API server logs every asyncio task's await chain, with local names and types but no text values, and every thread's stack. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .../entrypoints/serve/utils/test_task_dump.py | 58 ++++ tests/v1/engine/test_stall_diagnostics.py | 191 +++++++++++++ vllm/entrypoints/launchers/launcher.py | 3 + vllm/entrypoints/serve/utils/task_dump.py | 115 ++++++++ vllm/envs.py | 18 ++ vllm/v1/engine/async_llm.py | 24 ++ vllm/v1/engine/core.py | 40 +++ vllm/v1/engine/output_processor.py | 3 + vllm/v1/engine/stall_diagnostics.py | 251 ++++++++++++++++++ 9 files changed, 703 insertions(+) create mode 100644 tests/entrypoints/serve/utils/test_task_dump.py create mode 100644 tests/v1/engine/test_stall_diagnostics.py create mode 100644 vllm/entrypoints/serve/utils/task_dump.py create mode 100644 vllm/v1/engine/stall_diagnostics.py diff --git a/tests/entrypoints/serve/utils/test_task_dump.py b/tests/entrypoints/serve/utils/test_task_dump.py new file mode 100644 index 000000000000..4b2468b891f1 --- /dev/null +++ b/tests/entrypoints/serve/utils/test_task_dump.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import threading + +from vllm.entrypoints.serve.utils.task_dump import format_asyncio_tasks + + +def test_dump_follows_a_request_through_async_generators(): + """The dump reaches the awaited event through nested async generators and + shows local names and types, but not the values of text locals.""" + + async def main() -> str: + released = asyncio.Event() + + async def engine_outputs(request_id: str): + prompt = "private prompt text" + num_tokens = 228652 + await released.wait() + yield prompt, num_tokens + + async def stream_chunks(request_id: str): + async for output in engine_outputs(request_id): + yield output + + async def respond(): + async for _ in stream_chunks("chatcmpl-held"): + pass + + task = asyncio.create_task(respond(), name="held-response") + await asyncio.sleep(0) + dump = format_asyncio_tasks() + released.set() + await task + return dump + + stop = threading.Event() + helper = threading.Thread(target=stop.wait, name="tokenizer-worker") + helper.start() + try: + dump = asyncio.run(main()) + finally: + stop.set() + helper.join() + + section = dump.split("Task 'held-response':")[1].split(" Task ")[0] + order = [ + section.index(name) + for name in ("respond", "stream_chunks", "engine_outputs", "Event.wait") + ] + assert order == sorted(order) + assert "request_id='chatcmpl-held'" in section + assert "num_tokens=228652" in section + assert "prompt: str" in section + assert "private prompt text" not in dump + assert "task is blocked on RequestOutput: + stats = RequestStateStats(num_generation_tokens=7, queued_ts=10.0) + stats.scheduled_ts = 10.25 + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=list(range(100)), + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[1], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop" if finished else None, + ) + ], + finished=finished, + metrics=stats, + num_cached_tokens=96, + ) + + +def _timeline(now: float, **stages: float) -> RequestTimeline: + timeline = RequestTimeline("req-1", now - stages.pop("arrival_age"), now - 5.0) + for stage, age in stages.items(): + setattr(timeline, stage, now - age) + return timeline + + +def test_hold_before_submission_is_reported(caplog_vllm): + now = time.time() + # Rendered 125 s ago, but the generator only started 5 s ago. + timeline = _timeline( + now, + arrival_age=125.0, + processed=4.9, + submitted=4.8, + first_output=1.0, + final_output=0.1, + ) + timeline.output(_output("req-1", finished=True)) + timeline.final_output = now - 0.1 + timeline.close("finished", threshold=60.0, always=False) + + record = next(r for r in caplog_vllm.records if "req-1" in r.getMessage()) + message = record.getMessage() + assert record.levelname == "WARNING" + assert "Request req-1 finished after 125." in message + assert "arrival to generator start 120.00 s" in message + assert "engine queue 0.25 s" in message + assert "prompt 100 tokens (96 cached)" in message + assert "7 generated" in message + assert "finish reason stop" in message + + +def test_fast_request_is_not_reported(caplog_vllm): + now = time.time() + timeline = _timeline( + now, + arrival_age=5.5, + processed=4.9, + submitted=4.8, + first_output=1.0, + final_output=0.1, + ) + timeline.close("finished", threshold=60.0, always=False) + assert "req-1" not in caplog_vllm.text + + +def test_request_aborted_before_any_output_is_reported(caplog_vllm): + now = time.time() + timeline = _timeline(now, arrival_age=90.0, processed=4.9, submitted=4.8) + timeline.started = now - 89.9 + timeline.close("aborted", threshold=60.0, always=False) + assert "Request req-1 aborted after 90." in caplog_vllm.text + assert "submission to first output -" in caplog_vllm.text + + +def test_handler_stopping_after_final_output_counts_as_finished(caplog_vllm): + now = time.time() + timeline = _timeline(now, arrival_age=5.5, processed=4.9, submitted=4.8) + timeline.output(_output("req-1", finished=True)) + timeline.close("aborted", threshold=60.0, always=True) + assert "Request req-1 finished after" in caplog_vllm.text + + +def test_generate_records_each_stage(caplog_vllm): + async def add_request(request_id, prompt, params, **kwargs): + collector = RequestOutputCollector(RequestOutputKind.DELTA, request_id) + collector.put(_output(request_id, finished=True)) + return collector + + engine = SimpleNamespace( + add_request=add_request, + log_requests=False, + _request_stall_warning_s=60.0, + _log_request_timeline=True, + ) + prompt = {"type": "token", "prompt_token_ids": [1, 2], "arrival_time": 0.0} + prompt["arrival_time"] = time.time() - 2.0 + + async def consume(): + return [ + out + async for out in AsyncLLM.generate( + engine, prompt, SamplingParams(), "req-1" + ) + ] + + outputs = asyncio.run(consume()) + assert len(outputs) == 1 and outputs[0].finished + message = next( + r.getMessage() for r in caplog_vllm.records if "req-1" in r.getMessage() + ) + assert message.startswith("Request req-1 finished after 2.") + assert "arrival to generator start 2." in message + for stage in ( + "input processing", + "engine submission", + "submission to first output", + "first to final output", + "final output to stream end", + ): + assert f"{stage} -" not in message + + +def test_engine_loop_watchdog_reports_a_blocked_step(caplog_vllm): + watchdog = EngineLoopWatchdog(threshold=0.2, poll_interval=0.02) + + def step_waiting_on_a_transfer(): + time.sleep(0.6) + + try: + watchdog.busy("engine step") + step_waiting_on_a_transfer() + watchdog.idle() + finally: + watchdog.stop() + text = caplog_vllm.text + assert "Engine core has spent 0 s in one engine step" in text + assert "step_waiting_on_a_transfer" in text + assert "Engine core engine step finished after 0.6" in text + + +def test_engine_loop_watchdog_ignores_an_idle_loop(caplog_vllm): + watchdog = EngineLoopWatchdog(threshold=0.1, poll_interval=0.02) + try: + watchdog.busy("engine step") + watchdog.idle() + time.sleep(0.3) + finally: + watchdog.stop() + assert "Engine core" not in caplog_vllm.text + + +def test_add_request_waiting_in_the_input_queue_is_reported(caplog_vllm): + added = [] + engine = SimpleNamespace( + _loop_watchdog=None, + _stall_warning_s=1.0, + _add_received={"req-1": time.monotonic() - 5.0}, + _reject_add_in_shutdown=lambda request: False, + add_request=lambda request, wave: added.append((request.request_id, wave)), + ) + request = SimpleNamespace(request_id="req-1") + EngineCoreProc._handle_client_request( + engine, EngineCoreRequestType.ADD, (request, 0) + ) + assert added == [("req-1", 0)] + assert engine._add_received == {} + assert "Request req-1 waited 5." in caplog_vllm.text diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 09b0ea97811e..db026cff96a9 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -19,6 +19,7 @@ log_non_default_args, log_version_and_model, ) +from vllm.entrypoints.serve.utils.task_dump import log_asyncio_tasks from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.tool_parsers import ToolParserManager @@ -134,6 +135,8 @@ async def dummy_shutdown() -> None: loop.add_signal_handler(signal.SIGINT, signal_handler) loop.add_signal_handler(signal.SIGTERM, signal_handler) + # `kill -USR1 ` logs where every request task is suspended. + loop.add_signal_handler(signal.SIGUSR1, log_asyncio_tasks) async def handle_shutdown() -> None: await shutdown_event.wait() diff --git a/vllm/entrypoints/serve/utils/task_dump.py b/vllm/entrypoints/serve/utils/task_dump.py new file mode 100644 index 000000000000..ccc57cec3906 --- /dev/null +++ b/vllm/entrypoints/serve/utils/task_dump.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Show where every asyncio task of the API server is suspended. + +A request held at an ``await`` is on no thread's stack: py-spy and +faulthandler only show the event loop waiting in ``select``. This walks each +task's chain of suspended coroutines and async generators down to the object +it waits on. Frames list local variable names and types; values are shown +only for numbers and request IDs, so prompts and outputs stay out of logs. + +Send ``SIGUSR1`` to the API server process to log the dump. +""" + +import asyncio +import gc +import inspect +import sys +import threading +import traceback +from types import FrameType +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Chains are short in practice; this only bounds a pathological cycle. +_MAX_DEPTH = 100 +_ID_NAMES = ("request_id", "sub_request_id", "external_req_id", "session_id") + + +def _describe_frame(frame: FrameType) -> str: + code = frame.f_code + name = getattr(code, "co_qualname", code.co_name) + local_parts = [] + for local_name, value in frame.f_locals.items(): + if isinstance(value, bool | int | float) or value is None: + local_parts.append(f"{local_name}={value!r}") + elif local_name in _ID_NAMES and isinstance(value, str): + local_parts.append(f"{local_name}={value[:100]!r}") + else: + local_parts.append(f"{local_name}: {type(value).__qualname__}") + return ( + f" {name} ({code.co_filename}:{frame.f_lineno})\n" + f" locals: {', '.join(local_parts) or '-'}" + ) + + +def _step(awaitable: Any) -> tuple[FrameType | None, Any]: + """Return the frame of a suspended awaitable and what it awaits next.""" + if inspect.iscoroutine(awaitable): + return awaitable.cr_frame, awaitable.cr_await + if inspect.isasyncgen(awaitable): + return awaitable.ag_frame, awaitable.ag_await + if inspect.isgenerator(awaitable): + return awaitable.gi_frame, awaitable.gi_yieldfrom + if type(awaitable).__name__.startswith("async_generator_"): + # Awaiting an async generator's next item awaits a wrapper that does + # not expose the generator, but references it. + for referent in gc.get_referents(awaitable): + if inspect.isasyncgen(referent): + return None, referent + return None, None + + +def _describe_task(task: asyncio.Task) -> str: + lines = [f" Task {task.get_name()!r}:"] + awaitable: Any = task.get_coro() + for _ in range(_MAX_DEPTH): + if awaitable is None: + break + frame, awaited = _step(awaitable) + if frame is not None: + lines.append(_describe_frame(frame)) + elif awaited is None: + lines.append(f" waits on {type(awaitable).__qualname__}") + break + awaitable = awaited + waiter = getattr(task, "_fut_waiter", None) + if waiter is not None: + # The future's callbacks name what resumes the task, e.g. another + # task that it waits for. + lines.append(f" task is blocked on {waiter!r}") + return "\n".join(lines) + + +def format_asyncio_tasks(loop: asyncio.AbstractEventLoop | None = None) -> str: + """Describe every task of ``loop`` and the stack of every other thread. + + Args: + loop: Event loop whose tasks are described; the running loop if None. + + Returns: + A multi-line report, one section per task and per thread. + """ + loop = loop or asyncio.get_running_loop() + tasks = sorted(asyncio.all_tasks(loop), key=lambda task: task.get_name()) + sections = [f"{len(tasks)} asyncio tasks:"] + sections.extend(_describe_task(task) for task in tasks) + current = threading.get_ident() + names = {thread.ident: thread.name for thread in threading.enumerate()} + for ident, frame in sys._current_frames().items(): + if ident == current: + continue + stack = "".join(traceback.format_stack(frame)).rstrip() + sections.append(f" Thread {names.get(ident, ident)!r}:\n{stack}") + return "\n".join(sections) + + +def log_asyncio_tasks() -> None: + """Signal handler: log where each task and thread of this process waits.""" + try: + logger.warning("Task dump requested by signal:\n%s", format_asyncio_tasks()) + except Exception: + logger.exception("Task dump failed") diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..f19629431b46 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -45,6 +45,8 @@ VLLM_LOGGING_COLOR: str = "auto" NO_COLOR: bool = False VLLM_LOG_STATS_INTERVAL: float = 10.0 + VLLM_REQUEST_STALL_WARNING_S: float = 60.0 + VLLM_LOG_REQUEST_TIMELINE: bool = False VLLM_TRACE_FUNCTION: int = 0 VLLM_USE_FLASHINFER_SAMPLER: bool = True VLLM_GLM53_MTP_DRAFT_HEAD: Literal["bf16", "nvfp4"] = "bf16" @@ -876,6 +878,20 @@ def _resolve_rust_cli_path() -> str | None: if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0 else 10.0 ), + # Warn when a request takes longer than this many seconds to get from + # rendering to the engine core, waits that long in the engine-core input + # queue, or leaves its final output unread by the client that long; also + # warn, with the loop's stack, when one engine-core loop iteration runs + # that long. 0 disables these warnings. + "VLLM_REQUEST_STALL_WARNING_S": lambda: max( + 0.0, float(os.getenv("VLLM_REQUEST_STALL_WARNING_S", "60")) + ), + # If set, the API server logs one line per request with the time each + # stage took: rendering, input processing, engine submission, first and + # final output, and the client reading the final output. + "VLLM_LOG_REQUEST_TIMELINE": lambda: ( + os.getenv("VLLM_LOG_REQUEST_TIMELINE", "0").lower() in ("1", "true") + ), # Trace function calls # If set to 1, vllm will trace function calls # Useful for debugging @@ -2441,6 +2457,8 @@ def compile_factors() -> dict[str, object]: "VLLM_LOGGING_CONFIG_PATH", "VLLM_LOGGING_COLOR", "VLLM_LOG_STATS_INTERVAL", + "VLLM_REQUEST_STALL_WARNING_S", + "VLLM_LOG_REQUEST_TIMELINE", "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "VLLM_TUNED_CONFIG_FOLDER", "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index c15aa3153a78..cb0d8c89a30a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -50,6 +50,7 @@ from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.v1.engine.stall_diagnostics import RequestTimeline, arrival_time from vllm.v1.executor import Executor from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult from vllm.v1.metrics.loggers import ( @@ -127,6 +128,8 @@ def __init__( init_tracer("vllm.llm_engine", tracing_endpoint) self.log_requests = log_requests + self._request_stall_warning_s = envs.VLLM_REQUEST_STALL_WARNING_S + self._log_request_timeline = envs.VLLM_LOG_REQUEST_TIMELINE custom_stat_loggers = list(stat_loggers or []) custom_stat_loggers.extend(load_stat_logger_plugin_factories()) @@ -681,6 +684,14 @@ async def generate( """ q: RequestOutputCollector | None = None + # The body runs only once the serving handler asks for output, so a + # hold before engine submission shows up as a late generator start. + timeline = ( + RequestTimeline(request_id, arrival_time(prompt), time.time()) + if self._request_stall_warning_s > 0 or self._log_request_timeline + else None + ) + outcome = "failed" try: q = await self.add_request( request_id, @@ -696,6 +707,9 @@ async def generate( reasoning_ended=reasoning_ended, reasoning_parser_kwargs=reasoning_parser_kwargs, ) + if timeline is not None: + timeline.processed = q.created + timeline.submitted = time.time() # The output_handler task pushes items into the queue. # This task pulls from the queue and yields to caller. @@ -709,13 +723,17 @@ async def generate( # own request cleanup based on finished. assert isinstance(out, RequestOutput) finished = out.finished + if timeline is not None: + timeline.output(out) if out is not STREAM_FINISHED: yield out + outcome = "finished" # If the request is disconnected by the client, generate() # is cancelled or the generator is garbage collected. So, # we abort the request if we end up here. except (asyncio.CancelledError, GeneratorExit): + outcome = "aborted" if q is not None: await self.abort(q.request_id, internal=True) if self.log_requests: @@ -760,6 +778,12 @@ async def generate( finally: if q is not None: q.close() + if timeline is not None: + timeline.close( + outcome, + self._request_stall_warning_s, + self._log_request_timeline, + ) def _run_output_handler(self): """Background loop: pulls from EngineCore and pushes to AsyncStreams.""" diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9cd3a3d33835..ce12bb2744b9 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -72,6 +72,7 @@ UtilityOutput, UtilityResult, ) +from vllm.v1.engine.stall_diagnostics import EngineLoopWatchdog from vllm.v1.engine.tensor_ipc import TensorIpcReceiver from vllm.v1.engine.utils import ( EngineHandshakeMetadata, @@ -1153,6 +1154,15 @@ def __init__( ): self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]() self.output_queue = queue.Queue[tuple[int, EngineCoreOutputs] | bytes]() + # Monotonic time at which the input thread queued each ADD request + # that the loop has not handled yet. + self._add_received: dict[str, float] = {} + self._stall_warning_s = envs.VLLM_REQUEST_STALL_WARNING_S + self._loop_watchdog = ( + EngineLoopWatchdog(self._stall_warning_s) + if self._stall_warning_s > 0 + else None + ) executor_fail_callback = lambda: self.input_queue.put_nowait( (EngineCoreRequestType.EXECUTOR_FAILED, b"") ) @@ -1515,6 +1525,12 @@ def is_running(self) -> bool: """Returns true if shutdown has not been requested.""" return self.shutdown_state == EngineShutdownState.RUNNING + def shutdown(self): + # Teardown is not a loop step; do not report it as a stalled one. + if self._loop_watchdog is not None: + self._loop_watchdog.stop() + super().shutdown() + @fault_tolerant_wrapper def run_busy_loop(self): """Core busy loop of the EngineCore.""" @@ -1557,6 +1573,8 @@ def _process_input_queue(self): logger.debug("EngineCore waiting for work.") waited = True block = self.process_input_queue_block + if block and self._loop_watchdog is not None: + self._loop_watchdog.idle() try: req = self.input_queue.get(block=block) self._handle_client_request(*req) @@ -1576,6 +1594,8 @@ def _process_input_queue(self): def _process_engine_step(self) -> bool: """Called only when there are unfinished local requests.""" + if self._loop_watchdog is not None: + self._loop_watchdog.busy("engine step") # Step the engine core. outputs, model_executed = self.step_fn() # Put EngineCoreOutputs into the output queue. @@ -1650,10 +1670,28 @@ def _handle_client_request( ) -> None: """Dispatch request from client.""" + if self._loop_watchdog is not None: + self._loop_watchdog.busy( + f"{request[2]} utility call" + if request_type == EngineCoreRequestType.UTILITY + else f"{request_type.name} client request" + ) if request_type == EngineCoreRequestType.WAKEUP: return elif request_type == EngineCoreRequestType.ADD: req, request_wave = request + received = self._add_received.pop(req.request_id, None) + if ( + received is not None + and self._stall_warning_s > 0 + and (waited := time.monotonic() - received) > self._stall_warning_s + ): + logger.warning( + "Request %s waited %.1f s in the engine-core input queue " + "before the scheduler received it", + req.request_id, + waited, + ) if self._reject_add_in_shutdown(req): return self.add_request(req, request_wave) @@ -1899,6 +1937,8 @@ def process_input_sockets( # aborting in the scheduler is idempotent. self.aborts_queue.put_nowait(request) + if request_type == EngineCoreRequestType.ADD: + self._add_received[request[0].request_id] = time.monotonic() # Push to input queue for core busy loop. self.input_queue.put_nowait((request_type, request)) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 257287243dbd..9fa7bee2a773 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import time from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass @@ -62,6 +63,8 @@ def __init__(self, output_kind: RequestOutputKind, request_id: str): self.request_id = request_id self.output: RequestOutput | PoolingRequestOutput | Exception | None = None self.ready = asyncio.Event() + # Wall-clock creation time; the request's inputs are processed by then. + self.created = time.time() self._input_stream_task: asyncio.Task | None = None diff --git a/vllm/v1/engine/stall_diagnostics.py b/vllm/v1/engine/stall_diagnostics.py new file mode 100644 index 000000000000..e50ede7b9e07 --- /dev/null +++ b/vllm/v1/engine/stall_diagnostics.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Name the stage where a request or the engine-core loop stops moving. + +The periodic stats line cannot tell a blocked engine from an idle one: both +stop producing tokens, and the stats logger goes quiet for both. A request +that waits outside the engine leaves no trace at all. These helpers log the +stage and its duration instead. +""" + +import sys +import threading +import time +import traceback +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + +logger = init_logger(__name__) + +# Innermost loop-thread frames included in a stall report. +_STACK_FRAMES = 40 + + +def _clock(timestamp: float) -> str: + """Format a wall-clock time like the log line prefix, with milliseconds.""" + return time.strftime("%H:%M:%S", time.localtime(timestamp)) + ( + f".{int(timestamp % 1 * 1000):03d}" + ) + + +def _span(start: float | None, end: float | None) -> str: + if start is None or end is None: + return "-" + return f"{end - start:.2f} s" + + +class RequestTimeline: + """Wall-clock times at which one request passed each API-server stage. + + Args: + request_id: Request whose stages are recorded. + arrival: The renderer's arrival stamp, taken when rendering started, + or None when the prompt carried none. + started: When the result generator started running, which is when + the serving handler first asked it for output. + + Stages the request never reached stay None. + """ + + __slots__ = ( + "request_id", + "arrival", + "started", + "processed", + "submitted", + "first_output", + "final_output", + "closed", + "last_output", + ) + + def __init__(self, request_id: str, arrival: float | None, started: float): + self.request_id = request_id + self.arrival = arrival + self.started = started + # Input processing finished and the output collector was created. + self.processed: float | None = None + # The engine-core client accepted the request. + self.submitted: float | None = None + self.first_output: float | None = None + self.final_output: float | None = None + # The result generator stopped: after the handler read the final + # output, or when the request was aborted or failed. + self.closed: float | None = None + self.last_output: RequestOutput | None = None + + def output(self, output: "RequestOutput") -> None: + """Record an output handed to the serving handler.""" + now = time.time() + if self.first_output is None: + self.first_output = now + if output.finished: + self.final_output = now + self.last_output = output + + def outside_engine(self) -> float: + """Longest stretch spent outside engine work, in seconds. + + That is rendering until the generator started, input processing plus + submission, or the handler taking the final output to the client. + """ + waits = [0.0] + if self.arrival is not None: + waits.append(self.started - self.arrival) + if self.submitted is not None: + waits.append(self.submitted - self.started) + if self.final_output is not None and self.closed is not None: + waits.append(self.closed - self.final_output) + return max(waits) + + def describe(self, outcome: str) -> str: + """Return a one-line account of every stage the request reached.""" + assert self.closed is not None + origin = self.arrival if self.arrival is not None else self.started + parts = [ + f"Request {self.request_id} {outcome} after " + f"{self.closed - origin:.2f} s (arrived {_clock(origin)})", + f"arrival to generator start {_span(self.arrival, self.started)}", + f"input processing {_span(self.started, self.processed)}", + f"engine submission {_span(self.processed, self.submitted)}", + f"submission to first output {_span(self.submitted, self.first_output)}", + f"first to final output {_span(self.first_output, self.final_output)}", + f"final output to stream end {_span(self.final_output, self.closed)}", + ] + output = self.last_output + if output is not None: + stats = output.metrics + if stats is not None and stats.queued_ts and stats.scheduled_ts: + parts.append( + f"engine queue {stats.scheduled_ts - stats.queued_ts:.2f} s" + ) + if output.prompt_token_ids is not None: + parts.append( + f"prompt {len(output.prompt_token_ids)} tokens " + f"({output.num_cached_tokens or 0} cached)" + ) + if stats is not None: + parts.append(f"{stats.num_generation_tokens} generated") + if output.outputs and output.outputs[0].finish_reason is not None: + parts.append(f"finish reason {output.outputs[0].finish_reason}") + return ", ".join(parts) + + def close(self, outcome: str, threshold: float, always: bool) -> None: + """Stamp the end of the request and log its stages when warranted. + + Args: + outcome: How the request ended, e.g. "finished" or "aborted". + threshold: Seconds above which a warning is logged: spent outside + engine work, or in total by a request that ended without any + output. 0 never warns. + always: Log every request at INFO level. + """ + self.closed = time.time() + if self.final_output is not None: + # A handler may stop iterating once it has the final output. + outcome = "finished" + if always: + logger.info("%s", self.describe(outcome)) + return + if threshold <= 0: + return + origin = self.arrival if self.arrival is not None else self.started + if self.outside_engine() > threshold or ( + self.first_output is None and self.closed - origin > threshold + ): + logger.warning("%s", self.describe(outcome)) + + +def arrival_time(prompt: Any) -> float | None: + """Return the renderer's arrival stamp carried by a prompt, if any.""" + value = ( + prompt.get("arrival_time") + if isinstance(prompt, dict) + else getattr(prompt, "arrival_time", None) + ) + return value if isinstance(value, float | int) else None + + +class EngineLoopWatchdog: + """Log the loop's stack when one engine-core loop step runs too long. + + The loop thread calls ``busy(stage)`` before each unit of work and + ``idle()`` before blocking for new work. A daemon thread checks the + current unit. Once a unit runs past the threshold it logs the loop + thread's stack, again after each further threshold, and the loop thread + logs the unit's total time when the unit ends. + + Args: + threshold: Seconds a single unit may run before it is reported. + poll_interval: Seconds between checks; defaults to a quarter of the + threshold, at most five seconds. + """ + + def __init__(self, threshold: float, poll_interval: float | None = None): + if threshold <= 0: + raise ValueError("The engine loop stall threshold must be positive") + self._threshold = threshold + self._poll_interval = poll_interval or min(5.0, threshold / 4) + # (stage, monotonic start, loop thread id) of the running unit. + self._current: tuple[str, float, int] | None = None + # The unit last reported, and how many times. + self._reported: tuple[str, float, int] | None = None + self._reports = 0 + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._watch, name="EngineLoopWatchdog", daemon=True + ) + self._thread.start() + + def busy(self, stage: str) -> None: + """Mark the start of a unit of loop work, ending the previous one.""" + self._finish() + self._current = (stage, time.monotonic(), threading.get_ident()) + + def idle(self) -> None: + """Mark that the loop is about to block waiting for new work.""" + self._finish() + self._current = None + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=self._poll_interval + 1) + + def _finish(self) -> None: + current = self._current + if current is not None and current is self._reported: + logger.warning( + "Engine core %s finished after %.1f s", + current[0], + time.monotonic() - current[1], + ) + + def _watch(self) -> None: + while not self._stop.wait(self._poll_interval): + current = self._current + if current is None: + continue + stage, started, thread_id = current + reports = self._reports if current is self._reported else 0 + elapsed = time.monotonic() - started + if elapsed < self._threshold * (reports + 1): + continue + frame = sys._current_frames().get(thread_id) + stack = ( + "".join(traceback.format_stack(frame)[-_STACK_FRAMES:]) + if frame is not None + else " (stack unavailable)\n" + ) + logger.warning( + "Engine core has spent %.0f s in one %s without finishing; " + "the engine is blocked, not idle. Loop thread stack:\n%s", + elapsed, + stage, + stack.rstrip(), + ) + self._reported = current + self._reports = reports + 1 From 75f660a9356f9b0b8aa608101919122e03339656 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 20:03:12 +0000 Subject: [PATCH 195/365] chore(release): describe the request stall diagnostics Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-866.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-866.json diff --git a/.lil/changes/vllm-866.json b/.lil/changes/vllm-866.json new file mode 100644 index 000000000000..78e3884f923d --- /dev/null +++ b/.lil/changes/vllm-866.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-866", + "category": "feature", + "summary": "Log where a request or the engine-core loop is held", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "A warning names the stage when a request spends more than VLLM_REQUEST_STALL_WARNING_S (default 60) outside engine work, or ends without output after that long; VLLM_LOG_REQUEST_TIMELINE=1 logs every request's stages.", + "The engine core warns about ADD requests that waited in its input queue and logs the loop stack when one step runs longer than the threshold.", + "SIGUSR1 to the API server logs every asyncio task's await chain and every thread's stack, without prompt or output text." + ], + "pull_requests": [866] +} From 5bf6e0d6adad3f322786fb7568cb4f8eb87dcdaf Mon Sep 17 00:00:00 2001 From: Carlos Augusto Date: Wed, 23 Sep 2026 14:08:24 -0600 Subject: [PATCH 196/365] [KV Connector] Track allocated KV blocks before early return OffloadingConnectorScheduler.update_state_after_alloc returned early when num_external_tokens == 0 without recording the newly allocated blocks in _current_batch_allocated_block_ids. Untracked blocks could be reallocated to another request while an in-flight CPU load still referenced them, producing cross-request KV corruption (answers mixing conversations) on MultiConnector / zero-token-load configurations. Fix: record every non-null allocation through _record_allocated_blocks before the early return, and drop the redundant per-group inline update. Attribution: this is a verbatim port of vllm-project/vllm#46960 (issue #46951), author yinli-systems, which is still open as a draft. The hunks are upstream's, not ours; carry patch-id c0f42a7bfc6ca4a3a2cf4de893d33e3b3529a2be to match them after upstream rebases or squash-merges. Drop this commit once the PR lands and karmic rebases onto it. It applies to karmic unchanged because the group-indexed traversal (blocks.blocks[group_config.group_idx]) that karmic carries is already the shape the rebased PR diff targets. Includes the PR's regression test (multiple cache groups, null block ID 0) in tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py. Live only under the complex OffloadingConnector (VLLM_USE_SIMPLE_KV_OFFLOAD=0); inert under SimpleCPUOffloadConnector. --- .../offloading_connector/test_scheduler.py | 22 +++++++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 13 +++++++---- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 9693c36bea98..a368673660f0 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2136,6 +2136,28 @@ def capture_fence(): assert runner.connector_scheduler._block_id_to_pending_jobs == {} +def test_update_state_after_alloc_records_zero_token_allocations(): + scheduler = object.__new__(OffloadingConnectorScheduler) + scheduler._current_batch_allocated_block_ids = set() + blocks = KVCacheBlocks( + ( + ( + KVCacheBlock(block_id=7), + KVCacheBlock(block_id=0), + ), + (KVCacheBlock(block_id=11),), + ) + ) + + scheduler.update_state_after_alloc( + SimpleNamespace(request_id="req-0"), + blocks, + num_external_tokens=0, + ) + + assert scheduler._current_batch_allocated_block_ids == {7, 11} + + def test_fence_at_build_store_jobs(request_runner): """A new prefill (no load -> update_state_after_alloc returns early) reusing a finished request's pending-store block is flushed by diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index ff31c12da537..8e5f6ccf1603 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1115,6 +1115,7 @@ def get_num_new_matched_tokens( def update_state_after_alloc( self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int ): + self._record_allocated_blocks(blocks) if num_external_tokens == 0: return @@ -1136,10 +1137,6 @@ def update_state_after_alloc( req_status.group_states, ): group_blocks = blocks.blocks[group_config.group_idx] - self._current_batch_allocated_block_ids.update( - block.block_id for block in group_blocks if block.block_id != 0 - ) - tokens_per_block = group_config.tokens_per_block tokens_per_chunk = group_config.tokens_per_chunk offload_keys = group_state.offload_keys @@ -1228,6 +1225,14 @@ def update_state_after_alloc( self._chunks_being_loaded.update(keys_to_load) req_status.partial_tail_boundary = None + def _record_allocated_blocks(self, blocks: KVCacheBlocks) -> None: + self._current_batch_allocated_block_ids.update( + block.block_id + for group_blocks in blocks.blocks + for block in group_blocks + if block.block_id != 0 + ) + def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: """ Update request states from the Scheduler's output. From 495eb9ba863911d7681b0458f2f6127f38034231 Mon Sep 17 00:00:00 2001 From: Carlos Augusto Date: Wed, 23 Sep 2026 14:08:24 -0600 Subject: [PATCH 197/365] [KV Connector] Keep group store ranges aligned across skipped CoW groups OffloadingConnectorScheduler._build_store_jobs builds group_store_ranges, one (start_chunk_idx, num_chunks) per KV cache group, by zipping config.kv_group_configs with req_status.group_states. A group whose spec is MambaSpec with mamba_cache_mode == "align" sets requires_cow_source=True and was skipped with a bare continue before the append, so those groups got no entry. The consumer then zips kv_group_configs (N) x group_states (N) x group_store_ranges (N minus the CoW count): zip truncates to the shortest, so group_sizes is built with the wrong length and every surviving group's chunk range is paired with the wrong group. The worker (vllm/v1/kv_offload/cpu/gpu_worker.py, transfer_async) asserts len(group_sizes) == len(self.layer_refs_per_group) and EngineCore dies. On Qwen3.8-Flash-Next (50 KV cache groups, 13 align-mode Mamba) the scheduler emits 37-entry group_sizes against the worker's 50-entry layer_refs_per_group, so the first request large enough to complete an offloadable chunk (~15k tokens) crashes the engine. This is why the complex connector cannot serve hybrid align-mode Mamba models at all. Fix: append a (0, 0) placeholder for CoW groups instead of a bare continue, so group_store_ranges stays one entry per group and positionally aligned. The consumer's inner loop over offload_keys[0:0] is empty and contributes a 0 at that position, which is correct: CoW groups persist through the boundary / partial-tail store paths, not the chunk path. The now-dead second `if group_config.requires_cow_source: continue` after the append is removed. The sibling builders (_build_aligned_boundary_store_jobs, _build_partial_tail_store_jobs) already build group_sizes as [0] * num_groups and are unaffected. Also assert at the _build_store_jobs producer exit that every KV group contributes exactly one store range, so a future pre-append continue names itself at the source instead of firing the worker's distant length assert. Live only under the complex OffloadingConnector; inert under SimpleCPUOffloadConnector, which does not use this scheduler. --- .../kv_connector/v1/offloading/scheduler.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 8e5f6ccf1603..521f8b4a15c6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1573,6 +1573,16 @@ def _build_store_jobs( ) if group_config.requires_cow_source: + # CoW (mamba align-mode) groups are persisted through the + # boundary / partial-tail store path, not the chunk path. + # Append a (0, 0) placeholder rather than a bare continue so + # group_store_ranges stays positionally aligned with + # kv_group_configs. The consumer zip below pairs each range + # with its own group; a shorter list truncates group_sizes + # AND misaligns every surviving group's chunk range, which + # trips the worker assert + # len(group_sizes) == len(layer_refs_per_group). + group_store_ranges.append((0, 0)) continue start_chunk_idx = group_state.next_stored_chunk_idx @@ -1613,9 +1623,6 @@ def _build_store_jobs( start_chunk_idx = min(start_chunk_idx, partial_segment_start_chunk) group_store_ranges.append((start_chunk_idx, num_chunks)) - if group_config.requires_cow_source: - continue - if num_chunks <= start_chunk_idx: continue offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks] @@ -1682,6 +1689,13 @@ def _build_store_jobs( continue new_offload_keys.append(offload_key) + assert len(group_store_ranges) == len(self.config.kv_group_configs), ( + "every KV group must contribute exactly one store range; a " + "pre-append continue would silently truncate the consumer " + "zip and misalign group_sizes for the worker, whose own " + "assert is far from the cause" + ) + if not new_offload_keys: req_status.advance_stored_idx(num_offloadable_tokens) continue From f26b5d3bab1ab360f88800bac515b675dc768156 Mon Sep 17 00:00:00 2001 From: Carlos Augusto Date: Wed, 23 Sep 2026 14:08:24 -0600 Subject: [PATCH 198/365] [KV Connector] Serve external prefix hits for align-mode Mamba hybrids With the complex OffloadingConnector enabled for a hybrid model that has align-mode Mamba groups (Qwen3.8-Flash-Next: 50 KV cache groups, 13 of them GDN recurrent state), multi-turn serving corrupts deterministically at temp=0: the cold turn (cached_tokens=0) answers correctly; every later turn with a prefix hit answers a generic greeting or nothing, with no assert and a healthy engine. The model reads its own context as empty. Inert under SimpleCPUOffloadConnector, which never stores Mamba groups. Two independent defects on the external-hit path: 1. Stale boundary capture. Align-mode Mamba migrates its state column once per step, pre-forward. A chunked-prefill step computes an entire chunk at once (e.g. tokens 6016..9101 in the step starting at nct=6016), so the column ends the step holding the chunk-END state; the intermediate retention-boundary state (9024) never exists in any column. The connector's deferred store DMA faithfully copied the chunk-end state and filed it under the boundary key. On resume, the request recomputed the chunk tail on top of a state that already included it: recurrent-state double-processing. 2. Untransported padded tail. QSA attention groups pad their pages; the compressed-K selector tail is the [unpadded, padded) region (192,512 B here), read by selection through compressed_block_table. Narrow transfer refs carried the unpadded page size, so that region was never DMA'd in either direction: a resumed request selected against an empty tail while every copied byte looked correct. Fix: - Boundary stores source dedicated capture blocks, written once by the boundary-checkpoint crossing copy and manager-pinned from offer until store ack or drop; captures recorded after their producing step completed are withheld. Pins release on every store drop path (out-of-bounds, allocation failure, key dedup, stale-ack drain with live-fence guard), and block-multiple geometry is asserted at record and config sites. - Attention transfer refs carry the full padded page (spec.page_size_bytes); store and load share one size source; canonical ref width comes from each layer's own group spec (a stale loop variable made the padded-tail transport silently no-op on hybrid layouts whose last group is Mamba); Mamba refs keep the unpadded width. - The QSA selector pool fresh flag is derived from the restored-checkpoint signal at add time and cleared on boundary (get_recurrent_checkpoint_fresh); the anchor seeds to prefix-1. GLM-Next mirrors the gating. Local hits and non-offload paths keep stock always-fresh behavior. - The external hit extends to the deepest boundary every group keys, never a half-restore. Tests ship with the patch: crossing-capture digest identity, pin release on store-ack and every offer-drop path, hit-gate convergence, padded-ref register width per backend plus a GPU<->CPU DMA roundtrip, QSA fresh gating on cold/resume/local-hit adds, store-key dedup, and the b12x GDN prefill checkpoint clone. Live serving validation (drizzle, RTX PRO 6000, temp=0, seed4+fixed nonce, GPU cache capped at ~49,955 tokens so a resend of eight ~10.4k-token prompts cannot fit in GPU): regular phase 9/9 correct, resume sends cached_tokens=9024, corrupted=0; evict phase serving=8/8 correct=8/8 with all B-resends at cached=9024 from the CPU tier. Pre-patch on the same model and config: the first store of a large prompt crashed the engine (fixed by the CoW alignment commit), and for requests that did complete stores, every prefix-hit turn served a deterministic wrong answer. Depends on: the block-tracking commit in this series, a port of vllm-project/vllm#46960 (external-hit serving assumes it for MultiConnector zero-token allocations). --- .../test_gdn_prefill_b12x_checkpoint_clone.py | 103 ++++++ tests/models/test_glm5next_model_state.py | 20 +- .../test_mamba_boundary_capture.py | 345 ++++++++++++++++++ .../test_mamba_boundary_capture_digest.py | 237 ++++++++++++ .../offloading_connector/test_scheduler.py | 115 ++++++ .../unit/offloading_connector/test_worker.py | 228 ++++++++++++ tests/v1/kv_offload/cpu/test_manager.py | 25 ++ .../test_qwen4_exp_qsa_resume_fresh_gating.py | 265 ++++++++++++++ .../kv_transfer/kv_connector/v1/base.py | 10 + .../kv_connector/v1/multi_connector.py | 7 + .../kv_connector/v1/offloading/scheduler.py | 93 ++++- .../kv_connector/v1/offloading/worker.py | 14 +- .../kv_connector/v1/offloading_connector.py | 6 +- vllm/model_executor/layers/mamba/abstract.py | 25 +- vllm/models/glm5next/model_state.py | 18 +- vllm/models/qwen4_exp/nvidia/model_state.py | 49 ++- vllm/v1/core/kv_cache_manager.py | 14 + vllm/v1/core/sched/scheduler.py | 6 + vllm/v1/core/single_type_kv_cache_manager.py | 169 +++++++++ vllm/v1/kv_cache_interface.py | 7 + vllm/v1/worker/gpu/boundary_checkpoint.py | 5 + vllm/v1/worker/gpu/model_states/interface.py | 11 + 22 files changed, 1761 insertions(+), 11 deletions(-) create mode 100644 tests/kernels/mamba/test_gdn_prefill_b12x_checkpoint_clone.py create mode 100644 tests/v1/core/prefix_cache/test_mamba_boundary_capture.py create mode 100644 tests/v1/core/prefix_cache/test_mamba_boundary_capture_digest.py create mode 100644 tests/v1/worker/test_qwen4_exp_qsa_resume_fresh_gating.py diff --git a/tests/kernels/mamba/test_gdn_prefill_b12x_checkpoint_clone.py b/tests/kernels/mamba/test_gdn_prefill_b12x_checkpoint_clone.py new file mode 100644 index 000000000000..d05b7ff910f7 --- /dev/null +++ b/tests/kernels/mamba/test_gdn_prefill_b12x_checkpoint_clone.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# The b12x GDN prefill kernel must clone state at +# checkpoint_offset == query_len (a step-end boundary crossing) into the +# designated checkpoint slot, matching the reference oracle. The +# strictly-interior offset runs as the control. +import pytest +import torch + + +def _setup(lengths, offsets, ckpt_slots, max_seqs): + from b12x.testing.delta_prefill_cases import PrefillCase, make_inputs + + device = torch.device("cuda") + case = PrefillCase("gdn", 2, 6, lengths) + max_tokens = sum(lengths) + tensors = make_inputs( + case, device=device, max_tokens=max_tokens, max_seqs=max_seqs + ) + # Distinct non-final, non-initial checkpoint slots (pool has 3*max_seqs+1). + tensors["checkpoint_state_indices"][: len(lengths)] = torch.tensor( + ckpt_slots, dtype=torch.int32, device=device + ) + tensors["checkpoint_offsets"][: len(lengths)] = torch.tensor( + offsets, dtype=torch.int32, device=device + ) + return case, tensors + + +def _run_case(lengths, offsets, ckpt_slots, max_seqs=2): + from b12x.testing.delta_prefill_cases import ( + assert_close, + oracle, + prepared_binding, + run_binding, + ) + + case, tensors = _setup(lengths, offsets, ckpt_slots, max_seqs) + with prepared_binding( + case, + tensors, + max_tokens=sum(lengths), + max_seqs=max_seqs, + checkpoint_export=True, + null_state_index=0, + ) as binding: + run_binding("gdn", binding) + torch.cuda.synchronize() + expected_output, expected_pool = oracle(case, tensors, null_state_index=0) + assert_close( + "output", + binding.output[: sum(lengths)], + expected_output[: sum(lengths)], + ratio=1e-2, + ) + for slot, offset in zip(ckpt_slots, offsets): + assert_close( + f"state[{slot}]@offset={offset}", + binding.recurrent_state[slot], + expected_pool[slot], + ratio=5e-3, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_v1_checkpoint_clone_interior_offset(): + """Control: interior offset (32 < query_len 64) exports state@32.""" + _run_case(lengths=(64,), offsets=(32,), ckpt_slots=(4,)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_v1_checkpoint_clone_at_query_end(): + """The design's step-end case: offset == query_len (64) must export + state@64 into the checkpoint slot (distinct from the final slot).""" + _run_case(lengths=(64,), offsets=(64,), ckpt_slots=(4,)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_v1_checkpoint_clone_at_query_end_multi_seq(): + """Step-end clone for every sequence in a packed query.""" + _run_case(lengths=(64, 48), offsets=(64, 48), ckpt_slots=(4, 5)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_v1_checkpoint_clone_offset_exceeds_query_len_no_write(): + """offset > query_len: reference writes nothing; kernel must leave the + slot untouched (design gate never emits this, guard against corruption).""" + from b12x.testing.delta_prefill_cases import prepared_binding, run_binding + + case, tensors = _setup((64,), (96,), (4,), max_seqs=2) + with prepared_binding( + case, + tensors, + max_tokens=64, + max_seqs=2, + checkpoint_export=True, + null_state_index=0, + ) as binding: + saved = binding.recurrent_state[4].clone() + run_binding("gdn", binding) + torch.cuda.synchronize() + torch.testing.assert_close( + binding.recurrent_state[4], saved, rtol=0, atol=0 + ) diff --git a/tests/models/test_glm5next_model_state.py b/tests/models/test_glm5next_model_state.py index ae40b1def157..b965d1fd3a95 100644 --- a/tests/models/test_glm5next_model_state.py +++ b/tests/models/test_glm5next_model_state.py @@ -293,7 +293,14 @@ def test_glm5next_recycled_and_rebound_state_is_fresh(monkeypatch) -> None: lambda self: None, ) - state.add_request(5, SimpleNamespace(num_computed_tokens=0)) + state.add_request( + 5, + SimpleNamespace( + num_computed_tokens=0, + boundary_checkpoint=None, + boundary_checkpoint_blocks=None, + ), + ) assert state.selector_state_is_fresh_gpu[5] assert state.selector_committed_num_accepted_tokens_gpu[5] == 1 @@ -330,7 +337,9 @@ def test_glm5next_rejects_unaligned_fresh_prefix( state.add_request( 3, SimpleNamespace( - num_computed_tokens=prefix_length, boundary_checkpoint=None + num_computed_tokens=prefix_length, + boundary_checkpoint=None, + boundary_checkpoint_blocks=None, ), ) @@ -347,7 +356,12 @@ def test_glm5next_accepts_pool_aligned_fresh_prefix(monkeypatch) -> None: ) state.add_request( - 3, SimpleNamespace(num_computed_tokens=4, boundary_checkpoint=None) + 3, + SimpleNamespace( + num_computed_tokens=4, + boundary_checkpoint=None, + boundary_checkpoint_blocks=None, + ), ) assert calls == [3] diff --git a/tests/v1/core/prefix_cache/test_mamba_boundary_capture.py b/tests/v1/core/prefix_cache/test_mamba_boundary_capture.py new file mode 100644 index 000000000000..6787b91e15ae --- /dev/null +++ b/tests/v1/core/prefix_cache/test_mamba_boundary_capture.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Frozen boundary capture for align-mode Mamba + OffloadingConnector. + +The connector's boundary store must never source the request's live running +column: a chunked step's column holds the state at the CHUNK end, and the +deferred DMA would read it only after the next forward migrated the column +onward. With capture enabled, each retention crossing the step actually +computes is copied into a dedicated single-writer pool block before the +column advances; that block carries the boundary hash and is the sole store +source, pinned until the store DMA ack. + +These tests pin the scheduler-side contract (design invariants I1/I4/I5): +offers appear only after promotion, the offered id is the capture block and +never the column, the capture carries the boundary hash, and pins are +released exactly on store ack / drop / reset. +""" + +import pytest +import torch + +from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import ( + KVCacheBlockCopy, + init_none_hash, + make_block_hash_with_group_id, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _free_ids(manager): + return { + block.block_id + for block in manager.block_pool.free_block_queue.get_all_free_blocks() + } + + +def _make_capture_manager(*, boundary_capture: bool, num_blocks: int = 16): + """Hybrid manager with mamba block size 4 and dense retention. + + A 4-token chunk ending on the block boundary is the toy analogue of a + retention-grid crossing: cache_blocks records it, and only the next + schedule pass may promote the capture. + """ + block_size = 4 + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + boundary_capture=boundary_capture, + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + ) + return manager, block_size + + +def _mamba_manager(manager): + return manager.coordinator.single_type_managers[1] + + +def test_capture_offer_names_the_frozen_block_not_the_column(): + """I1/I4: the (group, block, B) hand-off appears only at promotion, and + its block is the dedicated capture -- never the live column the next + forward migrates.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + req = make_request("p", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks["p"][0] + + # Recorded, not yet offered: the producing step is still the current one. + assert "p" in mgr._pending_captures + assert manager.take_boundary_state_offloads() == {} + assert manager.take_kv_cache_block_copies() == ([], []) + + # Next schedule pass promotes: copy rail + offer, same scheduler output. + req.num_computed_tokens = 4 + manager.new_step_starts() + + ((group_id, capture_id, boundary),) = manager.take_boundary_state_offloads()["p"] + assert group_id == 1 + assert boundary == 4 + assert capture_id != column.block_id + + copies, retained = manager.take_kv_cache_block_copies() + assert copies == [ + KVCacheBlockCopy(src_block_id=column.block_id, dst_block_id=capture_id) + ] + assert {b.block_id for b in retained} == {column.block_id, capture_id} + + # The capture itself carries the boundary hash for its group (local hits + # past the last live column resume from it too). The live column keeps the + # same key until it is re-hashed, so assert on the capture's own fields. + block_hash = req.block_hashes[4 // block_size - 1] + capture = manager.block_pool.blocks[capture_id] + assert capture.block_hash == make_block_hash_with_group_id(block_hash, 1) + assert capture.block_hash_num_tokens == 4 + + # Pinned (I4) until the store ack: dropping the rail refs leaves exactly + # the manager pin; the ack releases the last one. + manager.block_pool.free_blocks(retained) + assert capture.ref_cnt == 1 + mgr.release_boundary_capture(capture_id) + assert capture.ref_cnt == 0 + assert capture_id in _free_ids(manager) + + manager.free(req) + + +def test_capture_survives_request_teardown_until_store_ack(): + """A promoted capture is pool-owned, not request-owned: aborting the + producer before the copy step is processed must not release the frozen + content or its hash; the store ack does.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + req = make_request("r", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks["r"][0] + column_id = column.block_id + + req.num_computed_tokens = 4 + manager.new_step_starts() + ((_, capture_id, boundary),) = manager.take_boundary_state_offloads()["r"] + assert boundary == 4 + + # The request dies before the copy step is processed (abort). The capture + # stays pinned: its content is owed to the CPU tier, and it is no longer + # reachable through any column. + manager.free(req) + assert manager.block_pool.blocks[capture_id].ref_cnt >= 1 + assert manager.block_pool.blocks[capture_id].block_hash is not None + + # Scheduler processed the copy step: rail refs drop; then the store ack. + _, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + assert manager.block_pool.blocks[capture_id].ref_cnt == 1 + mgr.release_boundary_capture(capture_id) + assert capture_id in _free_ids(manager) + # The column's rail ref went with the same release path. + assert manager.block_pool.blocks[column_id].ref_cnt == 0 + + +def test_no_capture_without_enablement_offers_live_column(): + """The capture path is opt-in: a deployment without boundary_capture + keeps the legacy live-column hand-off and enqueues no copy.""" + manager, block_size = _make_capture_manager(boundary_capture=False) + mgr = _mamba_manager(manager) + + req = make_request("x", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks["x"][0] + + ((group_id, block_id, boundary),) = manager.take_boundary_state_offloads()["x"] + assert (group_id, block_id, boundary) == (1, column.block_id, 4) + assert manager.take_kv_cache_block_copies() == ([], []) + assert mgr._pending_captures == {} + manager.free(req) + + +def test_release_is_noop_for_unknown_and_non_capture_ids(): + """The connector may pass fenced ids from any store kind; only capture + pins react.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + mgr.release_boundary_capture(12345) # unknown id: no-op + + req = make_request("n", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks["n"][0] + column_id = column.block_id + + req.num_computed_tokens = 4 + manager.new_step_starts() + ((_, capture_id, _),) = manager.take_boundary_state_offloads()["n"] + _, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + + # A live column id is never in the pin dict: releasing it must not touch + # the request's reference. + before = manager.block_pool.blocks[column_id].ref_cnt + assert before >= 1 + mgr.release_boundary_capture(column_id) + assert manager.block_pool.blocks[column_id].ref_cnt == before + mgr.release_boundary_capture(capture_id) + assert manager.block_pool.blocks[capture_id].ref_cnt == 0 + manager.free(req) + + +def test_release_all_drops_pins_and_pending_records(): + """Cache reset releases every pin and every record-time column ref.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + req = make_request("a", list(range(12)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + first_column = mgr.req_to_blocks["a"][0] + + req.num_computed_tokens = 4 + manager.new_step_starts() + ((_, capture_id, _),) = manager.take_boundary_state_offloads()["a"] + _, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + assert manager.block_pool.blocks[capture_id].ref_cnt == 1 + + # Second chunk [4,8): records a new crossing that is never promoted. + assert manager.allocate_slots(req, 4) is not None + assert "a" in mgr._pending_captures + second_column = mgr.req_to_blocks["a"][1] + assert second_column.ref_cnt >= 2 # table + record-time ref + + mgr.release_all_boundary_captures() + assert mgr._capture_pinned == {} + assert mgr._pending_captures == {} + assert manager.block_pool.blocks[capture_id].ref_cnt == 0 + # The record-time ref on the second column went with the dropped record. + assert second_column.ref_cnt == 1 + manager.free(req) + assert first_column.block_id in _free_ids(manager) + + +def test_crossing_capture_keyed_identically_across_diverging_prompts(): + """Key-level identity: two prompts sharing the prefix through B=4 + diverge after it; each crossing capture is registered under the SAME + boundary key (the chain hash at 4), so a consumer of the store for one + prompt resumes the other. The captures are distinct blocks -- content + is copied independently of either live column.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + capture_ids = [] + for name, tail in (("p1", 8), ("p2", 9)): + req = make_request(name, [0, 1, 2, 3, tail, tail + 1], block_size, sha256) + # Both must PRODUCE a capture at 4: without the skip the second + # prompt would simply resume from the first prompt's frozen block. + req.skip_reading_prefix_cache = name == "p2" + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + req.num_computed_tokens = 4 + manager.new_step_starts() + ((_, capture_id, boundary),) = manager.take_boundary_state_offloads()[name] + assert boundary == 4 + capture_ids.append(capture_id) + _, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + + first, second = capture_ids + assert first != second # dedicated block per producer, never shared + # Both prompts hashed the same chain at 4 tokens (shared prefix). + key = req.block_hashes[4 // block_size - 1] + keyed = make_block_hash_with_group_id(key, 1) + assert manager.block_pool.blocks[first].block_hash == keyed + assert manager.block_pool.blocks[second].block_hash == keyed + # Both pins drop only on the store acks. + mgr.release_boundary_capture(first) + mgr.release_boundary_capture(second) + assert manager.block_pool.blocks[first].ref_cnt == 0 + assert manager.block_pool.blocks[second].ref_cnt == 0 + + +def test_late_record_after_load_credit_is_withheld(): + """Patch A: a crossing whose producing step already completed must not + be recorded. When cache_blocks only catches up to the boundary after the + optimistic mirror already sits at it (delayed connector-load credit, the + probe-#4 late-record class: nct==bnd at record, delta=prompt_final-bnd at + drain), promotion would copy a column a later forward already advanced + past the boundary -- state@prompt_final keyed at the boundary, the + original mislabel defect. The boundary simply keeps no capture and hits + fall back to recompute.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + req = make_request("d", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + # The producing step runs with caching delayed (connector-load shape). + assert ( + manager.allocate_slots( + req, 4, num_computed, computed_blocks, delay_cache_blocks=True + ) + is not None + ) + assert mgr._pending_captures == {} + + # Mirror advances to the boundary; a later pass catches cache_blocks up. + req.num_computed_tokens = 4 + manager.coordinator.cache_blocks(req, 4) + assert mgr._pending_captures == {}, "late catch-up must not record" + + manager.new_step_starts() + assert manager.take_boundary_state_offloads() == {} + manager.free(req) + + +def test_in_step_record_still_captured(): + """Control for the gate: the step being dispatched now still has the + boundary ahead of it (nct < boundary) -- the record proceeds and the + next pass promotes exactly as before.""" + manager, block_size = _make_capture_manager(boundary_capture=True) + mgr = _mamba_manager(manager) + + req = make_request("c", list(range(8)), block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + assert "c" in mgr._pending_captures + manager.free(req) diff --git a/tests/v1/core/prefix_cache/test_mamba_boundary_capture_digest.py b/tests/v1/core/prefix_cache/test_mamba_boundary_capture_digest.py new file mode 100644 index 000000000000..cd9b3e6d3758 --- /dev/null +++ b/tests/v1/core/prefix_cache/test_mamba_boundary_capture_digest.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Crossing capture digest identity. + +Two prompts sharing the prefix through boundary B, differing after it, each +run through the real scheduler-side capture pipeline (record -> promote -> +copy rail) against simulated GPU state tensors. The frozen capture block's +bytes at B MUST be identical across the two prompts (deterministic +recurrence over the shared prefix) and MUST differ from each prompt's +chunk-end column bytes after the next step migrates the column onward -- +the exact value a live-column store would mislabel as state@B. Without +boundary capture the hand-off names the live column, so the store-time +digest is the post-migration value: this test pins the defect. +""" + +import hashlib + +import pytest +import torch + +from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import init_none_hash +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +BLOCK_SIZE = 4 +NUM_BLOCKS = 16 +PAGE_BYTES = 32 + + +def _make_capture_manager(boundary_capture: bool): + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=BLOCK_SIZE, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + boundary_capture=boundary_capture, + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=BLOCK_SIZE, + ) + # Two fake "layers" of per-block state: row index == pool block id. + kv_caches = [ + torch.zeros(NUM_BLOCKS, PAGE_BYTES, dtype=torch.uint8) + for _ in range(2) + ] + return manager, kv_caches + + +def _write_state(kv_caches, block_id, pattern: bytes) -> None: + for cache in kv_caches: + cache[block_id] = torch.frombuffer( + bytearray(pattern * (PAGE_BYTES // len(pattern))), + dtype=torch.uint8, + ).reshape(PAGE_BYTES) + + +def _digest(kv_caches, block_id) -> str: + h = hashlib.sha256() + for cache in kv_caches: + h.update(cache[block_id].numpy().tobytes()) + return h.hexdigest() + + +def _produce(manager, kv_caches, name, token_ids, state_at_b: bytes): + """Run one producer prompt through record -> promote -> copy rail, + simulating the worker: the producer's forward writes state@B into the + live column, the rail copies it, then the NEXT step's forward migrates + (clobbers) the column with the chunk-end state past B.""" + mgr = manager.coordinator.single_type_managers[1] + req = make_request(name, token_ids, BLOCK_SIZE, sha256) + req.skip_reading_prefix_cache = True # always produce, never resume + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks[name][0] + + # Producer forward: column holds state@B exactly (chunk ends at B). + _write_state(kv_caches, column.block_id, state_at_b) + + req.num_computed_tokens = 4 + manager.new_step_starts() # promote: rail + offer + ((group_id, offered_id, boundary),) = manager.take_boundary_state_offloads()[name] + assert group_id == 1 and boundary == 4 + copies, retained = manager.take_kv_cache_block_copies() + assert copies and offered_id == copies[0].dst_block_id + + # Worker applies the rail copy (before the next precopy/forward). + copy_kv_cache_blocks_inplace(kv_caches, NUM_BLOCKS, copies) + capture_digest = _digest(kv_caches, offered_id) + + # Next step's forward migrates the column: state@B is gone from it. + _write_state(kv_caches, column.block_id, b"\xee\xff") + column_digest_after = _digest(kv_caches, column.block_id) + + manager.block_pool.free_blocks(retained) + return offered_id, capture_digest, column_digest_after + + +def test_capture_digest_identical_across_prompts_and_not_the_column(): + """The mislabel test: frozen captures at B agree across prompts and + differ from each prompt's post-migration column bytes.""" + manager, kv_caches = _make_capture_manager(boundary_capture=True) + # Deterministic function of the shared 4-token prefix -> same bytes. + state_at_b = b"\x0a\x14\x1e\x28" + + id1, cap1, col1 = _produce(manager, kv_caches, "p1", [0, 1, 2, 3, 8, 9], state_at_b) + id2, cap2, col2 = _produce( + manager, kv_caches, "p2", [0, 1, 2, 3, 100, 101], state_at_b + ) + + assert id1 != id2 # dedicated block per producer + assert cap1 == cap2 # same prefix -> same state@B (key-identity holds) + # Neither capture equals either prompt's chunk-end column value after + # migration -- the pre-capture store would have shipped exactly those. + assert cap1 != col1 + assert cap2 != col2 + assert col1 == col2 # both columns clobbered identically past B + # The capture survived the column clobber: bytes are still state@B. + assert _digest(kv_caches, id1) == cap1 + assert _digest(kv_caches, id2) == cap2 + + +def test_legacy_live_column_store_mislabels_after_migration(): + """Without capture the hand-off names the live column: at store time the + deferred DMA would read the post-migration value under the state@B key. + This pins the defect the capture path removes (pre-change behavior).""" + manager, kv_caches = _make_capture_manager(boundary_capture=False) + mgr = manager.coordinator.single_type_managers[1] + + req = make_request("x", [0, 1, 2, 3, 8, 9], BLOCK_SIZE, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, 4, num_computed, computed_blocks) is not None + column = mgr.req_to_blocks["x"][0] + _write_state(kv_caches, column.block_id, b"\x0a\x14\x1e\x28") + + ((group_id, offered_id, boundary),) = manager.take_boundary_state_offloads()["x"] + assert (group_id, offered_id, boundary) == (1, column.block_id, 4) + # No copy rail on the legacy path. + assert manager.take_kv_cache_block_copies() == ([], []) + + state_at_b_digest = _digest(kv_caches, offered_id) + # The next forward migrates the column before the deferred DMA reads it. + _write_state(kv_caches, column.block_id, b"\xee\xff") + assert _digest(kv_caches, offered_id) != state_at_b_digest + + +def test_all_recurrent_groups_of_align_config_carry_capture(): + """Enablement completeness: every MambaSpec that reaches the kv-cache + config of an align-mode + OffloadingConnector deployment carries + boundary_capture=True, so no recurrent group's manager falls back to the + legacy live-column offer path under a boundary-colliding key. + + Guards the single-site enablement (MambaBase.get_kv_cache_spec): a future + spec-construction site or a get_kv_cache_spec override that rebuilds the + spec without super() would silently split capture coverage across groups. + """ + from types import SimpleNamespace + + from vllm.config.cache import CacheConfig + from vllm.model_executor.layers.mamba.abstract import MambaBase + + class _Layer(MambaBase): + def get_state_shape(self): + return ((1, 8), (1, 8, 4)) + + mamba_type = property(lambda self: None) + + def get_state_dtype(self): + return (torch.float32, torch.float32) + + def _spec_for(connector, offload_size, align, boundary_ckpt=False): + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + mamba_block_size=64, + mamba_page_size_padded=None, + mamba_cache_mode=align, + use_kda_recoverssm=False, + kv_offloading_size=offload_size, + ), + kv_transfer_config=connector, + num_speculative_tokens=0, + use_request_boundary_checkpoints=boundary_ckpt, + ) + return _Layer().get_kv_cache_spec(vllm_config) + + connector = SimpleNamespace( + has_connector=lambda name: name == "OffloadingConnector" + ) + # Align + complex offloading connector: capture on. + for mode in ("align",): + spec = _spec_for(connector, 1.0, mode) + assert spec.boundary_capture is True, f"{mode} group missed capture" + # Any gate missing -> capture off: boundary stores keep the live-column + # behavior and external hits stop at the last checkpoint. + assert _spec_for(connector, None, "align").boundary_capture is False + assert _spec_for(None, 1.0, "align").boundary_capture is False + assert _spec_for( + SimpleNamespace(has_connector=lambda name: False), 1.0, "align" + ).boundary_capture is False + assert _spec_for(connector, 1.0, "all").boundary_capture is False + assert ( + _spec_for(connector, 1.0, "align", boundary_ckpt=True).boundary_capture + is False + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index a368673660f0..882d76b4d87c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -4777,3 +4777,118 @@ def test_annotated_eagle_group_does_not_starve_coarse_sibling(self): False, ] assert self._roundtrip_served_tokens(scheduler) == 16 + + +def _make_capture_hybrid_scheduler(): + """OffloadingConnector scheduler over an align-Mamba hybrid whose mamba + group declares frozen boundary capture (the serve gate opens).""" + from dataclasses import replace + + vllm_config = _make_vllm_config() + vllm_config.speculative_config = None + kv_cache_config = _make_mamba_hybrid_kv_cache_config() + kv_cache_config = replace( + kv_cache_config, + kv_cache_groups=[ + kv_cache_config.kv_cache_groups[0], + KVCacheGroupSpec( + ["mamba_layer"], + replace( + kv_cache_config.kv_cache_groups[1].kv_cache_spec, + boundary_capture=True, + ), + ), + ], + ) + spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) + + +def test_external_hit_served_when_capture_enabled(): + """An align-Mamba deployment whose groups declare frozen boundary capture + extends external hits to the deepest boundary every group keys, instead of + zeroing them.""" + scheduler = _make_capture_hybrid_scheduler() + + request = _make_partial_tail_request(scheduler) + scheduler.manager.lookup.return_value = LookupResult.HIT + tokens, load_async = scheduler.get_num_new_matched_tokens(request, 0) + assert tokens > 0 and load_async + + +def test_hit_gate_extends_to_keyed_boundary(): + """With capture enabled the hit ends at the deepest boundary every group + keys. The mamba group participates with a one-chunk window over its + boundary key: with only its first chunk stored, the hit stops at one + chunk -- it is neither zeroed nor extended past the stored key. + """ + scheduler = _make_capture_hybrid_scheduler() + request = _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.num_locally_computed_tokens = 0 + req_status.update_offload_keys() + + # Store only chunk 0 of every group (attention keys h0, mamba state key + # at the first chunk end). + stored = {state.offload_keys[0] for state in req_status.group_states} + + def _lookup(key, req_context): + return LookupResult.HIT if key in stored else LookupResult.MISS + + scheduler.manager.lookup.side_effect = _lookup + tokens, load_async = scheduler.get_num_new_matched_tokens(request, 0) + assert tokens == 16 and load_async # one chunk of 16 tokens, not 0 + + +def test_capture_pin_released_on_store_ack_and_on_drop(): + """The connector releases the manager pin on the capture block when the + store job completes, and on every offer-skip path (unknown request, + save-window filter, dedup).""" + scheduler = _make_capture_hybrid_scheduler() + released: list[int] = [] + scheduler.bind_boundary_capture_releaser(released.append) + _make_partial_tail_request(scheduler) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + # Boundary beyond the save window is filtered -> pin dropped immediately. + jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 77, 10 ** 6)]}) + assert jobs == {} + assert released == [77] + + # Accepted offer -> job created; the ack releases the pin. + jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 88, 16)]}) + [job_id] = jobs + assert released == [77] # not before the workers report completion + output = KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: scheduler.config.num_workers} + ) + ) + req_status = scheduler._req_status["req"] + req_status.finished_signaled = True # keep _req_status alive for the ack + scheduler.update_connector_output(output) + assert released == [77, 88] + + # Unknown request drains its offers too. + jobs = scheduler._build_aligned_boundary_store_jobs({"ghost": [(1, 99, 16)]}) + assert jobs == {} + assert released == [77, 88, 99] + + +def test_capture_offer_skips_stores_release_pins(): + """prepare_store allocation failure and dedup (key already stored) both + drop the pin: no job will ever read the capture block.""" + scheduler = _make_capture_hybrid_scheduler() + released: list[int] = [] + scheduler.bind_boundary_capture_releaser(released.append) + _make_partial_tail_request(scheduler) + + scheduler.manager.prepare_store.return_value = None + assert scheduler._build_aligned_boundary_store_jobs({"req": [(1, 42, 16)]}) == {} + assert released == [42] + + scheduler.manager.prepare_store.return_value = generate_store_output([]) + assert scheduler._build_aligned_boundary_store_jobs({"req": [(1, 43, 16)]}) == {} + assert released == [42, 43] diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 366624e1d3b1..656408e30262 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -636,3 +636,231 @@ def test_register_kv_caches_uniform_type(backend): # opaque mapping rather than a certified, parallelism-agnostic one assert group_refs[0].mapping.parallelism_agnostic assert not group_refs[1].mapping.parallelism_agnostic + + +@pytest.mark.parametrize("backend", ATTN_BACKENDS) +def test_register_kv_caches_padded_attention_ref_carries_full_row(backend): + """The narrow-branch refs must transport the FULL padded row for + attention pages whose spec carries a padded tail. + + The QSA compressed-K tail lives in [unpadded, padded) of every padded + attention page and the consumer reads it through + compressed_block_table, so a ref sized to the unpadded page never + transports it. + + Also pins the non-regression half: an attention spec WITHOUT a padded + tail keeps the identical ref width (page == unpadded), and the CPU + tensor row of the padded group is the padded page (the store/load DMA + sizes now match it). + """ + from vllm.v1.worker.utils import AttentionGroup + + backend_cls = AttentionBackendEnum[backend].get_class() + + TAIL_BYTES = 4096 + plain_layer = "model.layers.0.self_attn" + padded_layer = "model.layers.1.self_attn" + plain_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + padded_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + page_size_padded=plain_spec.page_size_bytes + TAIL_BYTES, + ) + assert padded_spec.unpadded_page_size_bytes == plain_spec.page_size_bytes + assert padded_spec.page_size_bytes == plain_spec.page_size_bytes + TAIL_BYTES + + # Two independent groups, each with its own dense per-block run over + # one shared backing allocation (the model runner's convention). + window = plain_spec.page_size_bytes + padded_spec.page_size_bytes + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor( + size=window * NUM_BLOCKS, + layers=[plain_layer], + layer_stride=plain_spec.page_size_bytes * NUM_BLOCKS, + block_stride=plain_spec.page_size_bytes, + ), + KVCacheTensor( + size=window * NUM_BLOCKS, + layers=[padded_layer], + layer_stride=padded_spec.page_size_bytes * NUM_BLOCKS, + block_stride=padded_spec.page_size_bytes, + offset=plain_spec.page_size_bytes * NUM_BLOCKS, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=[plain_layer], kv_cache_spec=plain_spec + ), + KVCacheGroupSpec( + layer_names=[padded_layer], kv_cache_spec=padded_spec + ), + ], + ) + attn_groups = [ + [ + AttentionGroup( + backend=backend_cls, + layer_names=[plain_layer], + kv_cache_spec=plain_spec, + kv_cache_group_id=0, + ), + ], + [ + AttentionGroup( + backend=backend_cls, + layer_names=[padded_layer], + kv_cache_spec=padded_spec, + kv_cache_group_id=1, + ), + ], + ] + kv_caches = _allocate_kv_caches( + kv_cache_config, + attn_groups, + device=torch.device(f"{DEVICE_TYPE}:0"), + ) + + worker, spec = _make_worker(kv_cache_config) + worker.register_kv_caches(kv_caches) + + canonical = spec.get_worker.call_args[0][0] + assert isinstance(canonical, CanonicalKVCaches) + + plain_refs = canonical.group_data_refs[0] + padded_refs = canonical.group_data_refs[1] + # Non-regression: no padded tail -> ref width unchanged (== page == + # unpadded). + for ref in plain_refs: + assert ref.page_size_bytes == plain_spec.page_size_bytes + # The fix: the padded group's ref transports the whole row. + for ref in padded_refs: + assert ref.page_size_bytes == padded_spec.page_size_bytes + # The mapping still certifies the UNPADDED content width, so the + # [unpadded, padded) tail stays distinguishable from the payload. + assert ref.mapping is not None + assert ref.mapping.local_page_size_bytes == ( + padded_spec.unpadded_page_size_bytes + ) + # The block tensor row is the padded page on both sides. + for block_tensor in canonical.tensors: + assert block_tensor.tensor.shape[1] in ( + plain_spec.page_size_bytes, + padded_spec.page_size_bytes, + ) + + +@pytest.mark.skipif( + DEVICE_TYPE != "cuda", reason="Real GPU DMA + pinned CPU tier required" +) +def test_padded_attention_tail_survives_store_load_roundtrip(): + """End-to-end: the compressed-K tail bytes land in the CPU tier on + store and come back on load. + + Drives the REAL register_kv_caches path (so the narrow refs are what + the worker transports) through a real CPUOffloadingWorker: fill the + padded GPU row with distinct bytes, store, wipe the GPU row, load + back, and require the FULL row -- [unpadded, padded) tail included + -- to match.""" + from vllm.v1.worker.utils import AttentionGroup + from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec + from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker + + backend_cls = AttentionBackendEnum["FLASH_ATTN"].get_class() + + TAIL_BYTES = 4096 + layer = "model.layers.0.self_attn" + spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + page_size_padded=( + 2 * BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * get_dtype_size(DTYPE) + + TAIL_BYTES + ), + ) + unpadded = spec.unpadded_page_size_bytes + padded = spec.page_size_bytes + assert padded - unpadded == TAIL_BYTES + + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor( + size=padded * NUM_BLOCKS, + layers=[layer], + layer_stride=padded * NUM_BLOCKS, + block_stride=padded, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=[layer], kv_cache_spec=spec) + ], + ) + attn_groups = [ + [ + AttentionGroup( + backend=backend_cls, + layer_names=[layer], + kv_cache_spec=spec, + kv_cache_group_id=0, + ), + ] + ] + kv_caches = _allocate_kv_caches( + kv_cache_config, attn_groups, device=torch.device("cuda:0") + ) + + worker, spec_mock = _make_worker(kv_cache_config) + worker.register_kv_caches(kv_caches) + canonical = spec_mock.get_worker.call_args[0][0] + # The transport width the roundtrip below must observe. + assert canonical.group_data_refs[0][0].page_size_bytes == padded + + offload_worker = CPUOffloadingWorker( + kv_caches=canonical, + blocks_per_chunk=1, + num_cpu_chunks=4, + ) + try: + gpu_row = canonical.tensors[0].tensor[3] # GPU block 3 + # Distinct, position-dependent bytes so a zero tail is detectable. + pattern = ( + torch.arange(padded, dtype=torch.int64) % 251 + ).to(torch.int8) + gpu_row.copy_(pattern.to(gpu_row.device)) + + assert offload_worker.submit_store( + 1, + GPULoadStoreSpec([3], group_sizes=(1,), block_indices=(0,)), + CPULoadStoreSpec([1]), + ) + offload_worker.wait({1}) + + # Wipe the GPU row; only the CPU tier can bring the tail back. + gpu_row.zero_() + torch.cuda.synchronize() + + assert offload_worker.submit_load( + 2, + CPULoadStoreSpec([1]), + GPULoadStoreSpec([3], group_sizes=(1,), block_indices=(0,)), + ) + offload_worker.wait({2}) + + reloaded = gpu_row.cpu() + # Regression guard: [unpadded, padded) must leave the GPU on store. + tail = reloaded[unpadded:padded] + assert tail.abs().sum().item() != 0, "tail bytes were never transported" + torch.testing.assert_close(reloaded, pattern.cpu()) + finally: + offload_worker.shutdown() diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index a79f480ff551..1a783e59b4f0 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -1156,3 +1156,28 @@ def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: assert len(received) == 1 assert received[0][0] == keys assert received[0][1] is ctx + + +def test_second_store_of_in_flight_key_is_deduped(): + """A key registered by prepare_store but not yet completed must dedup a + second store of the same key: the in-flight chunk is already claimed, so + a concurrent producer (two requests sharing a boundary prefix) must not + allocate a second chunk or overwrite the pending one. This is the + first-writer-wins guarantee the boundary-store path relies on for + same-key producers.""" + manager = make_cpu_manager(num_chunks=4) + + out1 = manager.prepare_store(to_keys([7]), _EMPTY_REQ_CTX) + assert out1 is not None and out1.keys_to_store == to_keys([7]) + + # Same key offered again while the first store is still in flight. + out2 = manager.prepare_store(to_keys([7]), _EMPTY_REQ_CTX) + assert out2 is not None + assert out2.keys_to_store == [], "in-flight key must dedup, not overwrite" + assert manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING + + # The first store completing does not disturb the dedup decision. + manager.complete_store(to_keys([7]), _EMPTY_REQ_CTX) + assert manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT + out3 = manager.prepare_store(to_keys([7]), _EMPTY_REQ_CTX) + assert out3 is not None and out3.keys_to_store == [] diff --git a/tests/v1/worker/test_qwen4_exp_qsa_resume_fresh_gating.py b/tests/v1/worker/test_qwen4_exp_qsa_resume_fresh_gating.py new file mode 100644 index 000000000000..911dd6b88b42 --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_qsa_resume_fresh_gating.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The QSA pool fresh-reset must key on a real selector-pool restore. + +The QSA selector pools (raw ring / logical tags / RoPE / anchor) are per-slot +module buffers, reset when ``qsa_state_is_fresh_gpu`` is set for a slot +(b12x_qsa.py). Only the boundary-checkpoint rail restores them: it collects +``get_recurrent_checkpoint_tensors`` and copies the slot's rows, so it alone +may clear the flag -- and it does so itself, or the next forward would +re-zero what it just copied. + +A connector external-resume restores KV pages, not these pools. The ring is +a bounded speculative-interval window (``raw_ring_capacity`` rows); the +committed prefix's selector coverage is the compressed-K page tail, which the +padded connector transport does restore. So the reset discards no prefix +state, and it re-derives the anchor as ``first_position - accepted`` = +``prefix - 1``, the value the checkpoint path seeds by hand. A resumed +prefix without a checkpoint therefore keeps the fresh flag, which clears the +prior slot owner's rows. + +Regression shape: ``add_request`` keyed the clear on ``num_computed_tokens > +0``, so a recycled slot still holding another request's rows was read as if it +were this request's committed selector state. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState + +MAX_REQS = 8 + + +def _bare_qsa_state() -> Qwen4ExpModelState: + """A Qwen4ExpModelState with only the fields add_request touches.""" + state = Qwen4ExpModelState.__new__(Qwen4ExpModelState) + state.max_num_reqs = MAX_REQS + state.uses_qsa = True + state.qsa_state_slot_ids = torch.arange(MAX_REQS, dtype=torch.int32) + state._qsa_default_slot_ids = state.qsa_state_slot_ids.clone() + state.qsa_state_is_fresh = torch.ones(MAX_REQS, dtype=torch.bool) + state.qsa_state_is_fresh_gpu = torch.ones(MAX_REQS, dtype=torch.bool) + state.qsa_num_accepted_tokens = torch.ones(MAX_REQS, dtype=torch.int32) + state.qsa_committed_num_accepted_tokens_gpu = torch.ones( + MAX_REQS, dtype=torch.int32 + ) + state._qsa_draft_is_prefilling = torch.zeros(MAX_REQS, dtype=torch.bool) + state._qsa_draft_is_prefilling_gpu = torch.zeros(MAX_REQS, dtype=torch.bool) + return state + + +class _FakeQSALayer: + """Minimal stand-in for the QSA attention module surface add_request + consumes: the anchor pool + the setter.""" + + def __init__(self, max_seqs: int = MAX_REQS): + self._raw_interval_start_positions = torch.full( + (max_seqs,), -1, dtype=torch.int64 + ) + + def set_recurrent_checkpoint_anchor(self, slot: int, anchor) -> None: + self._raw_interval_start_positions[slot] = anchor + + +def _attach_fake_layers(state: Qwen4ExpModelState) -> list[_FakeQSALayer]: + layers = [_FakeQSALayer(), _FakeQSALayer()] + state.model = SimpleNamespace( + modules=lambda: iter(layers), + ) + return layers + + +def _new_req( + num_computed_tokens: int, + boundary_checkpoint=None, + boundary_checkpoint_blocks=None, +) -> SimpleNamespace: + return SimpleNamespace( + num_computed_tokens=num_computed_tokens, + boundary_checkpoint=boundary_checkpoint, + boundary_checkpoint_blocks=boundary_checkpoint_blocks, + ) + + +def test_cold_add_keeps_fresh_true(monkeypatch: pytest.MonkeyPatch) -> None: + """Control: a cold add (no restored prefix) owns a recycled slot.""" + state = _bare_qsa_state() + layers = _attach_fake_layers(state) + monkeypatch.setattr( + MambaHybridModelState, "add_request", lambda self, i, r: None + ) + + state.add_request(3, _new_req(0)) + + assert bool(state.qsa_state_is_fresh_gpu[3]) + assert int(state.qsa_committed_num_accepted_tokens_gpu[3]) == 1 + # No anchor seeding for a cold add. + for layer in layers: + assert int(layer._raw_interval_start_positions[3]) == -1 + + +def test_external_resume_add_keeps_fresh_true( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A connector external-resume restores KV pages only. Its slot may be + recycled from another request, so the pools must be reset; the reset + re-derives the same anchor a checkpoint restore would have seeded.""" + state = _bare_qsa_state() + state.qsa_state_is_fresh_gpu[5].fill_(False) # prior owner cleared + layers = _attach_fake_layers(state) + monkeypatch.setattr( + MambaHybridModelState, "add_request", lambda self, i, r: None + ) + + state.add_request(5, _new_req(9024)) + + assert bool(state.qsa_state_is_fresh_gpu[5]) + assert int(state.qsa_committed_num_accepted_tokens_gpu[5]) == 1 + # No anchor seeding: the fresh reset owns that write. + for layer in layers: + assert int(layer._raw_interval_start_positions[5]) == -1 + + +def test_local_hit_add_keeps_fresh_true( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A local prefix-cache hit without a checkpoint has the same gap: the + producer's raw rows lived in the producer's slot, not in the restored + pages, so the resumed slot stays fresh.""" + state = _bare_qsa_state() + _attach_fake_layers(state) + monkeypatch.setattr( + MambaHybridModelState, "add_request", lambda self, i, r: None + ) + + state.add_request(0, _new_req(6016)) + + assert bool(state.qsa_state_is_fresh_gpu[0]) + + +def test_boundary_checkpoint_add_clears_fresh_and_seeds_anchor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The one resume with selector-pool continuity: the checkpoint restore + writes this slot's pools, so the flag stays clear and the anchor seeds to + the checkpoint boundary - 1, the value the skipped fresh-reset would have + produced (the rail seeds the same value; this keeps the pre-forward + state right on its own).""" + state = _bare_qsa_state() + state.qsa_state_is_fresh_gpu[5].fill_(True) + layers = _attach_fake_layers(state) + monkeypatch.setattr( + MambaHybridModelState, "add_request", lambda self, i, r: None + ) + + checkpoint = SimpleNamespace(num_tokens=9024) + state.add_request( + 5, + _new_req( + 9024, + boundary_checkpoint=checkpoint, + boundary_checkpoint_blocks=((1,),), + ), + ) + + assert not bool(state.qsa_state_is_fresh_gpu[5]) + assert int(state.qsa_committed_num_accepted_tokens_gpu[5]) == 1 + for layer in layers: + assert int(layer._raw_interval_start_positions[5]) == 9023 + + +def test_checkpoint_without_allocation_keeps_fresh_true( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The rail restores only when both fields are present: without an + allocation its add_request returns before the pool copy, so the gate must + not clear a flag the restore will not re-clear.""" + state = _bare_qsa_state() + state.qsa_state_is_fresh_gpu[5].fill_(False) + _attach_fake_layers(state) + monkeypatch.setattr( + MambaHybridModelState, "add_request", lambda self, i, r: None + ) + + state.add_request( + 5, + _new_req(9024, boundary_checkpoint=SimpleNamespace(num_tokens=9024)), + ) + + assert bool(state.qsa_state_is_fresh_gpu[5]) + + +def test_fresh_hook_returns_the_gpu_flag_tensor() -> None: + """The boundary-restore path consumes the flag through the interface + hook; it must be the live GPU tensor, not a copy.""" + state = _bare_qsa_state() + assert state.get_recurrent_checkpoint_fresh() is state.qsa_state_is_fresh_gpu + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_boundary_restore_clears_fresh_flag() -> None: + """The restore branch of BoundaryCheckpointState.add_request must clear + the fresh flag. Drives the real method with a stub state carrying only + the attributes the restore branch reads.""" + from vllm.v1.worker.gpu.boundary_checkpoint import BoundaryCheckpointState + + device = torch.device("cuda") + slot = 2 + + fresh_gpu = torch.ones(MAX_REQS, dtype=torch.bool, device=device) + acceptance_gpu = torch.ones(MAX_REQS, dtype=torch.int32, device=device) + model_state = SimpleNamespace( + get_recurrent_checkpoint_fresh=lambda: fresh_gpu, + get_recurrent_checkpoint_acceptance=lambda: acceptance_gpu, + model_config=SimpleNamespace(max_model_len=16384), + ) + + # One auxiliary tensor restored from pool block 1 into the slot row. + aux = torch.zeros(MAX_REQS, 8, dtype=torch.uint8, device=device) + pool = torch.zeros(4, 64, dtype=torch.uint8, device=device) + pool[1, 0:8] = 7 + aux_metadata = torch.tensor( + [[aux.data_ptr(), aux.stride(0) * aux.element_size(), 8, 0]], + dtype=torch.int64, + device=device, + ) + + checkpoint = SimpleNamespace( + num_tokens=9024, + auxiliary_block_ids=(1,), + ) + st = BoundaryCheckpointState.__new__(BoundaryCheckpointState) + st.device = device + st.model_state = model_state + st.seen = torch.zeros(MAX_REQS, 4, dtype=torch.int32, device=device) + st.metadata = torch.zeros(MAX_REQS, 7, dtype=torch.int32, device=device) + st.blocks = torch.zeros(MAX_REQS, 4, 8, dtype=torch.int32, device=device) + st.mamba_blocks = torch.zeros(MAX_REQS, 4, 1, dtype=torch.int32, device=device) + st.stop_tokens = torch.zeros(MAX_REQS, 128, dtype=torch.int32, device=device) + st.mamba_group_ids = [0] + st.auxiliary_metadata = aux_metadata + st.pool = pool + st.target_modules = [] + + st.add_request( + slot, + SimpleNamespace( + boundary_checkpoint_blocks=((1, 0, 0, 0, 0, 0, 0, 0),), + sampling_params=SimpleNamespace( + max_tokens=64, min_tokens=1, stop_token_ids=[], eos_token_id=None + ), + prompt_len=9059, + recurrent_instruction_boundary=None, + recurrent_prefill_tail_boundary=None, + boundary_checkpoint=checkpoint, + ), + ) + + assert not bool(fresh_gpu[slot]) + # The restore still performs its other writes (acceptance reset). + assert int(acceptance_gpu[slot]) == 1 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index a7b6c2aa14f8..b5666641a8b0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -477,6 +477,16 @@ def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: self._kv_cache_manager = kv_cache_manager self.bind_gpu_block_pool(kv_cache_manager.block_pool) + def bind_boundary_capture_releaser(self, releaser: Callable[[int], None]) -> None: + """Bind the cache manager's capture-pin releaser. + + A connector whose boundary stores source frozen capture blocks must + release the manager pin when a store job completes or the offer is + dropped. Connectors that never source capture blocks ignore it (the + manager's pin dict stays empty, so a missed release cannot leak). + """ + return None + @classmethod def supports_request_boundary_checkpoints(cls, config: "VllmConfig") -> bool: """Whether configuration selects a complete external checkpoint adapter.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index d291d3852804..606b8322bf31 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -261,6 +261,13 @@ def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: for c in self._connectors: c.bind_gpu_block_pool(gpu_block_pool) + def bind_boundary_capture_releaser(self, releaser: Callable[[int], None]) -> None: + # OffloadingConnector as a MultiConnector child still owns capture + # pins (enablement checks has_connector, which sees children): every + # child gets the releaser; the base no-op absorbs the rest. + for c in self._connectors: + c.bind_boundary_capture_releaser(releaser) + # We must override the base class method here because we need to bind # the metadata to each connector in the order of the connectors in the # MultiKVConnectorMetadata. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 521f8b4a15c6..4ce18152e6f1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field, replace from itertools import chain, islice from typing import Any, NamedTuple @@ -97,7 +97,11 @@ class GroupOffloadConfig(NamedTuple): # Cached manager class for this group's KV cache spec, resolved at init time manager_cls: type[SingleTypeKVCacheManager] # Partial-tail data for this group comes from the scheduler's CoW hand-off - # rather than the request block table. + # rather than the request block table. Invariant (asserted in from_spec and + # at capture record time): a frozen capture boundary is a state-block + # multiple, and partial-tail configs force chunk == block, so the aligned + # builder alone owns capture offers and its release-before-DMA branch never + # abandons a pinned source to a partial-tail claim. requires_cow_source: bool = False # True for EAGLE/MTP draft-model attention groups. The trailing chunk # of these groups is volatile and lacks a stable hash, so it must @@ -314,6 +318,18 @@ def from_spec( and not any(config.is_eagle_group for config in kv_group_configs) and vllm_config.parallel_config.decode_context_parallel_size == 1 ) + # Drift guard for the capture-ownership invariant (see + # GroupOffloadConfig.requires_cow_source): capture boundaries are state + # -block multiples and the partial-tail builder only claims non-block + # multiples, so the aligned builder's release-before-DMA branch can + # never strand a pinned capture source to a partial-tail claim. + assert not supports_partial_tail or all( + config.tokens_per_chunk == config.tokens_per_block + for config in kv_group_configs + ), ( + "partial tails require one physical block per offload chunk so " + "frozen capture boundaries stay out of the partial-tail builder" + ) if retention_interval is not None: if retention_interval < 0: @@ -601,6 +617,19 @@ def _sliding_window_sort_key(i: int) -> int: for config in self.config.kv_group_configs if config.requires_cow_source ) + # Align-mode Mamba boundary stores source FROZEN capture blocks + # (single_type_kv_cache_manager boundary capture): a dedicated pool + # block written exactly once by a post-forward copy at the step whose + # forward ended at the retention boundary B, then fenced like any + # other store source. With capture enabled on every recurrent group, + # the external hit extends to the largest B where all groups key -- + # which is exactly what the per-group convergence in + # `_lookup_complete_chunks` computes (mamba groups participate with a + # one-chunk window over their boundary state key). No half-restore. + # Without frozen capture the boundary store still sources the request's + # live running column, which the deferred DMA reads after the next + # forward clobbers it: align-Mamba deployments must enable + # boundary_capture on every recurrent group. self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -618,6 +647,9 @@ def _sliding_window_sort_key(i: int) -> int: # Threshold value for stale jobs. All job ids >= _stale_job_threshold are # active jobs. self._stale_job_threshold: int = 0 + # job_id -> capture-pinned source blocks of store jobs orphaned by a + # cache reset; released when the worker's stale ack drains them (I4). + self._stale_capture_blocks: dict[int, list[int]] = {} self._jobs: dict[int, TransferJobStatus] = {} # block_id -> pending store job_ids. Used to track jobs that needs @@ -626,9 +658,20 @@ def _sliding_window_sort_key(i: int) -> int: # protected by their ref_cnt) and for sliding window blocks (which can # be freed before a request finishes). self._block_id_to_pending_jobs: dict[int, set[int]] = {} + # Set by the scheduler after the KV cache manager exists: releases the + # manager pin on a frozen capture block once its store completed or the + # offer was dropped (I4). No-op for block ids the manager never pinned. + self._release_boundary_capture: Callable[[int], None] | None = None self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + def bind_boundary_capture_releaser(self, releaser: Callable[[int], None]) -> None: + self._release_boundary_capture = releaser + + def _release_capture(self, block_id: int) -> None: + if self._release_boundary_capture is not None: + self._release_boundary_capture(block_id) + def _maybe_observe_lookup_async_delay( self, req_status: RequestOffloadState ) -> None: @@ -1300,12 +1343,17 @@ def _build_aligned_boundary_store_jobs( for req_id, entries in handoffs.items(): req_status = self._req_status.get(req_id) if req_status is None: + # Offered after the request's last job drained: nothing will + # store these blocks, so drop their capture pins now. + for _, block_id, _ in entries: + self._release_capture(block_id) continue req = req_status.req max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens) for group_idx, block_id, boundary in entries: config_idx = config_idx_by_group.get(group_idx) if config_idx is None: + self._release_capture(block_id) continue group_config = self.config.kv_group_configs[config_idx] if ( @@ -1313,6 +1361,7 @@ def _build_aligned_boundary_store_jobs( or boundary > max_boundary or boundary % group_config.tokens_per_chunk != 0 ): + self._release_capture(block_id) continue key = self._make_boundary_key(req, group_idx, boundary) @@ -1321,8 +1370,12 @@ def _build_aligned_boundary_store_jobs( self._connector_stats.increase_counter( _ConnectorMetricName.ALLOCATION_FAILURE ) + self._release_capture(block_id) continue if not store_output.keys_to_store: + # Deduped: the CPU tier already holds this key (an + # earlier producer stored it). No DMA is submitted. + self._release_capture(block_id) continue job_id = self._generate_job_id() @@ -1408,6 +1461,10 @@ def _build_partial_tail_store_jobs( and block_idx >= len(group_states[group.group_idx].block_ids) for group in self.config.kv_group_configs ): + # Abandoned hand-off: drop capture pins among the CoW sources + # (no-op for table blocks), mirroring the aligned builder. + for block_id in cow_blocks.values(): + self._release_capture(block_id) continue keys = [ self._make_boundary_key(req, group.group_idx, boundary) @@ -1426,8 +1483,14 @@ def _build_partial_tail_store_jobs( self._connector_stats.increase_counter( _ConnectorMetricName.ALLOCATION_FAILURE ) + for block_id in cow_blocks.values(): + self._release_capture(block_id) continue if not store_output.keys_to_store: + # Deduped: the CPU tier already holds these keys; no DMA is + # submitted, so any capture-sourced CoW pin drops here. + for block_id in cow_blocks.values(): + self._release_capture(block_id) continue for group_config, key in zip(self.config.kv_group_configs, keys): @@ -1930,6 +1993,18 @@ def update_connector_output(self, connector_output: KVConnectorOutput): job_id, self._stale_job_threshold, ) + # The release loop below never runs for this orphan: drain its + # recorded capture pins, skipping blocks a live (post-reset) + # job now fences — block ids are reused across resets. + stale_blocks = self._stale_capture_blocks.pop(job_id, ()) + if stale_blocks: + live_fenced: set[int] = set() + for job in self._jobs.values(): + live_fenced.update(job.fenced_block_ids or ()) + live_fenced.update(job.deferred_fence_block_ids or ()) + for block_id in stale_blocks: + if block_id not in live_fenced: + self._release_capture(block_id) continue job_status = self._jobs[job_id] job_status.pending_count -= count @@ -1940,6 +2015,11 @@ def update_connector_output(self, connector_output: KVConnectorOutput): req_status = self._req_status[job_status.req_id] if job_status.is_store: self.manager.complete_store(job_status.keys, req_status.req_context) + # Every worker's D2H read of the source blocks is done (I4): + # release the manager pin on any frozen capture sources. + # No-op for block ids the manager never pinned. + for block_id in job_status.fenced_block_ids or (): + self._release_capture(block_id) else: self.manager.complete_load(job_status.keys, req_status.req_context) if self._chunks_being_loaded: @@ -2058,6 +2138,15 @@ def reset_cache(self) -> None: # Discard jobs and save job_counter to be able to discard worker responses self._stale_job_threshold = self._job_counter + # Stale acks skip the release loop; record orphaned store fences so a + # reset that bypassed the manager's release-all can't leak pins (I4). + self._stale_capture_blocks.update( + { + job_id: list(job.fenced_block_ids) + for job_id, job in self._jobs.items() + if job.is_store and job.fenced_block_ids + } + ) self._jobs.clear() self._block_id_to_pending_jobs.clear() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 14e826e5525c..8649065f0d93 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, + KVCacheSpec, MambaSpec, UniformTypeKVCacheSpecs, group_kernel_blocks, @@ -82,6 +83,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} + # layer_name -> KV cache spec of the layer + layer_specs: dict[str, KVCacheSpec] = {} for kv_cache_group in selected_groups: assert not kv_cache_group.host_resident num_blocks = kv_cache_config.num_blocks @@ -95,6 +98,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): layer_kv_cache_spec = per_layer_specs.get( layer_name, group_kv_cache_spec ) + layer_specs[layer_name] = layer_kv_cache_spec ref = group_kernel_blocks(kv_caches[layer_name], num_blocks) page = layer_kv_cache_spec.page_size_bytes elem_size = ref.element_size() @@ -194,7 +198,15 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): block_data_refs[layer_name].append( CanonicalKVCacheRef( tensor_idx=curr_tensor_idx, - page_size_bytes=(unpadded_page_size_bytes[layer_name]), + # Attention refs transport the FULL padded page: + # the compressed-K tail lives in + # [unpadded, padded) and the consumer reads it. + # Mamba refs keep the unpadded width. + page_size_bytes=( + page_size_bytes[layer_name] + if isinstance(layer_specs[layer_name], AttentionSpec) + else unpadded_page_size_bytes[layer_name] + ), mapping=mapping, ) ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 73570881753d..2eb0f56118b6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable +from collections.abc import Callable, Iterable from functools import cached_property from typing import Any @@ -98,6 +98,10 @@ def __init__( spec, vllm_config, kv_cache_config ) + def bind_boundary_capture_releaser(self, releaser: "Callable[[int], None]") -> None: + if self.connector_scheduler is not None: + self.connector_scheduler.bind_boundary_capture_releaser(releaser) + def shutdown(self) -> None: if self.connector_worker is not None: self.connector_worker.shutdown() diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 6e7455a499a5..33e3802f3644 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -72,6 +72,26 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: mamba_block_size = vllm_config.cache_config.mamba_block_size assert mamba_block_size is not None page_size_padded = vllm_config.cache_config.mamba_page_size_padded + cache_config = vllm_config.cache_config + kv_transfer_config = vllm_config.kv_transfer_config + # Frozen boundary capture: the align-mode column-migration defect + # (a deferred connector boundary store sourcing a live column the + # next forward clobbers) only exists when the complex offloading + # connector moves state off-GPU. With it enabled, the manager copies + # each retention-grid crossing into a dedicated single-writer pool + # block before the column advances, and that block -- never the + # column -- is the store source. The copy rail is backend-agnostic + # (post-forward D2D of the whole page), so no scan rework is needed. + boundary_capture = ( + cache_config.mamba_cache_mode == "align" + and cache_config.kv_offloading_size is not None + and kv_transfer_config is not None + and kv_transfer_config.has_connector("OffloadingConnector") + # Atomic request-boundary checkpoints own the same columns with + # their own capture path; the two features are mutually exclusive + # (kv_offloading_size already asserts it) -- never double-capture. + and not vllm_config.use_request_boundary_checkpoints + ) return MambaSpec( shapes=tuple(self.get_state_shape()), dtypes=self.get_state_dtype(), @@ -79,12 +99,13 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: page_size_padded=page_size_padded, mamba_type=self.mamba_type, tp_replicated=self.is_kv_cache_tp_replicated, - mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, + mamba_cache_mode=cache_config.mamba_cache_mode, + boundary_capture=boundary_capture, # RecoverSSM verifies the whole window off one checkpoint, so it # never writes the baseline's per-draft-token state slots. num_speculative_blocks=( 0 - if vllm_config.cache_config.use_kda_recoverssm + if cache_config.use_kda_recoverssm else vllm_config.num_speculative_tokens ), ) diff --git a/vllm/models/glm5next/model_state.py b/vllm/models/glm5next/model_state.py index 2aea3d27a514..c50e21058087 100644 --- a/vllm/models/glm5next/model_state.py +++ b/vllm/models/glm5next/model_state.py @@ -139,7 +139,20 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: if self.uses_pooled_selector: # The scheduler may recycle this request-state slot while selector # raw-ring tags and its interval anchor still belong to the prior owner. - self.selector_state_is_fresh_gpu[req_index].fill_(True) + # Only a boundary-checkpoint restore carries the producer's committed + # selector pools into this slot; the OffloadingConnector restores KV + # pages, never these per-slot buffers. A restored prefix without a + # checkpoint still owns a recycled slot, so it keeps the fresh reset. + # The condition mirrors BoundaryCheckpointState.add_request, whose + # restore branch -- the only writer of these pools -- needs both + # fields before it clears the flag. + restored_selector_state = ( + new_req_data.boundary_checkpoint is not None + and new_req_data.boundary_checkpoint_blocks is not None + ) + self.selector_state_is_fresh_gpu[req_index].fill_( + not restored_selector_state + ) self.selector_committed_num_accepted_tokens_gpu[req_index].fill_(1) def reset_kv_cache_state(self) -> None: @@ -157,6 +170,9 @@ def get_recurrent_checkpoint_tensors(self) -> tuple[torch.Tensor, ...]: def get_recurrent_checkpoint_acceptance(self) -> torch.Tensor: return self.selector_committed_num_accepted_tokens_gpu + def get_recurrent_checkpoint_fresh(self) -> torch.Tensor: + return self.selector_state_is_fresh_gpu + def _prepare_selector_state( self, input_batch: InputBatch, diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index 183468978fea..e727721dc9c1 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -177,8 +177,52 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: # RoPE, and anchor pools still contain the prior owner's data. The # flag remains set through the complete next model forward so every # QSA layer independently resets its slot. - self.qsa_state_is_fresh_gpu[req_index].fill_(True) + # + # Only the boundary-checkpoint restore carries the producer's + # committed selector pools into this slot: the complex + # OffloadingConnector restores KV pages (main/compressed cache, + # align-Mamba columns), never these per-slot module buffers. A + # resumed prefix without a checkpoint -- local hit or connector + # external-resume -- still owns a recycled slot, so it keeps the + # fresh reset. The condition mirrors the rail's restore branch + # (BoundaryCheckpointState.add_request), the only pool writer. + checkpoint = new_req_data.boundary_checkpoint + restored_selector_state = ( + checkpoint is not None + and new_req_data.boundary_checkpoint_blocks is not None + ) + self.qsa_state_is_fresh_gpu[req_index].fill_(not restored_selector_state) self.qsa_committed_num_accepted_tokens_gpu[req_index].fill_(1) + if restored_selector_state: + # The fresh-reset's anchor formula (first_position - + # accepted, in the QSA reset kernel) is skipped with the + # flag down, so seed the value it would have produced: + # prefix_len - 1. The commit kernel overwrites it during + # the first forward (the anchor self-heals), so this keeps + # the pre-forward selector state consistent, not just the + # post-forward one. + anchor = checkpoint.num_tokens - 1 + for module in self._qsa_state_modules(): + module.set_recurrent_checkpoint_anchor(req_index, anchor) + + def _qsa_state_modules(self) -> tuple[Any, ...]: + """QSA layer modules exposing the per-request selector pools. + + Mirrors ``BoundaryCheckpointState.target_modules``: the raw ring / + logical-tag / RoPE / anchor pools are per-layer module buffers + indexed by request slot, so seeding or reading one request's + selector state goes through every module that publishes the + recurrent-checkpoint accessors. + """ + modules = self.__dict__.get("_qsa_state_module_cache") + if modules is None: + modules = tuple( + module + for module in self.model.modules() + if hasattr(module, "set_recurrent_checkpoint_anchor") + ) + self._qsa_state_module_cache = modules + return modules def get_recurrent_checkpoint_tensors(self) -> tuple[torch.Tensor, ...]: return ( @@ -189,6 +233,9 @@ def get_recurrent_checkpoint_tensors(self) -> tuple[torch.Tensor, ...]: def get_recurrent_checkpoint_acceptance(self) -> torch.Tensor: return self.qsa_committed_num_accepted_tokens_gpu + def get_recurrent_checkpoint_fresh(self) -> torch.Tensor: + return self.qsa_state_is_fresh_gpu + def _prepare_qsa_state( self, input_batch: InputBatch, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 84812b69baed..4ba5e0eb46b9 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -1095,6 +1095,10 @@ def reset_prefix_cache(self) -> bool: bool: True if the prefix cache is successfully reset, False otherwise. """ + # Capture pins are not request-owned; the pool reset releases them + # even while stores are still pending. + for mgr in self.coordinator.single_type_managers: + mgr.release_all_boundary_captures() if not self.coordinator.reset_prefix_cache(): return False if self.log_stats: @@ -1349,6 +1353,16 @@ def take_boundary_state_offloads( ) return offloads + def release_boundary_capture(self, block_id: int) -> None: + """Release a frozen capture block's manager pin (store ack / drop). + + Only the Mamba manager that pinned ``block_id`` acts; the call is a + no-op for every other manager and for block ids never pinned, so a + connector may pass it fenced ids from any store kind. + """ + for mgr in self.coordinator.single_type_managers: + mgr.release_boundary_capture(block_id) + def finalize_partial_tail_offloads( self, request: Request ) -> list[tuple[int, int, int]]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 6e5ea6d15b33..29ac17b28f57 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -368,6 +368,12 @@ def __init__( if self.connector is not None: self.connector.bind_kv_cache_manager(self.kv_cache_manager) self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + # Frozen boundary capture: the connector releases the manager pin + # on a capture block when its store DMA acks or the offer is + # dropped (I4). No-op for block ids the manager never pinned. + self.connector.bind_boundary_capture_releaser( + self.kv_cache_manager.release_boundary_capture + ) if self.kv_cache_manager.boundary_checkpoints is not None: self.connector.bind_boundary_checkpoint_cache(self.kv_cache_manager) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index bf53b82ffe4e..968b7b67fd68 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -962,6 +962,14 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: def new_step_starts(self) -> None: return None + def release_boundary_capture(self, block_id: int) -> None: + """Drop a frozen-capture pin. Only Mamba managers own capture blocks.""" + return None + + def release_all_boundary_captures(self) -> None: + """Release every frozen-capture pin (cache reset). No-op by default.""" + return None + class FullAttentionManager(SingleTypeKVCacheManager): supports_fine_grained_hash_lookup: ClassVar[bool] = True @@ -1907,6 +1915,18 @@ def __init__( self.mamba_cache_mode == "align" and kv_cache_spec.num_prefill_checkpoint_blocks > 0 ) + # Frozen boundary capture is align-mode only; MambaSpec carries the + # enablement decision (align + OffloadingConnector + kv_offloading). + # The bookkeeping exists for every mamba manager: the cache-manager + # release paths consult it unconditionally. + self.boundary_capture = ( + kv_cache_spec.boundary_capture and self.mamba_cache_mode == "align" + ) + # Recorded crossings (promoted at the next schedule pass) and the + # manager pins on promoted capture blocks (released on store ack): + # dedicated single-writer blocks holding state@B. + self._pending_captures: dict[str, list[tuple[Request, KVCacheBlock, int]]] = {} + self._capture_pinned: dict[int, KVCacheBlock] = {} # Mamba checkpoints follow Eagle's global replay boundary. self.drop_eagle_checkpoint_block = False self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() @@ -1927,6 +1947,112 @@ def __init__( self._producer_partial_tail_reqs: dict[str, tuple[KVCacheBlock, int]] = {} self._packed_prefill_checkpoint_reqs: set[str] = set() + def release_boundary_capture(self, block_id: int) -> None: + """Drop the manager pin on a frozen capture block once its store job + completed (or the connector dropped the offer). Safe against + in-flight DMA races: the connector fences block reuse for pending + store jobs, and the copy rail already retained the endpoints until + the copy step was processed.""" + block = self._capture_pinned.pop(block_id, None) + if block is not None: + self.block_pool.free_blocks([block]) + + def release_all_boundary_captures(self) -> None: + """Release every capture pin and pending record (cache reset).""" + blocks = list(self._capture_pinned.values()) + self._capture_pinned.clear() + for entries in self._pending_captures.values(): + blocks.extend(column for _, column, _ in entries) + self._pending_captures.clear() + self.block_pool.free_blocks(blocks) + + def _record_boundary_capture( + self, request: Request, column_block: KVCacheBlock, boundary: int + ) -> None: + """Remember that the step scheduled right now will write state at + ``boundary`` into ``column_block``. Promoted (copy + offer) at the + start of the next schedule pass, before the column can advance.""" + # Invariant (connector-aligned-builder ownership): a capture boundary + # is a state-block multiple, so the offload chunk grid claims it. + assert boundary % self.block_size == 0 + # Hold the column against retirement until promotion (next pass). + self.block_pool.touch([column_block]) + self._pending_captures.setdefault(request.request_id, []).append( + (request, column_block, boundary) + ) + + def _promote_boundary_captures(self) -> None: + """Turn recorded crossings into frozen capture blocks. + + Runs at the start of the schedule pass after the producing step was + dispatched. The producing forward has completed (or is queued ahead of + everything on the worker's compute stream), so: + - the copy rail (``kv_cache_block_copies``, executed in the next + worker step after zeroing and BEFORE ``preprocess_state``'s precopy + and the forward) reads the column exactly at state@boundary; + - the connector offer ships with the same scheduler output, and the + store DMA is submitted one further step later, chained behind the + copy on the transfer stream. + The capture block is a dedicated pool block: never in a request's + column table, so precopy/postprocess/spec relocation cannot touch it, + and it is hashed at the boundary so local hits can resume past the + last live column too. + """ + pending, self._pending_captures = self._pending_captures, {} + for req_id, entries in pending.items(): + req_blocks = self.req_to_blocks.get(req_id, []) + for request, column_block, boundary in entries: + column_idx = boundary // self.block_size - 1 + if ( + column_block.is_null + or column_block.block_hash is None + or column_block.block_hash_num_tokens != boundary + or column_idx >= len(req_blocks) + or req_blocks[column_idx] is not column_block + ): + # The column was retired, re-hashed, CoW-moved, or the + # producing step was rewound before committing state@B. + # (Never gate on request.num_computed_tokens here: under + # async scheduling it is an optimistic mirror that lags + # the dispatched step.) + self.block_pool.free_blocks([column_block]) + continue + try: + capture_block = self.block_pool.get_new_blocks(1)[0] + except ValueError: + # Pool pressure: skip the capture. The boundary is simply + # not stored; hits fall back to recompute (correct, slower). + self.block_pool.free_blocks([column_block]) + continue + # Ref plan (I4): the capture block keeps its allocation ref + # as a manager pin until the store DMA ack releases it + # (release_boundary_capture); it can then never be re-allocated + # as a copy destination, a live column, or zero-filled while + # its own content is still owed to the CPU tier. `touch` adds + # the copy-rail ref on each endpoint (released when the copy + # step is processed by take_kv_cache_block_copies / + # _free_cow_retained_blocks); the record-time ref on the + # column is dropped here. + self.block_pool.touch([column_block, capture_block]) + self._pending_cow_copies.append((column_block, capture_block)) + self._capture_pinned[capture_block.block_id] = capture_block + block_hash = self.block_pool.cache_partial_block( + request=request, + block=capture_block, + num_tokens=boundary, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + replace_existing_hashes=True, + ) + if block_hash is not None: + # Content lands with the copy at the start of this step's + # execution; defer same-step local hits (CoW precedent). + self.cached_blocks_this_step.add(block_hash) + self._pending_boundary_state_offloads.append( + (req_id, self.kv_cache_group_id, capture_block, boundary) + ) + self.block_pool.free_blocks([column_block]) + @classmethod def find_longest_cache_hit( cls, @@ -2463,6 +2589,14 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: for entry in self._pending_boundary_state_offloads if entry[0] != request_id ] + # Recorded-but-not-yet-promoted crossings can never ship a copy + # now that the column table is going away: drop them along with + # the record-time pin on each column. (Already-promoted captures + # are pool-owned pins held for the store ack, not for this + # request, and stay pinned.) + dropped = self._pending_captures.pop(request_id, ()) + if dropped: + self.block_pool.free_blocks([col for _, col, _ in dropped]) return super().pop_blocks_for_free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: @@ -2515,6 +2649,35 @@ def cache_blocks( self.cached_blocks_this_step.add(block.block_hash) if self.mamba_cache_mode == "align": assert block.block_hash_num_tokens is not None + if self.boundary_capture: + # Frozen capture: the step scheduled right now is the + # producer for this boundary only when the chunk ends + # exactly at it and this is the last (live) column the + # forward writes. Offer the live column never; the + # dedicated capture block is the sole store source. + # The producing step is the one that COMPUTES the + # boundary, i.e. its dispatched chunk still has the + # boundary ahead of it (nct < boundary). When the + # mirror already sits at the boundary, the crossing + # was produced by an EARLIER step (connector-load + # delay: the producing allocate ran with + # delay_cache_blocks and cache_blocks only catches + # up now). Promoting that record copies the column + # AFTER the later step's forward advanced it past + # the boundary -- state@prompt_final stored under + # key@boundary, the original mislabel defect. + # Reject: the boundary simply keeps no capture and + # hits fall back to recompute (correct, slower). + if ( + idx == num_cached_blocks_after - 1 + and num_tokens % self.block_size == 0 + and block.block_hash_num_tokens == num_tokens + and request.num_computed_tokens < num_tokens + ): + self._record_boundary_capture( + request, block, num_tokens + ) + continue # Offer every retained boundary with its exact block. # The connector filters against its save window, which may # extend past the original prompt during resumed prefill. @@ -2529,6 +2692,12 @@ def cache_blocks( def new_step_starts(self) -> None: self.cached_blocks_this_step.clear() + if self.boundary_capture: + # Promote last step's crossings before this pass allocates for the + # new step: the recorded live column still holds state@B until the + # producing forward's successor retires it, and the copy rail ships + # with this step's output. + self._promote_boundary_captures() def _cache_partial_tail_block( self, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 66ed299afbb9..6aaffa5e57bd 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1084,6 +1084,13 @@ class MambaSpec(KVCacheSpec): # False: the state is sharded across TP ranks (e.g. GDN). True: every TP # rank holds the full state (e.g. the replicated PLE conv state). tp_replicated: bool = False + # Frozen boundary capture: at every retention-grid crossing the step + # commits, the manager copies the live state column into a dedicated + # single-writer pool block before the column advances, so a connector + # boundary store never sources a column the next forward can clobber. + # Independent of `num_prefill_checkpoint_blocks` (the in-forward kernel + # export), which this model's backend cannot address safely. + boundary_capture: bool = False def prefill_checkpoint_indices( self, num_computed_tokens: int, num_tokens: int diff --git a/vllm/v1/worker/gpu/boundary_checkpoint.py b/vllm/v1/worker/gpu/boundary_checkpoint.py index 447502b5f851..06902fbdf72f 100644 --- a/vllm/v1/worker/gpu/boundary_checkpoint.py +++ b/vllm/v1/worker/gpu/boundary_checkpoint.py @@ -674,6 +674,11 @@ def add_request(self, slot: int, request: "NewRequestData") -> None: # Prefill consumes a canonical committed prefix, not the previous # verification interval's acceptance count. self.model_state.get_recurrent_checkpoint_acceptance()[slot] = 1 + # The pools this restore wrote are 'restored, not recycled'. + # The selector fresh flag gates the next forward's pool + # zeroing, so a restore that leaves it set would re-zero what + # it just copied -- clear it for the restored slot. + self.model_state.get_recurrent_checkpoint_fresh()[slot] = False for module in self.target_modules: if hasattr(module, "set_recurrent_checkpoint_anchor"): module.set_recurrent_checkpoint_anchor( diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 0a55b6118cc9..275b38c9c556 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -175,6 +175,17 @@ def get_recurrent_checkpoint_acceptance(self) -> torch.Tensor: """Return per-request selector acceptance for boundary MTP replay.""" raise NotImplementedError + def get_recurrent_checkpoint_fresh(self) -> torch.Tensor: + """Return the per-request selector fresh flag for boundary restores. + + Any restore path must clear it: the pools the restore writes are + 'restored, not recycled', and the next forward's fresh-reset would + otherwise re-zero them. Models without a selector fresh flag raise; + the boundary path only calls this when a checkpoint (and therefore a + restorable pool set) exists. + """ + raise NotImplementedError + def postprocess_state( self, idx_mapping: torch.Tensor, From ce25a437a0badeafb75b1a3d73f7429352307ff6 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:42:56 -0400 Subject: [PATCH 199/365] feat(b12x): serve DiffKV decode and verify through paged_decode The B12X attention backend routes decode and speculative-verify batches (max_query_len <= 1 + num_speculative_tokens) to b12x.attention.paged_decode when b12x provides it and the layer geometry is supported, and appends those layers' K/V through the same prepared plan (paged_decode.write_kv). - VLLM_B12X_PAGED_DECODE=auto (default) enables it for layers with unequal Q/K and V head dims (MiMo-V2's 192/128) and for non-causal FP8-KV layers (speculative drafters), which the paged kernels serve only through extend or the BF16-only gather path. "1" enables every supported layer, "0" disables it. Other layers keep their current paths. - paged_decode reads per-request query lengths from the device, so ragged verify batches (adaptive verification trims drafts on device) keep their split-KV schedule. Builders whose layers all use it report supports_varlen_decode_cudagraph, which adaptive verification requires. - Plans are declared per batch bucket in the preparation units; FP8 layers pass their per-layer scales; prefill keeps the paged extend. Tests: tests/v1/attention/test_b12x.py (routing, graph replay with ragged query lengths, KV append byte-equal to the Triton DiffKV writer). Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- tests/v1/attention/test_b12x.py | 162 ++++++++++++++++++- vllm/envs.py | 7 + vllm/utils/b12x.py | 6 + vllm/v1/attention/backends/b12x.py | 248 ++++++++++++++++++++++++++++- 4 files changed, 417 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_b12x.py b/tests/v1/attention/test_b12x.py index b754394b7425..72ad69cade92 100644 --- a/tests/v1/attention/test_b12x.py +++ b/tests/v1/attention/test_b12x.py @@ -757,9 +757,11 @@ def test_b12x_cache_update_and_graph_replay( ) impl.forward(layer, query, key, value, cache, None, output) (unit,) = impl.get_b12x_preparation_units(layer, workload) + # Paged plans bake in the q strides; paged_decode reads them at launch. assert all( request.plan.invocation["operands"]["q"]["strides"] == query.stride() for request in unit.requests + if request.plan.component_id == "attention.paged" ) init_workspace_manager(device) try: @@ -770,7 +772,25 @@ def run(): impl.do_kv_cache_update(layer, key, value, cache, slots) return impl.forward(layer, query, key, value, cache, metadata, output) + pristine = cache.clone() run() + if impl._paged_decode is not None: + # b12x paged_decode.write_kv replaces the Triton DiffKV writer: + # identical cache bytes. + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_diffkv, + ) + + triton_reshape_and_cache_flash_diffkv( + key, + value, + pristine.transpose(1, 2), + slots, + impl.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + assert torch.equal(cache.view(torch.uint8), pristine.view(torch.uint8)) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): run() @@ -812,11 +832,148 @@ def unexpected_plan(*args, **kwargs): reset_workspace_manager() +@pytest.mark.parametrize("q_lens", [[8, 3], [1, 1]]) +def test_b12x_noncausal_fp8_drafter_uses_paged_decode( + default_vllm_config, q_lens, monkeypatch +) -> None: + """Non-causal windowed FP8-KV layers (speculative drafters) run decode and + ragged query blocks through b12x paged_decode, including graph replay.""" + from b12x.attention.paged.reference import attention_reference + from b12x.preparation import PreparationSession + + from tests.v1.attention.utils import dense_kv_cache_views + from vllm.utils.b12x import B12xWorkload, get_b12x_paged_decode + from vllm.v1.worker.workspace import ( + init_workspace_manager, + reset_workspace_manager, + ) + + _require_b12x_paged_attention() + if get_b12x_paged_decode() is None: + pytest.skip("b12x.attention.paged_decode is not available") + device = torch.device("cuda", torch.accelerator.current_device_index()) + config = default_vllm_config + config.model_config = SimpleNamespace( + dtype=torch.bfloat16, max_model_len=2048, is_hybrid=False + ) + config.cache_config.block_size = 64 + config.scheduler_config.max_num_seqs = 2 + config.scheduler_config.max_num_batched_tokens = 64 + config.speculative_config = SimpleNamespace(num_speculative_tokens=7) + config.attention_config.use_non_causal = True + sinks = torch.linspace(-1, 1, 16, device=device, dtype=torch.float32) + impl = B12xPagedAttentionImpl( + num_heads=16, + head_size=128, + scale=128**-0.5, + num_kv_heads=2, + alibi_slopes=None, + sliding_window=1024, + kv_cache_dtype="fp8_e4m3", + sinks=sinks, + ) + impl.process_weights_after_loading(torch.bfloat16) + assert impl._paged_decode is not None + spec = B12xPagedAttentionBackend.customize_spec( + FullAttentionSpec( + block_size=64, num_kv_heads=2, head_size=128, dtype=torch.float8_e4m3fn + ) + ) + num_pages = 40 + raw = torch.zeros( + num_pages * 2 * spec.page_size_bytes, dtype=torch.uint8, device=device + ) + cache, peer = dense_kv_cache_views(raw, spec, num_pages, 2, KVCacheLayout.BLHNC) + peer.fill_(7) + layer = SimpleNamespace( + kv_cache=cache, + _k_scale=torch.tensor(0.5, device=device), + _v_scale=torch.tensor(0.25, device=device), + ) + context_lens = [1200, 131] + seq_lens = [c + q for c, q in zip(context_lens, q_lens)] + gen = torch.Generator(device=device).manual_seed(0) + pages = torch.randperm(num_pages - 1, generator=gen, device=device)[: 40 - 2] + table = torch.zeros((2, 20), dtype=torch.int32, device=device) + table[0, :20] = (pages[:20] + 1).to(torch.int32) + table[1, :3] = (pages[20:23] + 1).to(torch.int32) + slots = torch.cat( + [ + table[i, torch.arange(length, device=device) // 64].long() * 64 + + torch.arange(length, device=device) % 64 + for i, length in enumerate(seq_lens) + ] + ) + offsets = [0, q_lens[0], sum(q_lens)] + metadata = b12x.B12xPagedMetadata( + num_actual_tokens=sum(q_lens), + max_query_len=max(q_lens), + query_start_loc=torch.tensor(offsets, dtype=torch.int32, device=device), + max_seq_len=max(seq_lens), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + block_table=table, + slot_mapping=slots, + causal=False, + ) + key = torch.randn(sum(seq_lens), 2, 128, device=device, dtype=torch.bfloat16) + value = torch.randn(sum(seq_lens), 2, 128, device=device, dtype=torch.bfloat16) + query = torch.randn(sum(q_lens), 16, 128, device=device, dtype=torch.bfloat16) + output = torch.empty(sum(q_lens), 16, 128, device=device, dtype=torch.bfloat16) + workload = B12xWorkload( + stage="state", + token_counts=(1, 2, 16), + fixed_token_counts=(1, 2), + output_dtype=torch.bfloat16, + max_tokens=64, + max_seqs=2, + max_model_len=2048, + ) + impl.forward(layer, query, key, value, cache, None, output) + (unit,) = impl.get_b12x_preparation_units(layer, workload) + assert {r.plan.component_id for r in unit.requests} == {"attention.paged_decode"} + init_workspace_manager(device) + try: + with PreparationSession(device=device, autotune=False) as session: + session.prepare(unit.requests).close() + impl.do_kv_cache_update(layer, key, value, cache, slots) + # Only the new rows are written per step; the whole history here. + impl.forward(layer, query, key, value, cache, metadata, output) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + impl.forward(layer, query, key, value, cache, metadata, output) + for factor in (1.0, -0.5): + query.mul_(factor) + output.fill_(torch.nan) + graph.replay() + references = [] + for i, length in enumerate(q_lens): + start, end = sum(seq_lens[:i]), sum(seq_lens[: i + 1]) + fp8 = torch.float8_e4m3fn + k = (key[start:end].float() / 0.5).to(fp8).float() * 0.5 + v = (value[start:end].float() / 0.25).to(fp8).float() * 0.25 + reference, _ = attention_reference( + query[offsets[i] : offsets[i + 1]].float(), + k, + v, + causal=False, + window_left=1023, + attention_sink_bias=sinks, + ) + references.append(reference) + torch.testing.assert_close( + output.float(), torch.cat(references), atol=2e-2, rtol=2e-2 + ) + assert torch.all(peer == 7) + finally: + reset_workspace_manager() + + def test_b12x_prepares_request_counts_for_verification_graphs(monkeypatch) -> None: from vllm.utils.b12x import B12xWorkload impl = object.__new__(B12xPagedAttentionImpl) - impl._noncausal = None + impl._noncausal = impl._paged_decode = None + impl._causal = True impl.dtype = torch.bfloat16 impl._max_num_seqs = 8 impl._verify_q_per_req = 4 @@ -880,7 +1037,8 @@ def test_b12x_prepares_hybrid_attention_for_allocated_table_width( ) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (12, 0)) impl = object.__new__(B12xPagedAttentionImpl) - impl._noncausal = None + impl._noncausal = impl._paged_decode = None + impl._causal = True impl.device = torch.device("cuda", 0) impl.dtype = impl.kv_torch_dtype = torch.bfloat16 impl.num_heads, impl.num_kv_heads = 32, 2 diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..cb8b231dae18 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -207,6 +207,7 @@ VLLM_B12X_MLA_CKV_GATHER: bool = False VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 + VLLM_B12X_PAGED_DECODE: Literal["auto", "0", "1"] = "auto" VLLM_PLE_TABLE_MEMORY: Literal["ram", "disk"] | None = None VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False @@ -1710,6 +1711,12 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": lambda: int( os.getenv("VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS", "524288") ), + # B12X attention decode/verify through b12x.attention.paged_decode: + # "auto" for layers with unequal Q/K and V head dims and for non-causal + # FP8-KV layers, "1" for every supported layer, "0" never. + "VLLM_B12X_PAGED_DECODE": env_with_choices( + "VLLM_B12X_PAGED_DECODE", "auto", ["auto", "0", "1"] + ), # Qwen3.8-Flash-Next PLE offload policy, resolved by vLLM for b12x. "VLLM_PLE_TABLE_MEMORY": env_with_choices( "VLLM_PLE_TABLE_MEMORY", diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index aa5c3f17968a..cfe30b68d2e6 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -260,6 +260,7 @@ def _import_submodule(module_name: str) -> ModuleType | None: module_name: _import_submodule(module_name) for module_name in ( "b12x.attention.paged", + "b12x.attention.paged_decode", "b12x.attention.sparse_mla", "b12x.attention.compressed_sparse_mla", "b12x.attention.dsa_indexer", @@ -347,6 +348,11 @@ def get_b12x_paged_attention() -> ModuleType | None: return _get_submodule("b12x.attention.paged") +def get_b12x_paged_decode() -> ModuleType | None: + """Split-KV decode/verify attention; absent in b12x releases without it.""" + return _get_submodule("b12x.attention.paged_decode") + + def get_b12x_qsa() -> ModuleType | None: return _get_submodule("b12x.attention.qsa") diff --git a/vllm/v1/attention/backends/b12x.py b/vllm/v1/attention/backends/b12x.py index de64ac22477a..63fe15cff2c7 100644 --- a/vllm/v1/attention/backends/b12x.py +++ b/vllm/v1/attention/backends/b12x.py @@ -12,7 +12,12 @@ import torch -from vllm.config import VllmConfig, get_current_vllm_config +import vllm.envs as envs +from vllm.config import ( + VllmConfig, + get_current_vllm_config, + get_layers_from_vllm_config, +) from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.platforms import current_platform @@ -22,6 +27,7 @@ B12xWorkload, PreparationResourceUnavailableError, get_b12x_paged_attention, + get_b12x_paged_decode, get_b12x_scratch_buffers, ) from vllm.utils.math_utils import cdiv @@ -266,6 +272,16 @@ class B12xPagedMetadata(AttentionMetadata): causal: bool = True +def _group_uses_paged_decode(vllm_config: VllmConfig, layer_names: list[str]) -> bool: + from vllm.model_executor.layers.attention import Attention + + layers = get_layers_from_vllm_config(vllm_config, Attention, layer_names) + return bool(layers) and all( + getattr(layer.impl, "_paged_decode", None) is not None + for layer in layers.values() + ) + + class B12xPagedMetadataBuilder(AttentionMetadataBuilder[B12xPagedMetadata]): """Metadata builder for b12x. @@ -294,6 +310,12 @@ def __init__( device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + # Ragged decode/verify batches (adaptive verification trims drafts on + # device) are graph-safe when every layer of the group runs them + # through paged_decode, which derives its schedule on device. + self.supports_varlen_decode_cudagraph = _group_uses_paged_decode( + vllm_config, layer_names + ) def build( self, @@ -477,8 +499,16 @@ def __init__( ) self._plans: dict[_PagedPlanKey, object] = {} self.supports_quant_query_input = False + # Decode/verify (ragged, up to one verifier row block per request) + # through b12x.attention.paged_decode; prefill keeps the paged extend. + self._causal = not vllm_config.attention_config.use_non_causal + self._paged_decode_max_q = max(1, self._verify_q_per_req) + self._paged_decode_plans: dict[tuple[int, int], object] = {} + self._paged_decode = self._select_paged_decode(default_block_size) + # Noncausal (drafter) query blocks never exceed one verifier row block, + # so paged_decode serves them all; otherwise use the gather path. self._noncausal = None - if vllm_config.attention_config.use_non_causal: + if not self._causal and self._paged_decode is None: from vllm.v1.attention.ops.b12x_noncausal import B12xNoncausalAttention self._noncausal = B12xNoncausalAttention(self) @@ -496,6 +526,188 @@ def __init__( self._extend_q_capacities, ) + def _select_paged_decode(self, page_size: int): + mode = envs.VLLM_B12X_PAGED_DECODE + decode = get_b12x_paged_decode() + if mode == "0" or decode is None or not decode.is_supported(self.device): + return None + fp8 = _is_b12x_fp8_kv_cache(self.kv_cache_dtype) + if not decode.supports( + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + kv_dtype=self.kv_torch_dtype, + max_q_per_req=self._paged_decode_max_q, + window_left=self.window_left, + causal=self._causal, + has_sinks=self._sinks_source is not None, + ): + return None + # "auto": layers the paged kernels serve only through extend for + # ragged verify (unequal Q/K and V head dims) or not at all + # (non-causal FP8 KV); other layers keep their existing path. + if mode == "auto" and not ( + self.head_size != self.output_head_size or (not self._causal and fp8) + ): + return None + logger.info_once( + "b12x paged_decode serves decode/verify for q_heads=%d kv_heads=%d " + "head_dim_qk=%d head_dim_vo=%d window_left=%d causal=%s.", + self.num_heads, + self.num_kv_heads, + self.head_size, + self.output_head_size, + self.window_left, + self._causal, + ) + return decode + + def _paged_decode_plan(self, batch: int, page_size: int): + decode = self._paged_decode + return decode.plan( + decode.Caps( + device=self.device, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + max_batch=batch, + max_q_per_req=self._paged_decode_max_q, + kv_dtype=self.kv_torch_dtype, + window_left=self.window_left, + causal=self._causal, + has_sinks=self._sinks_source is not None, + ) + ) + + def _paged_decode_prepared_call( + self, owner: torch.nn.Module, state: Any, plan, batch: int, page_size: int + ): + from b12x.preparation import PreparedCall + + key_cache, value_cache = self._kv_cache_views(owner.kv_cache) + total_q = batch * self._paged_decode_max_q + scratch = tuple( + torch.empty(spec.shape, dtype=spec.dtype, device=self.device) + for spec in plan.scratch_specs() + ) + q = torch.empty_strided( + (total_q, self.num_heads, self.head_size), + self._query_stride, + dtype=self.dtype, + device=self.device, + ) + output = torch.empty( + (total_q, self.num_heads, self.output_head_size), + dtype=self.dtype, + device=self.device, + ) + page_table = torch.zeros( + (batch, self._max_page_table_widths[page_size]), + dtype=torch.int32, + device=self.device, + ) + cache_seqlens = torch.full( + (batch,), page_size, dtype=torch.int32, device=self.device + ) + cu_seqlens_q = torch.arange( + 0, + total_q + 1, + self._paged_decode_max_q, + dtype=torch.int32, + device=self.device, + ) + k_descale, v_descale = self._paged_decode_descales(owner) + binding = state.bind( + plan=plan, + scratch=scratch, + q=q, + k_cache=key_cache, + v_cache=value_cache, + output=output, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + attention_sink_bias=self.sinks, + k_descale=k_descale, + v_descale=v_descale, + ) + return PreparedCall(run=lambda: state.run(binding), produce=lambda: q.zero_()) + + def _paged_decode_units(self, layer, page_size: int, batches): + plans = {} + requests = [] + for batch in batches: + plan = self._paged_decode_plan(batch, page_size) + plans[(page_size, batch)] = plan + call = ( + lambda state, plan=plan, batch=batch: self._paged_decode_prepared_call( + layer, state, plan, batch, page_size + ) + ) + requests.append( + plan.request( + name=f"attention.paged_decode.{id(layer):x}.p{page_size}.b{batch}", + prepare_call=call, + benchmark_call=call, + ) + ) + self._paged_decode_plans = plans + return requests + + def _paged_decode_descales(self, layer): + if not _is_b12x_fp8_kv_cache(self.kv_cache_dtype): + return None, None + # Per-layer dequant scales, read by the kernel as one value. + return layer._k_scale, layer._v_scale + + def _forward_paged_decode( + self, layer, q, out, key_cache, value_cache, attn_metadata, page_size + ) -> None: + page_table = _ensure_i32_contiguous(attn_metadata.block_table, "block_table") + cache_seqlens = _ensure_i32_contiguous(attn_metadata.seq_lens, "seq_lens") + cu_seqlens_q = _ensure_i32_contiguous( + attn_metadata.query_start_loc, "query_start_loc" + ) + num_reqs = int(cache_seqlens.shape[0]) + bucket = min( + ( + b + for (p, b) in self._paged_decode_plans + if p == page_size and b >= num_reqs + ), + default=None, + ) + if bucket is None: + # Undeclared by preparation (e.g. a profiling shape): declare it + # with its default configuration, materialized on first use. + bucket = num_reqs + self._paged_decode_plans[(page_size, bucket)] = self._paged_decode_plan( + bucket, page_size + ) + plan = self._paged_decode_plans[(page_size, bucket)] + k_descale, v_descale = self._paged_decode_descales(layer) + decode = self._paged_decode + decode.run( + decode.bind( + plan, + scratch=tuple(get_b12x_scratch_buffers(plan)), + q=q, + k_cache=key_cache, + v_cache=value_cache, + output=out, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + attention_sink_bias=self.sinks, + k_descale=k_descale, + v_descale=v_descale, + ) + ) + def _request_name(self, owner: object, key: _PagedPlanKey) -> str: mode, page_size, batch, total_q, q_stride, out_stride = key return ( @@ -725,7 +937,9 @@ def get_b12x_preparation_units( | {self._max_num_seqs} ) ) - for batch in batches: + if self._paged_decode is not None: + requests.extend(self._paged_decode_units(layer, page_size, batches)) + for batch in batches if self._paged_decode is None else (): key = ( "decode", page_size, @@ -796,7 +1010,8 @@ def get_b12x_preparation_units( ), ) ) - for batch in range(1, self._max_num_seqs + 1): + extend_batches = range(1, self._max_num_seqs + 1) if self._causal else () + for batch in extend_batches: for total_q in self._extend_q_capacities: if batch >= total_q: continue @@ -1108,6 +1323,14 @@ def forward( key_cache, value_cache = self._kv_cache_views(kv_cache) page_size = _kv_page_size(key_cache, value_cache) + if ( + self._paged_decode is not None + and int(attn_metadata.max_query_len) <= self._paged_decode_max_q + ): + self._forward_paged_decode( + layer, q, out, key_cache, value_cache, attn_metadata, page_size + ) + return output if not attn_metadata.causal: if self._noncausal is None: raise ValueError("B12X noncausal attention was not configured.") @@ -1173,6 +1396,23 @@ def do_kv_cache_update( ) -> None: if kv_cache.numel() == 0: return + if self._paged_decode is not None: + # Layers paged_decode serves append through its prepared plan. + if not self._paged_decode_plans: + raise RuntimeError("b12x paged_decode KV write before preparation") + key_cache, value_cache = self._kv_cache_views(kv_cache) + k_scale, v_scale = self._paged_decode_descales(layer) + self._paged_decode.write_kv( + next(iter(self._paged_decode_plans.values())), + key=key, + value=value, + k_cache=key_cache, + v_cache=value_cache, + slot_mapping=slot_mapping, + k_scale=k_scale, + v_scale=v_scale, + ) + return if self.head_size != self.output_head_size: from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash_diffkv, From 84422c8bac884431776249a2c2f4fef622526f7e Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:42:55 -0400 Subject: [PATCH 200/365] perf(b12x): fuse all-reduce and RMSNorm at MiMo V2 and DFlash norm sites The allreduce_rms pass rewrites all-reduce -> fused_add_rms_norm into b12x_fused_allreduce_add_rms_norm, but the fused one-shot kernel only runs for plans a model has declared, and the plans were keyed by id(norm_weight). No model declared them, and the compiled graph passes an alias of the declared Parameter anyway, so every call fell back to an out-of-place all-reduce, a copy and a separate RMSNorm (slower than no fusion). - Key the norm weight by data_ptr(). Unfused all-reduce plans are unchanged. - Add declare_b12x_fused_allreduce_rms_norm_sites(owner, norms, hidden, prefix) next to register_b12x_collective_describer, and declare MiMo V2 (post-attention norms, input norms of layers 1.., final norm: 96 sites for Flash) and the DFlash drafter (10 sites for a 5-layer draft). It is a no-op unless the native b12x PCIe transport is active. MiMo-V2.6-Flash TP4 (SM120), C1 decode: separate RMSNorm launches go from 101 to 10 per step, the post-all-reduce copies disappear, and step time drops by about 0.1 ms. The precision gate is at its noise floor. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../b12x_pcie_all_reduce.py | 10 ++++- vllm/distributed/parallel_state.py | 39 +++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 26 +++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 18 +++++++++ 4 files changed, 92 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py index e82d1510731f..247483687ca3 100644 --- a/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/b12x_pcie_all_reduce.py @@ -754,7 +754,15 @@ def prepare_dma(state, invocation=invocation): @staticmethod def _plan_key(operation, shape, dtype, strides, weight=None, epsilon=None): - norm = None if operation == "all_reduce" else (id(weight), epsilon) + # Key the fused RMSNorm weight by storage, not Python identity: a + # compiled graph passes an alias of the declared Parameter (e.g. + # ``weight.data``), so ``id()`` never matches and every fused call + # would fall back to all-reduce + copy + norm. + norm = ( + None + if operation == "all_reduce" + else (weight.data_ptr() if weight is not None else None, epsilon) + ) return operation, tuple(shape), dtype, tuple(strides), norm def _index_declared_plans(self) -> None: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index f5fc7bf9b4db..bbc40bea35d3 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -218,6 +218,45 @@ def register_b12x_collective_describer(owner, describe, *, group=None) -> bool: return True +def declare_b12x_fused_allreduce_rms_norm_sites( + owner, norms, hidden_size: int, name_prefix: str, *, group=None +) -> int: + """Declare a model's TP all-reduce -> residual-add RMSNorm pairs. + + The ``allreduce_rms`` fusion pass rewrites these pairs into + ``b12x_fused_allreduce_add_rms_norm``, whose fused kernel only runs for + declared (shape, norm weight, epsilon) plans; an undeclared call falls back + to all-reduce, a copy, and a separate RMSNorm. ``norms`` is an iterable of + ``(name, RMSNorm)`` for norms whose input is a TP all-reduce output. + Returns the number of declared sites (0 without the native transport). + """ + norms = list(norms) + if not norms: + return 0 + + def describe(requirements): + from vllm.distributed.device_communicators.b12x_pcie_all_reduce import ( + B12xPcieInvocation, + ) + + return tuple( + B12xPcieInvocation( + name=f"{name_prefix}.{name}.fused_ar_norm.m{rows}", + operation="all_reduce_fused_add_rms_norm", + shape=(rows, hidden_size), + dtype=requirements.output_dtype, + norm_weight=norm.weight, + epsilon=float(norm.variance_epsilon), + ) + for name, norm in norms + for rows in requirements.token_counts + ) + + if not register_b12x_collective_describer(owner, describe, group=group): + return 0 + return len(norms) + + def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: assert group_name in _groups, f"Group {group_name} is not found." group = _groups[group_name]() diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 78af26326ca9..862614e1764a 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -20,6 +20,9 @@ get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, ) +from vllm.distributed.parallel_state import ( + declare_b12x_fused_allreduce_rms_norm_sites, +) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -587,6 +590,29 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.norm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) else: self.norm = PPMissingLayer() + self._declare_fused_allreduce_rms_norm_sites() + + def _declare_fused_allreduce_rms_norm_sites(self) -> None: + # Each post-attention norm follows the o_proj all-reduce, and each later + # layer's input norm follows the previous layer's MLP/MoE all-reduce. + norms = [] + layers = islice(self.layers, self.start_layer, self.end_layer) + for i, layer in enumerate(layers): + norms.append( + ( + f"layers.{layer.layer_id}.post_attention_layernorm", + layer.post_attention_layernorm, + ) + ) + if i > 0: + norms.append( + (f"layers.{layer.layer_id}.input_layernorm", layer.input_layernorm) + ) + if isinstance(self.norm, RMSNorm): + norms.append(("norm", self.norm)) + declare_b12x_fused_allreduce_rms_norm_sites( + self, norms, self.config.hidden_size, "mimo_v2" + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3b5b01f5afb2..02fc03c0baf8 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -17,6 +17,9 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.distributed.parallel_state import ( + declare_b12x_fused_allreduce_rms_norm_sites, +) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm @@ -495,6 +498,21 @@ def __init__( self.config.hidden_size, eps=self.config.rms_norm_eps, ) + # Residual RMSNorms fed by a TP all-reduce: post-attention norms, input + # norms of layers 1.., and the final norm. + fused_norms = [] + for i, layer in enumerate(self.layers): + fused_norms.append( + (f"layers.{i}.post_attention_layernorm", layer.post_attention_layernorm) + ) + if i > 0: + fused_norms.append( + (f"layers.{i}.input_layernorm", layer.input_layernorm) + ) + fused_norms.append(("norm", self.norm)) + declare_b12x_fused_allreduce_rms_norm_sites( + self, fused_norms, self.config.hidden_size, "dflash" + ) # The context projection concatenates K/V weights from every draft # layer. It is not a LinearBase module because its output layout is a # DFlash-specific fusion. Serialized MXFP8 checkpoints retain an From 0f4b6cc62096d9ed7bbc35fe04de89e977009ebc Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:44:20 -0400 Subject: [PATCH 201/365] perf(b12x): serve decode-sized BF16 linears through bf16_gemv plans With VLLM_B12X_BF16_GEMV=1 under --linear-backend b12x, unquantized BF16 linears declare b12x gemm.bf16_gemv plans for the CUDA-graph token counts up to eight rows. The preparation driver autotunes each plan across b12x's GEMV kernels and cuBLAS (b12x's torch backend), so a layer only leaves cuBLAS where a native kernel measures faster; larger batches keep F.linear. The routing decision uses the layer's name, so it holds inside the compiled graph. Default off. Tests: tests/model_executor/kernels/test_b12x_linear.py (opt-in default, prepared plans under fullgraph torch.compile). Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../kernels/test_b12x_linear.py | 73 ++++++ vllm/envs.py | 4 + .../kernels/linear/b12x_unquantized.py | 218 ++++++++++++++++++ vllm/model_executor/layers/linear.py | 7 + vllm/model_executor/layers/utils.py | 7 + vllm/utils/b12x.py | 5 + 6 files changed, 314 insertions(+) create mode 100644 vllm/model_executor/kernels/linear/b12x_unquantized.py diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 23a6c8411111..891e38a9b3a1 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -2982,3 +2982,76 @@ def test_v41_unquantized_prepares_dtypes_and_exact_rows_before_replay( finally: graph.reset() session.close() + + +def test_b12x_unquantized_decode_linear_is_opt_in( + default_vllm_config, monkeypatch +) -> None: + """Without VLLM_B12X_BF16_GEMV=1, BF16 linears keep the stock F.linear path.""" + from vllm.model_executor.kernels.linear.b12x_unquantized import ( + maybe_attach_b12x_bf16_gemv, + ) + + monkeypatch.delenv("VLLM_B12X_BF16_GEMV", raising=False) + weight = torch.zeros(256, 512, dtype=torch.bfloat16) + layer = _WeakNamespace(prefix="test.bf16_gemv.off", weight=weight, bias=None) + assert not maybe_attach_b12x_bf16_gemv(layer) + assert not hasattr(layer, "b12x_bf16_gemv_plans") + + +@pytest.mark.parametrize("n,k", [(4096, 2048), (2560, 4096)]) +@torch.inference_mode() +def test_b12x_unquantized_decode_linear_uses_prepared_plans( + default_vllm_config, monkeypatch, tmp_path, n, k +) -> None: + """Decode-sized BF16 linears run prepared b12x plans; larger batches F.linear.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("requires SM12x") + from b12x._lib.runtime_control import kernel_resolution_guard + + monkeypatch.setenv("VLLM_B12X_BF16_GEMV", "1") + + from vllm.model_executor.kernels.linear.b12x_unquantized import ( + b12x_unquantized_gemm, + maybe_attach_b12x_bf16_gemv, + ) + + device = torch.device("cuda", torch.accelerator.current_device_index()) + torch.manual_seed(7) + weight = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05 + layer = _WeakNamespace(prefix=f"test.bf16_gemv.{n}x{k}", weight=weight, bias=None) + assert maybe_attach_b12x_bf16_gemv(layer) + session, _ = _prepare( + layer, device=device, counts=(1, 2, 4, 8, 16), cache_dir=tmp_path + ) + session.freeze() + assert sorted(layer.b12x_bf16_gemv_plans) == [1, 2, 4, 8] + + def check(source, actual): + expected = source.double() @ weight.double().T + torch.testing.assert_close(actual.double(), expected, rtol=1e-2, atol=1e-2) + + compiled = torch.compile( + lambda source: b12x_unquantized_gemm(layer, source, weight), + fullgraph=True, + dynamic=False, + ) + with session, kernel_resolution_guard("prepared b12x BF16 decode linear"): + for rows in (1, 3, 8, 16): + source = torch.randn(rows, k, device=device, dtype=torch.bfloat16) + check(source, b12x_unquantized_gemm(layer, source, weight)) + # Model forwards are traced: the routing check must not need the + # opaque layer name's value. + source = torch.randn(4, k, device=device, dtype=torch.bfloat16) + check(source, compiled(source)) + source = torch.randn(4, k, device=device, dtype=torch.bfloat16) + graph = torch.cuda.CUDAGraph() + try: + with session.capture(), torch.cuda.graph(graph): + actual = b12x_unquantized_gemm(layer, source, weight) + source.neg_() + graph.replay() + torch.accelerator.synchronize() + check(source, actual) + finally: + graph.reset() diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..2f2a374f8fae 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -191,6 +191,7 @@ VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_B12X_MOE_FP4_FORCE_A16: bool = False + VLLM_B12X_BF16_GEMV: bool = False VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE: Literal["0", "1", "all", "w13", "w2"] = "0" VLLM_DEFAULT_MOE_BACKEND: str = "auto" VLLM_B12X_DENSE_ACTIVATION_MODE: Literal["auto", "a16", "quantized"] = "auto" @@ -1656,6 +1657,9 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool( int(os.getenv("VLLM_B12X_MOE_FP4_FORCE_A16", "0")) ), + # Opt-in: under --linear-backend b12x, serve decode-sized (<= 8 row) BF16 + # linears through b12x gemm.bf16_gemv plans autotuned against cuBLAS. + "VLLM_B12X_BF16_GEMV": lambda: os.getenv("VLLM_B12X_BF16_GEMV", "0") == "1", # Select layer-wide activation scales for b12x NVFP4 MoE projections. "VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE": env_with_choices( "VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE", diff --git a/vllm/model_executor/kernels/linear/b12x_unquantized.py b/vllm/model_executor/kernels/linear/b12x_unquantized.py new file mode 100644 index 000000000000..6ec9e8bc03ec --- /dev/null +++ b/vllm/model_executor/kernels/linear/b12x_unquantized.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decode-sized unquantized linears through b12x ``gemm.bf16_gemv``. + +Under ``--linear-backend b12x``, BF16 linear layers declare b12x projection +plans for the CUDA-graph token counts up to ``MAX_ROWS``. Preparation +autotunes each plan across b12x's GEMV kernels and cuBLAS (b12x's ``torch`` +backend), so a layer leaves cuBLAS only where a native kernel measures faster; +a layer whose plans all keep cuBLAS traces a plain ``F.linear``. Larger +batches keep ``F.linear``. +""" + +from collections.abc import Sequence + +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config_or_none +from vllm.platforms import current_platform +from vllm.utils.b12x import ( + B12xPreparationUnit, + B12xWorkload, + b12x_layer, + b12x_layer_prefix, + get_b12x_bf16_gemv, + register_b12x_layer, + set_b12x_preparation_provider, +) +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) + +MAX_ROWS = 8 + + +def _aligned(tensor: torch.Tensor) -> bool: + return tensor.data_ptr() % 16 == 0 and tensor.stride(0) % 8 == 0 + + +def _plan(api, layer: torch.nn.Module, rows: int, bias: torch.Tensor | None): + n, k = map(int, layer.weight.shape) + return api.plan( + api.GemvQuery( + source_dtype="bfloat16", + weight_dtype="bfloat16", + max_rows=rows, + in_features=k, + out_features=n, + source_contiguous=True, + source_aligned=True, + weight_contiguous=True, + weight_aligned=True, + bias_dtype=None if bias is None else "bfloat16", + ) + ) + + +def _b12x_bf16_linear( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + layer_name: LayerNameType, +) -> torch.Tensor: + layer = b12x_layer(_resolve_layer_name(layer_name)) + rows = int(x.shape[0]) + plan = min( + ( + (capacity, plan) + for capacity, plan in layer.b12x_bf16_gemv_plans.items() + if capacity >= rows + ), + default=(None, None), + key=lambda item: item[0], + )[1] + if plan is None or rows == 0 or not x.is_contiguous() or not _aligned(x): + return torch.nn.functional.linear(x, weight, bias) + api = get_b12x_bf16_gemv() + assert api is not None + return api.mm(x, weight, plan=plan, bias=bias) + + +def _b12x_bf16_linear_fake( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + layer_name: LayerNameType, +) -> torch.Tensor: + del bias, layer_name + return x.new_empty((x.shape[0], weight.shape[0])) + + +direct_register_custom_op( + op_name="b12x_bf16_linear", + op_func=_b12x_bf16_linear, + fake_impl=_b12x_bf16_linear_fake, +) + + +def b12x_bf16_gemv_enabled() -> bool: + """Routing needs autotuning: only measured plans may leave cuBLAS.""" + config = get_current_vllm_config_or_none() + return ( + envs.VLLM_B12X_BF16_GEMV + and config is not None + and config.kernel_config.enable_b12x_autotune + and current_platform.is_cuda() + and current_platform.is_device_capability_family(120) + and get_b12x_bf16_gemv() is not None + ) + + +def maybe_attach_b12x_bf16_gemv(layer: torch.nn.Module) -> bool: + """Route the layer's decode-sized calls through b12x when its weight fits.""" + weight = getattr(layer, "weight", None) + bias = getattr(layer, "bias", None) + if ( + not b12x_bf16_gemv_enabled() + or not isinstance(weight, torch.Tensor) + or weight.dtype != torch.bfloat16 + or weight.ndim != 2 + or not weight.is_contiguous() + or not _aligned(weight) + or (bias is not None and bias.dtype != torch.bfloat16) + or getattr(layer, "b12x_preparation_suppressed", False) + ): + return False + name = b12x_layer_prefix(layer) + register_b12x_layer(name, layer) + layer.b12x_layer_name = _encode_layer_name(name) + # Plain-string copy for the trace-time routing check: torch.compile sees + # the opaque LayerName as a FakeScriptObject. + layer.b12x_bf16_gemv_name = name + layer.b12x_bf16_gemv_plans = {} + set_b12x_preparation_provider(layer, _B12xBf16GemvProvider()) + return True + + +@torch.compiler.assume_constant_result +def _selects_native_kernel(layer_name: str) -> bool: + """Whether preparation picked a b12x kernel over cuBLAS for any row count. + + Weight-stage preparation completes before the model is first traced, so a + layer whose plans all chose cuBLAS keeps a plain ``F.linear``. + """ + plans = b12x_layer(layer_name).b12x_bf16_gemv_plans + return any( + plan.selection is not None and plan.selection.config.backend != "torch" + for plan in plans.values() + ) + + +def b12x_unquantized_gemm( + layer: torch.nn.Module, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + if ( + getattr(layer, "b12x_bf16_gemv_plans", None) is None + or x.ndim != 2 + or not _selects_native_kernel(layer.b12x_bf16_gemv_name) + ): + return torch.nn.functional.linear(x, weight, bias) + return torch.ops.vllm.b12x_bf16_linear(x, weight, bias, layer.b12x_layer_name) + + +class _B12xBf16GemvProvider: + """Declares one autotuned projection plan per decode token count.""" + + def get_b12x_preparation_units( + self, + layer: torch.nn.Module, + workload: B12xWorkload, + ) -> Sequence[B12xPreparationUnit]: + api = get_b12x_bf16_gemv() + weight = layer.weight + if api is None or weight.is_meta: + return () + bias = getattr(layer, "bias", None) + k = int(weight.shape[1]) + prefix = _resolve_layer_name(layer.b12x_layer_name) + plans = layer.b12x_bf16_gemv_plans + requests = [] + for rows in sorted({r for r in workload.token_counts if r <= MAX_ROWS}): + plan = plans.get(rows) + if plan is None: + plan = plans[rows] = _plan(api, layer, rows, bias) + + def call(state, *, m=rows): + from b12x.preparation import PreparedCall + + x = torch.empty((m, k), dtype=torch.bfloat16, device=weight.device) + return PreparedCall( + run=lambda: state.run(x, weight, bias=bias), + produce=lambda: x.normal_(), + owners=(weight,) if bias is None else (weight, bias), + ) + + requests.append( + plan.request( + name=f"linear.bf16_gemv.{prefix}.m{rows}", + prepare_call=call, + benchmark_call=call, + ) + ) + if not requests: + return () + return ( + B12xPreparationUnit( + name="BF16_GEMV", + key=(prefix, workload.token_counts), + requests=tuple(requests), + stage="weights", + ), + ) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index e7eef58f1978..d1e385ba3040 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -215,6 +215,7 @@ def __init__(self) -> None: linear_backend = ( config.kernel_config.linear_backend if config is not None else "auto" ) + self._linear_backend = linear_backend self._gemm_impl = dispatch_unquantized_gemm(linear_backend) def create_weights( @@ -247,6 +248,12 @@ def create_weights( set_weight_attrs(weight, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self._linear_backend == "b12x" and current_platform.is_cuda(): + from vllm.model_executor.kernels.linear.b12x_unquantized import ( + maybe_attach_b12x_bf16_gemv, + ) + + maybe_attach_b12x_bf16_gemv(layer) if current_platform.is_cpu(): # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical, # so skip packing and use a plain fallback. diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index f483a619b293..8f3eb3610495 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -591,6 +591,13 @@ def dispatch_unquantized_gemm( elif not current_platform.is_cuda(): return default_unquantized_gemm + if linear_backend == "b12x": + from vllm.model_executor.kernels.linear.b12x_unquantized import ( + b12x_unquantized_gemm, + ) + + return b12x_unquantized_gemm + backend_spec = _FLASHINFER_BF16_BACKENDS.get(linear_backend) if backend_spec is None: return default_unquantized_gemm diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index aa5c3f17968a..aeffa61886c4 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -264,6 +264,7 @@ def _import_submodule(module_name: str) -> ModuleType | None: "b12x.attention.compressed_sparse_mla", "b12x.attention.dsa_indexer", "b12x.attention.qsa", + "b12x.gemm.bf16_gemv", "b12x.gemm.bf16_vocab_projection", "b12x.gemm.blockscaled", "b12x.gemm.mla_query_projection", @@ -299,6 +300,10 @@ def get_b12x_blockscaled() -> ModuleType | None: return _get_submodule("b12x.gemm.blockscaled") +def get_b12x_bf16_gemv() -> ModuleType | None: + return _get_submodule("b12x.gemm.bf16_gemv") + + def get_b12x_bf16_vocab_projection() -> ModuleType | None: return _get_submodule("b12x.gemm.bf16_vocab_projection") From 10d3e1d0d7e390ecc0df91bb065fabade1294181 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:42:56 -0400 Subject: [PATCH 202/365] fix(mimo): load fused FP8 QKV exactly when a rank owns several chunks MiMo V2.5/V2.6 checkpoints store fused QKV as num_key_value_heads (4) [Q | K | V] chunks, each with its own 128x128 block-scale grid. Below TP4 a rank owns several chunks, and the loader dequantized, de-interleaved and requantized them. On MiMo-V2.6-Flash-RL at TP2 that changes 40% of the global layers' K and 65% of their V projection weights (1.5-2% relative L2 error): a global chunk's 192 K + 128 V rows share a scale block, so the de-interleaved layout regroups rows under new scales. Keep every row on its checkpoint scale instead: - when Q, K and V each span whole scale blocks (all sliding-window layers), the de-interleaved layout is a block permutation of the checkpoint; - otherwise (global layers) the rank keeps each chunk's K and V rows together, zero-padded to whole blocks, as [Q_1 | ... | Q_g | K_1 V_1 pad | ... | K_g V_g pad], and the forward splits that layout (K costs one small copy per global layer). The requantizing path remains only for geometries neither covers. TP4 (one chunk per rank) is unchanged. Tests: tests/models/quantization/test_mimo_v2_fused_qkv_exact.py (18 cases); real-checkpoint audit at TP2: 0 of 2.87B QKV weights change (was 40% of global K and 65% of global V). Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../test_mimo_v2_fused_qkv_exact.py | 157 ++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 154 +++++++++++++++-- 2 files changed, 293 insertions(+), 18 deletions(-) create mode 100644 tests/models/quantization/test_mimo_v2_fused_qkv_exact.py diff --git a/tests/models/quantization/test_mimo_v2_fused_qkv_exact.py b/tests/models/quantization/test_mimo_v2_fused_qkv_exact.py new file mode 100644 index 000000000000..24537a2ba495 --- /dev/null +++ b/tests/models/quantization/test_mimo_v2_fused_qkv_exact.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiMo fused-QKV FP8 checkpoints load exactly at every TP size that splits them. + +MiMo V2.5/V2.6 store fused QKV as ``num_key_value_heads`` [Q | K | V] chunks +(the training TP layout), each with its own 128x128 block-scale grid. A rank +that owns several chunks must not requantize: global layers put K and V rows +(192 + 128) in a shared scale block, so requantizing their de-interleaved +layout changes the stored FP8 values. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.models.mimo_v2 as mimo + +pytestmark = pytest.mark.cpu_test + +BLOCK = 128 +CHUNKS = 4 +GEOMETRY = { + # name: (num_heads, num_kv_heads, head_dim, v_head_dim) + "global": (64, 4, 192, 128), + "sliding": (64, 8, 192, 128), +} + + +def _checkpoint(num_heads, num_kv_heads, head_dim, v_head_dim, hidden=256): + gen = torch.Generator().manual_seed(0) + rows = (num_heads * head_dim + num_kv_heads * (head_dim + v_head_dim)) // CHUNKS + scale_rows = -(-rows // BLOCK) + weight = (torch.randn(CHUNKS * rows, hidden, generator=gen) * 64).clamp(-448, 448) + scale = torch.rand(CHUNKS * scale_rows, hidden // BLOCK, generator=gen) + 0.5 + return weight.to(torch.float8_e4m3fn), scale, rows, scale_rows + + +def _dequant(weight, scale): + expanded = scale.double().repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1) + return weight.double() * expanded[: weight.shape[0], : weight.shape[1]] + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("kind", ["global", "sliding"]) +def test_fused_qkv_shards_keep_checkpoint_values(kind, tp_size): + heads, kv_heads, head_dim, v_head_dim = GEOMETRY[kind] + weight, scale, rows, scale_rows = _checkpoint(*GEOMETRY[kind]) + kv_chunk_rows = mimo._fused_qkv_kv_chunk_rows( + heads, kv_heads, head_dim, v_head_dim, tp_size, CHUNKS + ) + # Only global layers (K/V share a scale block) on ranks owning several + # chunks need the padded layout. + assert bool(kv_chunk_rows) == (kind == "global" and tp_size < CHUNKS) + + q_rows = heads // CHUNKS * head_dim + k_rows = kv_heads // CHUNKS * head_dim + q_size = heads // tp_size * head_dim + k_size = kv_heads // tp_size * head_dim + v_size = kv_heads // tp_size * v_head_dim + per_rank = CHUNKS // tp_size + for rank in range(tp_size): + w_rank, s_rank = mimo._shard_fp8_qkv_proj( + weight, + scale, + heads, + kv_heads, + head_dim, + v_head_dim, + rank, + tp_size, + CHUNKS, + kv_chunk_rows=kv_chunk_rows, + ) + out_rows = q_size + ( + per_rank * kv_chunk_rows if kv_chunk_rows else k_size + v_size + ) + assert w_rank.shape == (out_rows, weight.shape[1]) + assert s_rank.shape == (-(-out_rows // BLOCK), scale.shape[1]) + + # Split the rank's dequantized rows as the forward splits its output. + q, k, v = mimo._split_qkv( + _dequant(w_rank, s_rank).T, q_size, k_size, v_size, kv_chunk_rows + ) + chunks = [ + _dequant( + weight[c * rows : (c + 1) * rows], + scale[c * scale_rows : (c + 1) * scale_rows], + ) + for c in range(rank * per_rank, (rank + 1) * per_rank) + ] + # Bit-exact: the same FP8 values times the same block scales. + assert torch.equal(q, torch.cat([c[:q_rows] for c in chunks]).T) + assert torch.equal( + k, torch.cat([c[q_rows : q_rows + k_rows] for c in chunks]).T + ) + assert torch.equal(v, torch.cat([c[q_rows + k_rows :] for c in chunks]).T) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("kind", ["global", "sliding"]) +@pytest.mark.parametrize("fused", [True, False]) +def test_attention_projection_matches_rank_layout(monkeypatch, kind, tp_size, fused): + created = {} + + def recorder(name): + def build(*args, **kwargs): + created[name] = (args, kwargs) + return torch.nn.Identity() + + return build + + class Backend: + name = "STUB_DIFFKV" + + @staticmethod + def get_class(): + return Backend + + @staticmethod + def set_head_size_v(size): + pass + + @staticmethod + def get_name(): + return "STUB_DIFFKV" + + monkeypatch.setattr(mimo, "get_tensor_model_parallel_world_size", lambda: tp_size) + monkeypatch.setattr(mimo, "QKVParallelLinear", recorder("qkv")) + monkeypatch.setattr(mimo, "MergedColumnParallelLinear", recorder("merged")) + monkeypatch.setattr(mimo, "RowParallelLinear", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr(mimo, "get_rope", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr(mimo, "Attention", lambda *a, **k: torch.nn.Identity()) + monkeypatch.setattr( + mimo, + "get_current_vllm_config", + lambda: SimpleNamespace(attention_config=SimpleNamespace(backend=Backend)), + ) + heads, kv_heads, head_dim, v_head_dim = GEOMETRY[kind] + attn = mimo.MiMoV2Attention( + hidden_size=256, + num_heads=heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + v_head_dim=v_head_dim, + fused_qkv_chunks=CHUNKS if fused else 0, + prefix="model.layers.0.self_attn", + ) + if fused and kind == "global" and tp_size < CHUNKS: + assert set(created) == {"merged"} + _, output_sizes = created["merged"][0] + # Q rows, then every chunk's K|V rows padded to whole scale blocks. + assert output_sizes == [heads * head_dim, CHUNKS * 3 * BLOCK] + assert attn.kv_chunk_rows == 3 * BLOCK + else: + assert set(created) == {"qkv"} + assert attn.kv_chunk_rows == 0 diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 78af26326ca9..a9c366b9d58c 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -51,6 +51,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.model_executor.weight_transfer import copy_weight, materialize_weight from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -230,6 +231,7 @@ def __init__( cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, partial_rotary_factor: float = 1.0, + fused_qkv_chunks: int = 0, prefix: str = "", ) -> None: super().__init__() @@ -256,16 +258,38 @@ def __init__( self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=attention_bias, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - v_head_size=self.v_head_dim, + # A fused-QKV checkpoint stores ``fused_qkv_chunks`` [Q | K | V] + # chunks; see _fused_qkv_kv_chunk_rows for the rank layout. + self.kv_chunk_rows = _fused_qkv_kv_chunk_rows( + num_heads, + num_kv_heads, + head_dim, + self.v_head_dim, + tp_size, + fused_qkv_chunks, ) + if self.kv_chunk_rows: + self.qkv_proj = MergedColumnParallelLinear( + hidden_size, + [ + self.total_num_heads * self.head_dim, + fused_qkv_chunks * self.kv_chunk_rows, + ], + bias=attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + else: + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + v_head_size=self.v_head_dim, + ) self.o_proj = RowParallelLinear( self.total_num_heads * self.v_head_dim, @@ -342,7 +366,9 @@ def forward( hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q, k, v = _split_qkv( + qkv, self.q_size, self.k_size, self.v_size, self.kv_chunk_rows + ) q, k = self.rotary_emb(positions, q, k) # Apply v_scale before attention @@ -370,6 +396,12 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: max_position_embeddings = getattr(config, "max_position_embeddings", 32768) v_scale = getattr(config, "attention_value_scale", None) + fused_qkv_chunks = ( + config.num_key_value_heads + if getattr(config, "attention_projection_layout", None) == "fused_qkv" + and getattr(quant_config, "weight_block_size", None) + else 0 + ) if self.is_compressed_softmax_layer(): self.self_attn = MiMoV2Attention( @@ -389,6 +421,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: max_position_embeddings=max_position_embeddings, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + fused_qkv_chunks=fused_qkv_chunks, prefix=f"{prefix}.self_attn", ) else: @@ -406,6 +439,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: max_position_embeddings=max_position_embeddings, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + fused_qkv_chunks=fused_qkv_chunks, prefix=f"{prefix}.self_attn", ) @@ -462,6 +496,54 @@ def is_compressed_softmax_layer(self) -> bool: return self.config.hybrid_layer_pattern[self.layer_id] == 1 +def _fused_qkv_kv_chunk_rows( + num_heads: int, + num_kv_heads: int, + head_dim: int, + v_head_dim: int, + tp_size: int, + checkpoint_tp_size: int, + block: int = 128, +) -> int: + """Rows per checkpoint chunk of a rank's ``K | V`` fused-QKV section. + + Returns 0 when the de-interleaved ``[Q | K | V]`` rank layout is exact: + a rank owns at most one checkpoint chunk, or every chunk's Q, K and V + rows span whole FP8 scale blocks, so the layout is a block permutation. + + Otherwise, when a rank owns several chunks whose K and V rows share a + scale block (MiMo global layers: 192 K + 128 V rows), the rank keeps each + chunk's K and V rows together, zero-padded to whole blocks, so every row + keeps its checkpoint scale: + + [Q_1 | ... | Q_g | K_1 V_1 pad | ... | K_g V_g pad] + + and this returns the padded ``K | V`` rows per chunk. + """ + if checkpoint_tp_size <= tp_size or checkpoint_tp_size % tp_size: + return 0 + q_rows = num_heads // checkpoint_tp_size * head_dim + k_rows = num_kv_heads // checkpoint_tp_size * head_dim + v_rows = num_kv_heads // checkpoint_tp_size * v_head_dim + if q_rows % block or (k_rows % block == 0 and v_rows % block == 0): + return 0 + return cdiv(k_rows + v_rows, block) * block + + +def _split_qkv( + qkv: torch.Tensor, q_size: int, k_size: int, v_size: int, kv_chunk_rows: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a rank's fused QKV output (see _fused_qkv_kv_chunk_rows).""" + if not kv_chunk_rows: + return qkv.split([q_size, k_size, v_size], dim=-1) + q, kv = qkv.split([q_size, qkv.shape[-1] - q_size], dim=-1) + kv = kv.unflatten(-1, (-1, kv_chunk_rows)) + k_rows, v_rows = k_size // kv.shape[-2], v_size // kv.shape[-2] + k = kv[..., :k_rows].flatten(-2) + v = kv[..., k_rows : k_rows + v_rows].flatten(-2) + return q, k, v + + def _shard_fp8_qkv_proj( w_full: torch.Tensor, s_full: torch.Tensor, @@ -473,6 +555,7 @@ def _shard_fp8_qkv_proj( tp_size: int, checkpoint_tp_size: int, block: int = 128, + kv_chunk_rows: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """Shard the fp8 qkv_proj weights for ``tp_rank``. @@ -488,13 +571,13 @@ def _shard_fp8_qkv_proj( [Q_1 | Q_2 | ... | Q_g | K_1 | K_2 | ... | K_g | V_1 | V_2 | ... | V_g] When ``g == 1`` the rank's slice is already ``[Q | K | V]``, so a plain - chunk suffices. When ``g > 1`` we cannot reach the de-interleaved layout by - re-permuting the fp8 block scales: each scale covers a 128-row block, and - since K is 192 rows (1.5 blocks) a block straddles the K/V boundary, so no - whole-block permutation produces it. Instead we dequantize this rank's - groups to float (dropping the block constraint), reorder the rows into the - layout above (Q, K, and V then each span a whole number of blocks), and - re-quantize to fp8. + chunk suffices. When Q, K and V each span whole 128-row scale blocks, the + layout above is a block permutation of the checkpoint and is exact. When + a block straddles the K/V boundary (a 192-row K is 1.5 blocks), the rank + uses the padded ``[Q_1 | ... | Q_g | K_1 V_1 pad | ... | K_g V_g pad]`` + layout of ``_fused_qkv_kv_chunk_rows`` instead (``kv_chunk_rows``), again + exact. Only when neither applies do we dequantize this rank's groups to + float, reorder the rows into the layout above, and re-quantize to fp8. """ assert checkpoint_tp_size % tp_size == 0, ( "TP size must evenly split the checkpoint TP size." @@ -516,8 +599,42 @@ def _shard_fp8_qkv_proj( scale_rows_per_group = (rows_per_group + block - 1) // block assert w_full.shape[0] == rows_per_group * checkpoint_tp_size assert s_full.shape[0] == scale_rows_per_group * checkpoint_tp_size + groups = range(tp_rank * groups_per_rank, (tp_rank + 1) * groups_per_rank) + q_end = q_rows_per_group + k_end = q_end + k_rows_per_group + if kv_chunk_rows or all( + n % block == 0 for n in (q_rows_per_group, k_rows_per_group, v_rows_per_group) + ): + # Exact: every row keeps its checkpoint FP8 value and block scale. + def rows(g: int, start: int, stop: int) -> torch.Tensor: + base = g * rows_per_group + return materialize_weight(w_full[base + start : base + stop]) + + def scales(g: int, start: int, stop: int) -> torch.Tensor: + base = g * scale_rows_per_group + return materialize_weight( + s_full[base + start // block : base + cdiv(stop, block)] + ) + + if kv_chunk_rows: + pad = kv_chunk_rows - (rows_per_group - q_end) + weights = [rows(g, 0, q_end) for g in groups] + for g in groups: + kv = rows(g, q_end, rows_per_group).view(torch.uint8) + zeros = kv.new_zeros((pad, kv.shape[1])) # FP8 +0.0 + weights.append(torch.cat([kv, zeros]).view(w_full.dtype)) + scale_rows = [scales(g, 0, q_end) for g in groups] + scale_rows += [scales(g, q_end, rows_per_group) for g in groups] + else: + segments = ((0, q_end), (q_end, k_end), (k_end, rows_per_group)) + weights = [rows(g, *seg) for seg in segments for g in groups] + scale_rows = [scales(g, *seg) for seg in segments for g in groups] + return ( + torch.cat([w.view(torch.uint8) for w in weights]).view(w_full.dtype), + torch.cat(scale_rows), + ) qs, ks, vs = [], [], [] - for g_idx in range(tp_rank * groups_per_rank, (tp_rank + 1) * groups_per_rank): + for g_idx in groups: row_start = g_idx * rows_per_group scale_row_start = g_idx * scale_rows_per_group # Dequantize this group's weights. @@ -826,6 +943,7 @@ def _try_load_fp8_qkv_proj( tp_rank=tp_rank, tp_size=tp_size, checkpoint_tp_size=self.config.num_key_value_heads, + kv_chunk_rows=getattr(attn, "kv_chunk_rows", 0), ) sharded = {"weight": w_rank, "weight_scale_inv": s_rank} for kind, tensor in sharded.items(): From 7ba930357b719ac2da54ad94cdeb528fbdd7485d Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 18:38:57 -0400 Subject: [PATCH 203/365] perf(mimo): prefetch upcoming decode weights into L2 on a side stream VLLM_MIMO_L2_PREFETCH=1 prefetches upcoming decode weights into L2 on a side stream while device memory idles, using the GLM-5.3 L2 prefetch machinery: this layer's o_proj, router and post-attention norm at the attention op, and the next layer's input norm and qkv projection at the end of the MoE op, before its all-reduce. Prefetches are issued only inside FULL CUDA-graph decode captures and are cache hints, so numerics are unchanged; the last layer rejoins the side stream. The attention op and MoERunner._forward_impl gain optional hook attributes: both run inside opaque custom ops, so torch.compile never traces them. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- vllm/envs.py | 4 + .../layers/attention/attention.py | 5 ++ .../layers/fused_moe/runner/moe_runner.py | 10 ++- vllm/model_executor/models/mimo_v2.py | 82 +++++++++++++++++++ 4 files changed, 100 insertions(+), 1 deletion(-) diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..341554704dbd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -204,6 +204,7 @@ VLLM_MTP_NVFP4_LM_HEAD: bool = True VLLM_QWEN3_8_FLASH_NEXT_OVERLAP: bool = True VLLM_QWEN3_8_FLASH_NEXT_HC_TP: bool = True + VLLM_MIMO_L2_PREFETCH: bool = False VLLM_B12X_MLA_CKV_GATHER: bool = False VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 @@ -1698,6 +1699,9 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_QWEN3_8_FLASH_NEXT_HC_TP": lambda: bool( int(os.getenv("VLLM_QWEN3_8_FLASH_NEXT_HC_TP", "1")) ), + # MiMo-V2: prefetch upcoming decode weights into L2 on a side stream + # during FULL CUDA-graph decode (cache hints only; numerics unchanged). + "VLLM_MIMO_L2_PREFETCH": lambda: os.getenv("VLLM_MIMO_L2_PREFETCH", "0") == "1", # Gather DCP-sharded C4 records before B12X sparse-MLA prefill. This avoids # query replication plus the per-rank LSE combine and is opt-in while the # path is being qualified on GLM5Next. diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index d69e9ae9d1cb..56a4f0b4878a 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -799,6 +799,11 @@ def unified_attention_with_output( del kv_cache_dummy_dep layer_name = _resolve_layer_name(layer_name) attn_metadata, self, kv_cache, _ = get_attention_context(layer_name) + # Optional model-installed callback (L2 weight prefetch windows); inside + # this opaque op so torch.compile never traces it. + prefetch = getattr(self, "_l2_prefetch_hook", None) + if prefetch is not None: + prefetch(query.shape[0]) self.impl.forward( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 3dbb76215c2f..71695fbde772 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -964,6 +964,7 @@ def _forward_impl( router_logits, ) + num_tokens = hidden_states.shape[0] shared_output, hidden_states = self._apply_quant_method( hidden_states=hidden_states, router_logits=router_logits, @@ -972,10 +973,17 @@ def _forward_impl( shared_experts_overlapping=shared_experts_overlapping, ) - return self._maybe_combine( + output = self._maybe_combine( shared_output, hidden_states, ) + # Optional model-installed callback fired inside the opaque MoE op, + # after the experts and before the final all-reduce (L2 weight + # prefetch for torch.compiled models). + prefetch = getattr(self, "_l2_prefetch_post_experts_hook", None) + if prefetch is not None: + prefetch(num_tokens) + return output ######################################################### # diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 78af26326ca9..d75c932db5ac 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -6,6 +6,7 @@ import torch from torch import nn +import vllm.envs as envs from vllm.compilation.decorators import support_torch_compile from vllm.config import ( CacheConfig, @@ -545,6 +546,83 @@ def _shard_fp8_qkv_proj( ) +_L2_PREFETCH_WINDOW_BYTES = 20 * 1000 * 1000 + + +def _in_full_decode_capture() -> bool: + from vllm.config import CUDAGraphMode + from vllm.forward_context import get_forward_context, is_forward_context_available + + return torch.cuda.is_current_stream_capturing() and not ( + is_forward_context_available() + and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + + +def _install_l2_prefetch_windows( + layer: "MiMoV2FlashDecoderLayer", nxt: "MiMoV2FlashDecoderLayer | None" +) -> None: + """Prefetch upcoming decode weights into L2 while device memory idles. + + Window A (attention op): this layer's o_proj, router and post-attention + norm. Window C (end of the MoE op, before its all-reduce): the next layer's + input norm and qkv projection; the last layer rejoins the side stream. + Prefetches are issued only inside FULL CUDA-graph decode captures and are + cache hints: numerics are unchanged. + """ + from vllm.models.glm5next.nvidia import l2_prefetch + + plans: dict[str, object] = {} + + def build() -> None: + device = layer.post_attention_layernorm.weight.device + small = 16 + segments = l2_prefetch.segments_of(layer.self_attn.o_proj, "o_proj.") + if layer.is_layer_sparse: + segments += l2_prefetch.segments_of( + layer.mlp.gate, "gate.", min_bytes=small + ) + segments += l2_prefetch.segments_of( + layer.post_attention_layernorm, "post_norm.", min_bytes=small + ) + plans["a"] = l2_prefetch.make_plan(segments, _L2_PREFETCH_WINDOW_BYTES, device)[ + 0 + ] + plans["c"] = None + if nxt is not None: + segments = l2_prefetch.segments_of( + nxt.input_layernorm, "next.input_norm.", min_bytes=small + ) + l2_prefetch.segments_of(nxt.self_attn.qkv_proj, "next.qkv_proj.") + plans["c"] = l2_prefetch.make_plan( + segments, _L2_PREFETCH_WINDOW_BYTES, device + )[0] + l2_prefetch.L2Prefetcher.get(device) + + def window(name: str, final: bool): + def hook(num_tokens: int) -> None: + if not plans: + if torch.cuda.is_current_stream_capturing(): + return + build() + if not _in_full_decode_capture(): + return + device = layer.post_attention_layernorm.weight.device + prefetcher = l2_prefetch.L2Prefetcher.get(device) + prefetcher.issue(plans[name], num_tokens) + if final: + prefetcher.join() + + return hook + + object.__setattr__(layer.self_attn.attn, "_l2_prefetch_hook", window("a", False)) + if layer.is_layer_sparse: + object.__setattr__( + layer.mlp.experts, + "_l2_prefetch_post_experts_hook", + window("c", nxt is None), + ) + + @support_torch_compile class MiMoV2Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -579,6 +657,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + if envs.VLLM_MIMO_L2_PREFETCH: + layers = list(islice(self.layers, self.start_layer, self.end_layer)) + for layer, nxt in zip(layers, [*layers[1:], None]): + _install_l2_prefetch_windows(layer, nxt) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size From b85ea198050dec9a2402423db4d91ad5658db514 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:42:55 -0400 Subject: [PATCH 204/365] feat(dflash): shard the context projection and draft per vocab shard at TP > 1 Two per-step costs of DFlash under tensor parallelism: - The context projection fc (target aux hidden states -> draft hidden, e.g. 20480 x 4096) was a ReplicatedLinear, so every rank streamed the whole weight each step. It is now a ColumnParallelLinear with gather_output=True at TP > 1: each output keeps its full dot product (bit-identical), each rank reads 1/tp of the weight. MiMo-V2.6-Flash TP4 + DFlash K7: C1 step time -60 us, +20K tokens of KV capacity. - Probabilistic drafting all-gathered the drafter's [rows, vocab] logits (224 x 152576 BF16 at C32 K7, ~1.2 ms of a ~33 ms PCIe TP4 step). With VLLM_DFLASH_VOCAB_PARALLEL_DRAFT=1 each rank Gumbel-samples its vocab shard with the same absolute-token-id noise, and the per-rank winners are merged; verification uses the drawn token's logit and a global logsumexp combined from per-rank statistics, and rejection resampling runs per shard the same way. The all-gathers shrink to [rows, 5] and [reqs, 2]. It applies to plain vocab-parallel heads, standard rejection sampling and no watermarking; anything else keeps the full-vocab path. Default off. Tests: tests/v1/spec_decode/test_vocab_parallel_drafting.py simulates TP4 shards against the full-vocab kernels: identical draft tokens and cache shards, logsumexp within 1e-5 of FP64, identical rejection sampling (accept lengths, resampled and bonus tokens) across greedy and 0.6-1.3 temperatures, CUDA-graph pad rows and verification depths 0-7. MiMo-V2.6-Flash TP4 C32: 2182 -> 2263 tok/s. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../test_vocab_parallel_drafting.py | 190 ++++++++ vllm/envs.py | 6 + vllm/model_executor/models/qwen3_dflash.py | 23 +- .../gpu/spec_decode/dflash/speculator.py | 127 ++++- .../spec_decode/rejection_sampler_utils.py | 147 ++++-- .../worker/gpu/spec_decode/vocab_parallel.py | 432 ++++++++++++++++++ 6 files changed, 895 insertions(+), 30 deletions(-) create mode 100644 tests/v1/spec_decode/test_vocab_parallel_drafting.py create mode 100644 vllm/v1/worker/gpu/spec_decode/vocab_parallel.py diff --git a/tests/v1/spec_decode/test_vocab_parallel_drafting.py b/tests/v1/spec_decode/test_vocab_parallel_drafting.py new file mode 100644 index 000000000000..4248c26ad818 --- /dev/null +++ b/tests/v1/spec_decode/test_vocab_parallel_drafting.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Vocab-parallel draft sampling must reproduce the full-vocab path exactly. + +TP shards are simulated in one process: each shard runs the local kernels and +the per-shard results are stacked where the all-gather would put them. +""" + +import pytest +import torch + +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode import rejection_sampler_utils as rs +from vllm.v1.worker.gpu.spec_decode import vocab_parallel as vp + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +V, TP, STEPS, MAX_REQS = 16384 * 3, 4, 7, 8 +VL = V // TP + + +def _draft(trial: int, dtype: torch.dtype): + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(trial) + reqs = 6 + slots = torch.tensor([3, 0, 5, 1, 7, 2], device=dev, dtype=torch.int32) + temps = torch.zeros(MAX_REQS, device=dev) + temps[slots.long()] = torch.tensor([1.0, 0.7, 1.3, 0.0, 1.0, 0.6], device=dev) + seeds = torch.randint(0, 2**62, (MAX_REQS,), generator=g, device=dev) + rows = reqs * STEPS + 5 # plus CUDA-graph padding rows + idx = torch.full((rows,), -1, dtype=torch.int32, device=dev) + idx[: reqs * STEPS] = slots.repeat_interleave(STEPS) + cols = torch.arange(STEPS, device=dev, dtype=torch.int32).repeat(rows // STEPS + 1) + cols = cols[:rows].contiguous() + pos = torch.randint(10, 100000, (rows,), generator=g, device=dev) + logits = torch.randn((rows, V), generator=g, device=dev) * 2.5 + logits[:, :64] += 6.0 * torch.rand((rows, 64), generator=g, device=dev) + logits = logits.to(torch.bfloat16) + + full = torch.zeros((MAX_REQS, STEPS, V), dtype=dtype, device=dev) + ref = gumbel_sample( + logits, + idx, + temps, + seeds, + pos, + apply_temperature=True, + is_drafting=True, + logits_cache=full, + logits_cache_col=cols, + ) + caches = [ + vp.VPDraftCache( + torch.zeros((MAX_REQS, STEPS, VL), dtype=dtype, device=dev), r * VL, V, None + ) + for r in range(TP) + ] + stats = torch.stack( + [ + vp.draft_local( + logits[:, r * VL : (r + 1) * VL].contiguous(), + idx, + temps, + seeds, + pos, + cols, + caches[r], + ) + for r in range(TP) + ] + ) + for c in caches: + tokens, _, _ = vp.draft_combine(stats, idx, cols, c) + return dict( + slots=slots, + temps=temps, + seeds=seeds, + idx=idx, + cols=cols, + logits=logits, + full=full, + ref=ref, + caches=caches, + tokens=tokens, + ) + + +@pytest.mark.parametrize("trial", range(4)) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_vocab_parallel_draft_matches_full_vocab(trial, dtype): + d = _draft(trial, dtype) + valid = d["idx"] >= 0 + assert torch.equal(d["tokens"][valid], d["ref"][valid]) + for r, c in enumerate(d["caches"]): + assert torch.equal(c.logits, d["full"][:, :, r * VL : (r + 1) * VL]) + t = d["temps"][d["idx"][valid].long()] + t = torch.where(t > 0, t, torch.ones_like(t)).double() + lse_ref = torch.logsumexp(d["logits"][valid].double() / t[:, None], dim=-1) + rows = (d["idx"][valid].long(), d["cols"][valid].long()) + torch.testing.assert_close( + d["caches"][0].lse[rows].double(), lse_ref, atol=1e-5, rtol=0 + ) + tok_logit = ( + d["logits"][valid].float().gather(-1, d["ref"][valid][:, None]).squeeze(-1) + ) + assert torch.equal(d["caches"][0].token_logit[rows], tok_logit) + + +@pytest.mark.parametrize("trial", range(4)) +def test_vocab_parallel_rejection_matches_full_vocab(trial, monkeypatch): + d = _draft(100 + trial, torch.float32) + dev = d["logits"].device + g = torch.Generator(device=dev).manual_seed(trial) + reqs = d["slots"].numel() + depths = [7, 3, 0, 5, 7, 2] + draft_tok = d["tokens"][: reqs * STEPS].view(reqs, STEPS) + cu = [0] + for n in depths: + cu.append(cu[-1] + n + 1) + num_logits = cu[-1] + draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=dev) + exp_idx = torch.zeros(num_logits, dtype=torch.int32, device=dev) + exp_local = torch.zeros(num_logits, dtype=torch.int32, device=dev) + pos = torch.zeros(num_logits, dtype=torch.int64, device=dev) + target = torch.randn((num_logits, V), generator=g, device=dev) * 2.5 + for q, n in enumerate(depths): + s0 = cu[q] + draft_sampled[s0 + 1 : s0 + 1 + n] = draft_tok[q, :n] + exp_idx[s0 : s0 + n + 1] = d["slots"][q] + exp_local[s0 : s0 + n + 1] = torch.arange(n + 1, dtype=torch.int32, device=dev) + pos[s0 : s0 + n + 1] = torch.arange(n + 1, device=dev) + 5000 + 97 * q + for i in range(n + 1): + src = q * STEPS + min(i, STEPS - 1) + target[s0 + i] = 0.7 * d["logits"][src].float() + 0.3 * target[s0 + i] + cu_t = torch.tensor(cu, dtype=torch.int32, device=dev) + args = ( + draft_sampled, + cu_t, + pos, + d["slots"], + exp_idx, + exp_local, + d["temps"], + d["seeds"], + STEPS, + ) + ref_sampled, ref_n = rs.rejection_sample(target, d["full"], *args) + + caches = d["caches"] + + def resample( + target_logits, + lse, + step, + cu_num, + eidx, + dsamp, + temp, + seed, + p, + cache, + use_fp64=False, + ): + stats = torch.stack( + [ + vp.resample_local( + target_logits, + lse, + step, + cu_num, + eidx, + dsamp, + temp, + seed, + p, + c, + use_fp64, + ) + for c in caches + ] + ) + return vp.resample_combine(stats, use_fp64) + + monkeypatch.setattr(vp, "resample", resample) + for c in caches: + sampled, n = rs.rejection_sample(target, c.logits, *args) + assert torch.equal(n, ref_n) + for q in range(reqs): + k = int(ref_n[q]) + assert torch.equal(sampled[q, :k], ref_sampled[q, :k]) diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..728f49dc1d63 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -48,6 +48,7 @@ VLLM_TRACE_FUNCTION: int = 0 VLLM_USE_FLASHINFER_SAMPLER: bool = True VLLM_GLM53_MTP_DRAFT_HEAD: Literal["bf16", "nvfp4"] = "bf16" + VLLM_DFLASH_VOCAB_PARALLEL_DRAFT: bool = False VLLM_PP_LAYER_PARTITION: str | None = None VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" @@ -896,6 +897,11 @@ def _resolve_rust_cli_path() -> str | None: ["bf16", "nvfp4"], case_sensitive=False, ), + # Sample probabilistic DFlash drafts per vocab shard (exact: same draws and + # verification outcomes) instead of all-gathering the draft logits. + "VLLM_DFLASH_VOCAB_PARALLEL_DRAFT": lambda: ( + os.getenv("VLLM_DFLASH_VOCAB_PARALLEL_DRAFT", "0") == "1" + ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), # (CPU backend only) CPU key-value cache space. diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3b5b01f5afb2..dd5eab512f55 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, @@ -476,7 +477,13 @@ def __init__( ] ) if self.use_aux_hidden_state: - self.fc = ReplicatedLinear( + # Shard the context projection's output features across TP ranks + # and all-gather the result: every output element keeps its full + # dot product (bit-identical to a replicated layer) while each rank + # streams 1/tp of the [hidden, num_aux * hidden] weight per step. + tp_fc = get_tensor_model_parallel_world_size() > 1 + fc_cls = ColumnParallelLinear if tp_fc else ReplicatedLinear + self.fc = fc_cls( input_size=_get_dflash_fc_input_size( vllm_config, ), @@ -486,6 +493,7 @@ def __init__( quant_config=self.quant_config, prefix=maybe_prefix(prefix, "fc"), return_bias=False, + **({"gather_output": True} if tp_fc else {}), ) self.hidden_norm = RMSNorm( self.config.hidden_size, @@ -886,6 +894,19 @@ def compute_logits( logits_new[:, targets] = logits return logits_new + def compute_local_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """This rank's vocab-shard logits, not gathered, for vocab-parallel + drafting (the speculator validates the shard layout: no vocab padding, + no added vocab, no draft-to-target id map).""" + lp = self.logits_processor + logits = lp._apply_head(self.lm_head, hidden_states, None) + logits = logits[..., : self.lm_head.shard_indices.num_org_elements] + if lp.soft_cap is not None: + logits = torch.tanh(logits / lp.soft_cap) * lp.soft_cap + if lp.scale != 1.0: + logits = logits * lp.scale + return logits + def precompute_and_store_context_kv( self, context_states: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 8ae1ec9014b0..e209f6829412 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context @@ -183,7 +184,131 @@ def load_draft_model( target_model: nn.Module, target_attn_layer_names: set[str], ) -> nn.Module: - return load_dflash_model(target_model, self.vllm_config) + model = load_dflash_model(target_model, self.vllm_config) + self._setup_vocab_parallel_drafting(model) + return model + + def _setup_vocab_parallel_drafting(self, model: nn.Module) -> None: + """Keep draft logits vocab-sharded (VLLM_DFLASH_VOCAB_PARALLEL_DRAFT=1). + + Probabilistic drafts are Gumbel-sampled per vocab shard and combined + from per-rank winners (spec_decode.vocab_parallel), so the [rows, vocab] + logits all-gather disappears; the cache holds this rank's shard. Any + layout the sampler does not cover keeps the full-vocab path. + """ + from vllm.distributed import get_tensor_model_parallel_world_size + + self._vocab_parallel = None + if not envs.VLLM_DFLASH_VOCAB_PARALLEL_DRAFT or self.draft_logits is None: + return + head = getattr(model, "lm_head", None) + si = getattr(head, "shard_indices", None) + tp = get_tensor_model_parallel_world_size() + if ( + tp == 1 + or si is None + or not hasattr(model, "compute_local_logits") + or getattr(model, "draft_id_to_target_id", None) is not None + or self.draft_watermarker is not None + or self.speculative_config.rejection_sample_method != "standard" + or si.num_org_vocab_padding != 0 + or si.num_added_elements != 0 + or si.num_org_elements * tp != self.vocab_size + ): + logger.info_once( + "DFlash vocab-parallel drafting not applicable; " + "keeping the full-vocab draft path" + ) + return + from vllm.v1.worker.gpu.spec_decode.vocab_parallel import VPDraftCache + + dtype, fill = self.draft_logits_spec(self.vllm_config) + self.draft_logits = torch.full( + (self.max_num_reqs, self.num_speculative_steps, si.num_org_elements), + fill, + dtype=dtype, + device=self.device, + ) + self._vocab_parallel = VPDraftCache( + self.draft_logits, si.org_vocab_start_index, self.vocab_size, None + ) + logger.info_once( + "DFlash vocab-parallel drafting: shard [%d, %d) of %d, no draft-logits " + "all-gather", + si.org_vocab_start_index, + si.org_vocab_end_index, + self.vocab_size, + ) + + def sample_draft( + self, + hidden_states: torch.Tensor, + sample_src_positions: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + draft_step: torch.Tensor, + draft_logits: torch.Tensor | None, + ) -> torch.Tensor: + vp = getattr(self, "_vocab_parallel", None) + if vp is None or draft_logits is None or draft_logits is not vp.logits: + return super().sample_draft( + hidden_states, + sample_src_positions, + idx_mapping, + temperature, + seeds, + draft_step, + draft_logits, + ) + from vllm.v1.worker.gpu.spec_decode import vocab_parallel as vp_sampling + + local = self.model.compute_local_logits(hidden_states) + col = draft_step if draft_step.dim() > 0 else draft_step.expand(local.shape[0]) + tokens, maxes, sums = vp_sampling.draft_sample( + local, + idx_mapping, + temperature, + seeds, + sample_src_positions, + col, + vp, + use_fp64=self.use_fp64_gumbel, + ) + est = self.acceptance_estimator + if est is not None: + # Same predictor as acceptance_estimator.predict, fed with per-rank + # (max, sum-exp) partials of logits / T instead of vocab blocks. + from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( + _MAX_LOG_ODDS, + _predict_kernel, + ) + + num_tokens, nparts = maxes.shape + _predict_kernel[(num_tokens,)]( + est.features, + est.features.stride(0), + est.predictions, + est.predictions.stride(0), + self.draft_token_confidence_probs, + self.draft_token_confidence_probs.stride(0), + est.slope, + est.intercepts, + maxes, + maxes.stride(0), + sums, + sums.stride(0), + idx_mapping, + idx_mapping.stride(0), + draft_step, + num_tokens, + nparts, + per_token_step=draft_step.dim() > 0, + NUM_SPECULATIVE_STEPS=est.num_speculative_steps, + MAX_LOG_ODDS=_MAX_LOG_ODDS, + PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(nparts), + ) + return tokens def set_attn( self, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 9d64c1ecbcf4..f80c88ad52a5 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -534,10 +534,15 @@ def _rejection_kernel( local_residual_mass_ptr, local_residual_mass_stride, vocab_num_blocks, + # Vocab-parallel drafts: [max_num_reqs + 1, num_speculative_steps] + vp_token_logit_ptr, + vp_lse_ptr, + vp_stride, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, SYNTHETIC_MODE: tl.constexpr, USE_BLOCK_VERIFICATION: tl.constexpr, + VP_DRAFT: tl.constexpr = False, ): req_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64) @@ -628,32 +633,69 @@ def _rejection_kernel( tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled) else: # Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192 - target_logprob, draft_logprob, target_lse, draft_lse = ( - _compute_global_logprobs_and_logsumexp( - draft_sampled, - True, # mask - logit_idx, - req_state_idx, - i, - temp, - target_logits_ptr, - target_logits_stride, - target_local_max_ptr, - target_local_max_stride, - target_local_sumexp_ptr, - target_local_sumexp_stride, - draft_logits_ptr, - draft_logits_stride_0, - draft_logits_stride_1, - draft_local_max_ptr, - draft_local_max_stride, - draft_local_sumexp_ptr, - draft_local_sumexp_stride, - vocab_num_blocks, - PADDED_VOCAB_NUM_BLOCKS, - HAS_DRAFT_LOGITS, + if VP_DRAFT: + # Vocab-parallel drafts: q(d) from the drawn token's + # pre-temperature logit and the global logsumexp, both + # recorded when the draft was sampled. + target_logprob, _, target_lse, _ = ( + _compute_global_logprobs_and_logsumexp( + draft_sampled, + True, # mask + logit_idx, + req_state_idx, + i, + temp, + target_logits_ptr, + target_logits_stride, + target_local_max_ptr, + target_local_max_stride, + target_local_sumexp_ptr, + target_local_sumexp_stride, + draft_logits_ptr, + draft_logits_stride_0, + draft_logits_stride_1, + draft_local_max_ptr, + draft_local_max_stride, + draft_local_sumexp_ptr, + draft_local_sumexp_stride, + vocab_num_blocks, + PADDED_VOCAB_NUM_BLOCKS, + False, + ) + ) + draft_lse = tl.load(vp_lse_ptr + req_state_idx * vp_stride + i) + draft_logprob = ( + tl.load(vp_token_logit_ptr + req_state_idx * vp_stride + i) + / temp + - draft_lse + ) + else: + target_logprob, draft_logprob, target_lse, draft_lse = ( + _compute_global_logprobs_and_logsumexp( + draft_sampled, + True, # mask + logit_idx, + req_state_idx, + i, + temp, + target_logits_ptr, + target_logits_stride, + target_local_max_ptr, + target_local_max_stride, + target_local_sumexp_ptr, + target_local_sumexp_stride, + draft_logits_ptr, + draft_logits_stride_0, + draft_logits_stride_1, + draft_local_max_ptr, + draft_local_max_stride, + draft_local_sumexp_ptr, + draft_local_sumexp_stride, + vocab_num_blocks, + PADDED_VOCAB_NUM_BLOCKS, + HAS_DRAFT_LOGITS, + ) ) - ) if SYNTHETIC_MODE: rate = tl.load(synthetic_conditional_rates_ptr + i) accepted = u < rate @@ -1040,14 +1082,24 @@ def rejection_sample( watermarking_bytes = watermarking.view(torch.uint8) watermark_key_0 = watermark_key & 0xFFFFFFFF watermark_key_1 = watermark_key >> 32 + from vllm.v1.worker.gpu.spec_decode import vocab_parallel as vp_sampling + + vp = vp_sampling.lookup(draft_logits) + if vp is not None: + assert not use_block_verification, ( + "vocab-parallel drafts: standard verification only" + ) + assert contexts is None, "vocab-parallel drafts do not support watermarking" draft_logits_stride_0 = 0 draft_logits_stride_1 = 0 if has_draft_logits := draft_logits is not None: draft_logits_stride_0 = draft_logits.stride(0) draft_logits_stride_1 = draft_logits.stride(1) # In some cases (e.g. MiMo v2.5 Pro + DFlash) the target model's - # vocab size is larger than the draft's due to padding. - vocab_size = min(vocab_size, draft_logits.size(-1)) + # vocab size is larger than the draft's due to padding. A + # vocab-parallel draft cache holds one shard; the target stays full. + if vp is None: + vocab_size = min(vocab_size, draft_logits.size(-1)) # Compute the per-vocab-block logits stats, such as target argmax # (for greedy requests), and target max + softmax exponential @@ -1092,7 +1144,8 @@ def rejection_sample( vocab_size, num_speculative_steps, BLOCK_SIZE=VOCAB_BLOCK_SIZE, - HAS_DRAFT_LOGITS=has_draft_logits, + # Vocab-parallel drafts carry their logsumexp; only target stats here. + HAS_DRAFT_LOGITS=has_draft_logits and vp is None, ) # Precompute the running joint ratio and residual mass for block @@ -1212,13 +1265,51 @@ def rejection_sample( local_residual_mass, local_residual_mass.stride(0) if local_residual_mass is not None else 0, vocab_num_blocks, + vp.token_logit if vp is not None else target_rejected_logsumexp, + vp.lse if vp is not None else target_rejected_logsumexp, + vp.lse.stride(0) if vp is not None else 0, PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks, HAS_DRAFT_LOGITS=has_draft_logits, SYNTHETIC_MODE=synthetic_conditional_rates is not None, USE_BLOCK_VERIFICATION=use_block_verification, + VP_DRAFT=vp is not None, num_warps=1, ) + if vp is not None: + # Each rank Gumbel-samples its shard of the residual (or of the target + # at the bonus position); the per-rank winners play the role of the + # vocab blocks of _insert_resampled_kernel. + resampled_local_max, resampled_local_argmax = vp_sampling.resample( + target_logits, + target_rejected_logsumexp, + num_sampled, + cu_num_logits, + expanded_idx_mapping, + draft_sampled, + temperature, + seed, + pos, + vp, + use_fp64=use_fp64, + ) + resample_num_blocks = resampled_local_max.shape[1] + _insert_resampled_kernel[(num_reqs,)]( + sampled, + sampled.stride(0), + num_sampled, + resampled_local_argmax, + resampled_local_argmax.stride(0), + resampled_local_max, + resampled_local_max.stride(0), + resample_num_blocks, + cu_num_logits, + expanded_idx_mapping, + temperature, + PADDED_RESAMPLE_NUM_BLOCKS=triton.next_power_of_2(resample_num_blocks), + ) + return sampled, num_sampled + # Resample the rejected/bonus tokens. RESAMPLE_BLOCK_SIZE = 1024 resample_num_blocks = triton.cdiv(vocab_size, RESAMPLE_BLOCK_SIZE) diff --git a/vllm/v1/worker/gpu/spec_decode/vocab_parallel.py b/vllm/v1/worker/gpu/spec_decode/vocab_parallel.py new file mode 100644 index 000000000000..db0fbc14c26d --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/vocab_parallel.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Vocab-parallel draft sampling for speculative decoding (no logits all-gather). + +The drafter's lm_head is vocab-parallel: each TP rank holds logits for its +contiguous vocab shard. Gathering them (``[rows, vocab]`` per step) was the +largest collective of a C32 step. Sampling needs far less: + + * Gumbel-max draws key the noise by absolute token id (``gumbel.py``), so + each rank draws over its shard with the same noise the full-vocab kernel + uses, and the global argmax of the per-rank winners is the token the + full-vocab draw returns (ties aside, which have probability zero). + * Verification needs q(d) for the drawn token d and, on rejection, the + residual max(p - q, 0) over the vocabulary. q(d) comes from d's logit + (known to the winning rank) and the global normalizer (combined from + per-rank max / sum-exp); both are cached per (request, step). The + residual is formed and Gumbel-sampled per shard, then combined the same way. + +Every rank ends with identical tokens and statistics. Pre-temperature draft +logits stay cached per shard for the residual. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, tldevice, triton +from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax + +BLOCK = 1024 + +# draft-logits cache (sharded) data_ptr -> VPDraftCache +_REGISTRY: dict[int, VPDraftCache] = {} + + +class VPDraftCache: + """Side caches of a sharded draft-logits cache [reqs, steps, local_vocab].""" + + def __init__( + self, logits: torch.Tensor, vocab_start: int, vocab_size: int, tp_group + ): + reqs, steps, _ = logits.shape + self.logits = logits + self.vocab_start = int(vocab_start) + self.local_vocab = int(logits.shape[-1]) + self.vocab_size = int(vocab_size) + self.tp_group = tp_group + dev = logits.device + # Pre-temperature logit of the drawn token, and logsumexp(logits / T). + # Row ``reqs`` is a sink for CUDA-graph padding rows (slot -1). + self.num_reqs = int(reqs) + self.token_logit = torch.zeros( + (reqs + 1, steps), dtype=torch.float32, device=dev + ) + self.lse = torch.zeros((reqs + 1, steps), dtype=torch.float32, device=dev) + _REGISTRY[logits.data_ptr()] = self + + +def lookup(draft_logits: torch.Tensor | None) -> VPDraftCache | None: + if draft_logits is None: + return None + return _REGISTRY.get(draft_logits.data_ptr()) + + +@triton.jit +def _vp_draft_block_kernel( + blk_val_ptr, # [rows, nblk] fp32/fp64: noisy max + blk_idx_ptr, # [rows, nblk] int64: absolute token id + blk_max_ptr, # [rows, nblk] fp32: max(logit / T) + blk_sum_ptr, # [rows, nblk] fp32: sum exp(logit / T - blk_max) + nblk, + cache_ptr, # [reqs, steps, local_vocab] + cache_stride_0, + cache_stride_1, + col_ptr, # [rows] step column + logits_ptr, # [rows, local_vocab] + logits_stride, + idx_mapping_ptr, # [rows] request slot (-1: padding) + seeds_ptr, + pos_ptr, + temp_ptr, + local_vocab, + vocab_start, + BLOCK_SIZE: tl.constexpr, + USE_FP64: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + b = tl.program_id(1) + offs = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < local_vocab + logits = tl.load( + logits_ptr + row * logits_stride + offs, mask=mask, other=float("-inf") + ) + logits = logits.to(tl.float32) + req = tl.load(idx_mapping_ptr + row).to(tl.int64) + valid = req >= 0 + temp = tl.load(temp_ptr + req, mask=valid, other=0.0).to(tl.float32) + col = tl.load(col_ptr + row).to(tl.int64) + # Cache the pre-temperature logits (the rejection sampler divides by the + # same temperature on load, exactly as for the full-vocab cache). + tl.store( + cache_ptr + req * cache_stride_0 + col * cache_stride_1 + offs, + logits, + mask=mask & valid, + ) + seed = tl.load(seeds_ptr + req, mask=valid, other=0) + pos = tl.load(pos_ptr + row) + keys = vocab_start + offs + value, idx = gumbel_noised_argmax( + logits, + keys, + mask, + seed, + pos, + temp, + IS_DRAFTING=True, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=True, + ) + tl.store(blk_val_ptr + row * nblk + b, value) + tl.store( + blk_idx_ptr + row * nblk + b, (vocab_start + b * BLOCK_SIZE + idx).to(tl.int64) + ) + t = tl.where(temp > 0.0, temp, 1.0) + scaled = tl.where(mask, logits / t, float("-inf")) + m = tl.max(scaled, axis=0) + s = tl.where(m > float("-inf"), tl.sum(tl.exp(scaled - m)), 0.0) + tl.store(blk_max_ptr + row * nblk + b, m) + tl.store(blk_sum_ptr + row * nblk + b, s) + + +def draft_local( + local_logits: torch.Tensor, # [rows, local_vocab] + idx_mapping: torch.Tensor, # [rows] + temperature: torch.Tensor, # [max_reqs] + seeds: torch.Tensor, # [max_reqs] + pos: torch.Tensor, # [rows] + col: torch.Tensor, # [rows] + cache: VPDraftCache, + use_fp64: bool = False, +) -> torch.Tensor: + """This shard's part: caches the shard's logits and returns per-row + [noisy max, token id, its logit, max(logit/T), sumexp] as fp64 [rows, 5].""" + rows, local_vocab = local_logits.shape + nblk = triton.cdiv(local_vocab, BLOCK) + dev = local_logits.device + val_dtype = torch.float64 if use_fp64 else torch.float32 + blk_val = torch.empty((rows, nblk), dtype=val_dtype, device=dev) + blk_idx = torch.empty((rows, nblk), dtype=torch.int64, device=dev) + blk_max = torch.empty((rows, nblk), dtype=torch.float32, device=dev) + blk_sum = torch.empty((rows, nblk), dtype=torch.float32, device=dev) + _vp_draft_block_kernel[(rows, nblk)]( + blk_val, + blk_idx, + blk_max, + blk_sum, + nblk, + cache.logits, + cache.logits.stride(0), + cache.logits.stride(1), + col.contiguous(), + local_logits, + local_logits.stride(0), + idx_mapping.contiguous(), + seeds, + pos.contiguous(), + temperature, + local_vocab, + cache.vocab_start, + BLOCK_SIZE=BLOCK, + USE_FP64=use_fp64, + ) + # Per-shard winner (first max, i.e. lowest token id among ties). + best_blk = blk_val.argmax(dim=-1, keepdim=True) + best_val = blk_val.gather(-1, best_blk).squeeze(-1).to(torch.float64) + best_id = blk_idx.gather(-1, best_blk).squeeze(-1) + best_logit = local_logits.gather( + -1, (best_id - cache.vocab_start).unsqueeze(-1) + ).squeeze(-1) + row_max = blk_max.max(dim=-1).values + row_sum = ( + (blk_sum * torch.exp(blk_max - row_max.unsqueeze(-1))).nan_to_num_(0.0).sum(-1) + ) + # token ids < 2**53 are exact in fp64 + return torch.stack( + ( + best_val, + best_id.to(torch.float64), + best_logit.to(torch.float64), + row_max.to(torch.float64), + row_sum.to(torch.float64), + ), + dim=-1, + ) + + +def draft_combine( + gathered: torch.Tensor, # [tp, rows, 5] in vocab-shard order + idx_mapping: torch.Tensor, + col: torch.Tensor, + cache: VPDraftCache, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Global winner and normalizer; records (logit, logsumexp) per slot. + Returns (tokens, per-shard max [rows, tp], per-shard sumexp [rows, tp]).""" + tp, rows, _ = gathered.shape + winner = gathered[:, :, 0].argmax(dim=0) # first max = lowest vocab shard + pick = gathered.gather(0, winner.view(1, rows, 1).expand(1, rows, 5)).squeeze(0) + tokens = pick[:, 1].to(torch.int64) + tok_logit = pick[:, 2].to(torch.float32) + maxes = gathered[:, :, 3] # [tp, rows] + sums = gathered[:, :, 4] + gmax = maxes.max(dim=0).values + total = (sums * torch.exp(maxes - gmax)).nan_to_num_(0.0).sum(dim=0) + lse = (gmax + torch.log(total)).to(torch.float32) + req = idx_mapping.to(torch.int64) + slot = torch.where(req >= 0, req, torch.full_like(req, cache.num_reqs)) + colv = col.to(torch.int64) + cache.token_logit[slot, colv] = tok_logit + cache.lse[slot, colv] = lse + return ( + tokens, + maxes.t().to(torch.float32).contiguous(), + sums.t().to(torch.float32).contiguous(), + ) + + +def draft_sample( + local_logits, + idx_mapping, + temperature, + seeds, + pos, + col, + cache, + use_fp64: bool = False, +): + """Sample one draft token per row over the TP-sharded vocabulary: see + draft_local / draft_combine; one tiny all-gather instead of the logits'.""" + from vllm.distributed import tensor_model_parallel_all_gather + + stats = draft_local( + local_logits, idx_mapping, temperature, seeds, pos, col, cache, use_fp64 + ) + gathered = tensor_model_parallel_all_gather(stats.unsqueeze(0), dim=0) + return draft_combine(gathered, idx_mapping, col, cache) + + +@triton.jit +def _vp_resample_block_kernel( + out_val_ptr, # [reqs, nblk] + out_idx_ptr, # [reqs, nblk] absolute token id + nblk, + target_logits_ptr, # [num_logits, V] full (processed) + target_logits_stride, + target_lse_ptr, # [reqs] (rejected position), 0 if unused + draft_cache_ptr, # [max_reqs, steps, local_vocab] pre-temperature + draft_stride_0, + draft_stride_1, + draft_lse_cache_ptr, # [max_reqs, steps] + draft_lse_stride, + rejected_step_ptr, # [reqs] + cu_num_logits_ptr, # [reqs + 1] + expanded_idx_mapping_ptr, # [num_logits] + draft_sampled_ptr, # [num_logits] + temp_ptr, + seed_ptr, + pos_ptr, + local_vocab, + vocab_start, + BLOCK_SIZE: tl.constexpr, + USE_FP64: tl.constexpr, +): + """One shard's part of _resample_kernel (standard verification, draft + logits present): the residual log(max(p - q, 0)) at the rejected position, + or the target logits at the bonus position, Gumbel-sampled per block.""" + r = tl.program_id(0) + b = tl.program_id(1) + resample_idx = tl.load(rejected_step_ptr + r) + start = tl.load(cu_num_logits_ptr + r).to(tl.int64) + end = tl.load(cu_num_logits_ptr + r + 1) + tok_idx = start + resample_idx + req = tl.load(expanded_idx_mapping_ptr + tok_idx).to(tl.int64) + temp = tl.load(temp_ptr + req).to(tl.float32) + is_bonus = tok_idx == end - 1 + offs = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < local_vocab + keys = vocab_start + offs + if temp == 0.0 and not is_bonus: + tl.store(out_val_ptr + r * nblk + b, float("-inf")) + tl.store( + out_idx_ptr + r * nblk + b, (vocab_start + b * BLOCK_SIZE).to(tl.int64) + ) + return + rejected = tl.load(draft_sampled_ptr + tok_idx + 1, mask=not is_bonus, other=0) + target = tl.load( + target_logits_ptr + tok_idx * target_logits_stride + keys, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + if is_bonus or rejected < 0: + residual = target + else: + draft = ( + tl.load( + draft_cache_ptr + + req * draft_stride_0 + + resample_idx * draft_stride_1 + + offs, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + / temp + ) + t_lse = tl.load(target_lse_ptr + r) + d_lse = tl.load(draft_lse_cache_ptr + req * draft_lse_stride + resample_idx) + t_lp = target - t_lse + ratio = tl.exp((draft - d_lse) - t_lp) + residual = tl.where( + ratio < 1.0, t_lp + tldevice.log1p(-ratio), float("-inf") + ).to(tl.float32) + residual = tl.where(mask, residual, float("-inf")) + seed = tl.load(seed_ptr + req) + pos = tl.load(pos_ptr + tok_idx) + value, idx = gumbel_noised_argmax( + residual, + keys, + mask, + seed, + pos, + temp, + IS_DRAFTING=False, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, + ) + tl.store(out_val_ptr + r * nblk + b, value) + tl.store( + out_idx_ptr + r * nblk + b, (vocab_start + b * BLOCK_SIZE + idx).to(tl.int64) + ) + + +def resample_local( + target_logits: torch.Tensor, + target_rejected_lse: torch.Tensor, + rejected_step: torch.Tensor, + cu_num_logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + draft_sampled: torch.Tensor, + temperature: torch.Tensor, + seed: torch.Tensor, + pos: torch.Tensor, + cache: VPDraftCache, + use_fp64: bool = False, +) -> torch.Tensor: + """This shard's resample winner per request: fp64 [reqs, 2] (value, id).""" + reqs = int(cu_num_logits.shape[0]) - 1 + nblk = triton.cdiv(cache.local_vocab, BLOCK) + dev = target_logits.device + val_dtype = torch.float64 if use_fp64 else torch.float32 + val = torch.empty((reqs, nblk), dtype=val_dtype, device=dev) + idx = torch.empty((reqs, nblk), dtype=torch.int64, device=dev) + _vp_resample_block_kernel[(reqs, nblk)]( + val, + idx, + nblk, + target_logits, + target_logits.stride(0), + target_rejected_lse, + cache.logits, + cache.logits.stride(0), + cache.logits.stride(1), + cache.lse, + cache.lse.stride(0), + rejected_step, + cu_num_logits, + expanded_idx_mapping, + draft_sampled, + temperature, + seed, + pos, + cache.local_vocab, + cache.vocab_start, + BLOCK_SIZE=BLOCK, + USE_FP64=use_fp64, + ) + best = val.argmax(dim=-1, keepdim=True) + best_val = val.gather(-1, best).squeeze(-1).to(torch.float64) + best_idx = idx.gather(-1, best).squeeze(-1) + return torch.stack((best_val, best_idx.to(torch.float64)), dim=-1) + + +def resample_combine(gathered: torch.Tensor, use_fp64: bool = False): + """[tp, reqs, 2] -> ([reqs, tp] values, [reqs, tp] ids), shard order kept + (the insert kernel takes the first max).""" + val_dtype = torch.float64 if use_fp64 else torch.float32 + return ( + gathered[:, :, 0].t().contiguous().to(val_dtype), + gathered[:, :, 1].t().contiguous().to(torch.int64), + ) + + +def resample( + target_logits, + target_rejected_lse, + rejected_step, + cu_num_logits, + expanded_idx_mapping, + draft_sampled, + temperature, + seed, + pos, + cache, + use_fp64: bool = False, +): + """Per-request (value, token) candidates from every shard: ([reqs, tp], + [reqs, tp]) handed to _insert_resampled_kernel as its vocab blocks.""" + from vllm.distributed import tensor_model_parallel_all_gather + + stats = resample_local( + target_logits, + target_rejected_lse, + rejected_step, + cu_num_logits, + expanded_idx_mapping, + draft_sampled, + temperature, + seed, + pos, + cache, + use_fp64, + ) + gathered = tensor_model_parallel_all_gather(stats.unsqueeze(0), dim=0) + return resample_combine(gathered, use_fp64) From 5a98e01846997c11c94b81307b110c5af0621651 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:43:18 -0400 Subject: [PATCH 205/365] feat(dflash): support block-FP8 drafters in the fused context projection Online `fp8_per_block` quantization (speculative_config.quantization) quantizes the drafter's linears while they load, so the fused context K/V projection concatenated raw FP8 rows without their block scales and F.linear failed on a BF16 x FP8 matmul. Dequantize block-FP8 K/V rows with their scales for the fused projection; each layer keeps its own FP8 GEMMs, and serialized MXFP8 rows stay packed for their fused linear method. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- tests/v1/spec_decode/test_dflash2.py | 36 ++++++++++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 39 +++++++++++++++++++++- 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 10faa8b671a0..07788707a922 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -471,3 +471,39 @@ class MockGroup: # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass) assert len(model.layers) == 2 assert isinstance(model.layers[0], DFlash2Qwen3DecoderLayer) + + +def test_fused_context_projection_dequantizes_block_fp8_rows(): + """Online ``fp8_per_block`` drafters quantize while loading, so the fused + context projection receives FP8 rows; it must run on their dequantized + values (with their block scales), not on raw FP8 bytes.""" + from torch import nn + + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + torch.manual_seed(0) + layers, expected = [], [] + for _ in range(2): + projection = nn.Module() + projection.quant_method = None + scale = torch.rand(2, 2) + 0.5 + full_scale = scale.repeat_interleave(128, 0).repeat_interleave(128, 1) + dense = torch.randn(256, 256) + weight = (dense / full_scale).clamp(-448, 448).to(torch.float8_e4m3fn) + projection.weight = nn.Parameter(weight, requires_grad=False) + projection.weight_scale_inv = nn.Parameter(scale, requires_grad=False) + projection.weight_block_size = [128, 128] + layers.append( + SimpleNamespace( + qkv_proj=projection, + q_size=128, + k_norm=SimpleNamespace(weight=torch.ones(32)), + ) + ) + expected.append((weight.float() * full_scale)[128:].to(torch.bfloat16)) + model = object.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.hidden_norm = SimpleNamespace(weight=torch.ones(256, dtype=torch.bfloat16)) + model._build_context_kv_buffers(layers, has_bias=False) + assert model._fused_kv_weight.dtype == torch.bfloat16 + assert torch.equal(model._fused_kv_weight, torch.cat(expected)) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3b5b01f5afb2..d4d5dc34c72b 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -395,6 +395,38 @@ def forward( return hidden_states, residual +def _context_kv_rows( + attn: nn.Module, dtype: torch.dtype, *, packed_mxfp8: bool +) -> torch.Tensor: + """One layer's K/V projection rows for the fused context projection. + + Serialized MXFP8 rows stay packed for the fused projection's own linear + method. Block-FP8 rows (serialized, or online ``fp8_per_block``, which + quantizes while loading) are dequantized with their block scales, so the + fused projection runs in ``dtype`` while each layer keeps its FP8 GEMM. + """ + projection = attn.qkv_proj + rows = projection.weight[attn.q_size :] + scale = getattr(projection, "weight_scale_inv", None) + if ( + packed_mxfp8 + or scale is None + or rows.dtype + not in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + ): + return rows + block_n, block_k = getattr(projection, "weight_block_size", None) or (128, 128) + if attn.q_size % block_n: + raise ValueError("DFlash block-FP8 K/V rows must start on a scale block.") + scales = scale[attn.q_size // block_n :].float() + scales = scales.repeat_interleave(block_n, 0)[: rows.shape[0]] + scales = scales.repeat_interleave(block_k, 1)[:, : rows.shape[1]] + return (rows.float() * scales).to(dtype) + + @support_torch_compile class DFlashQwen3Model(nn.Module): decoder_layer_cls = DFlashQwen3DecoderLayer @@ -532,7 +564,12 @@ def _build_context_kv_buffers( self._hidden_norm_weight = self.hidden_norm.weight.data # KV projection weights: [num_layers * 2 * kv_size, hidden_size] - kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] + kv_weights = [ + _context_kv_rows( + a, self.hidden_norm.weight.dtype, packed_mxfp8=all(uses_mxfp8) + ) + for a in layers_attn + ] self._fused_kv_weight = torch.cat(kv_weights, dim=0) if all(uses_mxfp8): From addaaa328bac26f865deb35185b4cace8028e5f8 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:43:19 -0400 Subject: [PATCH 206/365] fix(spec-decode): keep the adaptive verification acceptance estimate recoverable The online acceptance estimator learns only from drafts that adaptive verification admits. A draft predicted near-certain carries almost no IRLS weight, so a data-poor refit that saw one confidently wrong draft divided its residual by ~the ridge and moved an intercept by about -20. Predicted acceptance then fell to ~1e-8, the budget admitted no drafts, and nothing was ever graded again: speculation stopped for the rest of a long request (reported as DFlash "falling off" during long thinking). - Bound each refit's coefficient change (trust region) and keep the slope and intercepts within fixed limits, off logistic saturation. - After EXPLORE_AFTER_IDLE_STEPS consecutive steps that schedule drafts but verify none, verify one draft per request for a step so the estimator keeps receiving evidence. Every rank makes the same choice. Only the number of verified drafts changes; rejection sampling keeps the output distribution exact. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../spec_decode/test_acceptance_estimator.py | 23 +++++++++++++++++ .../spec_decode/test_adaptive_verification.py | 19 ++++++++++++++ .../gpu/spec_decode/acceptance_estimator.py | 25 ++++++++++++++++--- .../gpu/spec_decode/adaptive_verification.py | 18 +++++++++++++ 4 files changed, 82 insertions(+), 3 deletions(-) diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py index a1f3da09402f..9ce7963d1e67 100644 --- a/tests/v1/spec_decode/test_acceptance_estimator.py +++ b/tests/v1/spec_decode/test_acceptance_estimator.py @@ -292,3 +292,26 @@ def test_feature_ignores_which_token_was_drawn(): atol=2e-3, rtol=2e-3, ) + + +def test_refit_stays_bounded_after_a_confidently_wrong_draft(): + """One confidently wrong draft must not end drafting. + + A draft predicted near certain carries almost no IRLS weight, so dividing + its residual by that weight plus the small ridge moved the position-0 + intercept by about -20 in one refit. Every later prediction then rounded + to ~1e-8, adaptive verification admitted no drafts, and no graded drafts + arrived to move it back. + """ + device = torch.device("cuda") + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 1 + before = estimator.intercepts.clone() + slots = torch.zeros(1, dtype=torch.int32, device=device) + features = torch.full((1, NUM_STEPS), 30.0, device=device) + _drive_round(estimator, features, torch.zeros(1, device=device), slots) + + typical = torch.sigmoid(estimator.slope * 3.0 + estimator.intercepts[0]) + assert typical.item() > 0.05 + step = (estimator.intercepts - before).abs().max().item() + assert step <= estimator.MAX_STEP + 1e-6 diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index a9578a1b1305..028e3540d27f 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -656,3 +656,22 @@ def test_tp_confidence_publication_keeps_graph_and_request_budgets_consistent(): torch.multiprocessing.spawn( _run_tp_confidence_consistency, args=(get_open_port(),), nprocs=2, join=True ) + + +def test_budget_probes_after_steps_without_drafts(): + """An estimator that stopped trusting the drafter must keep being graded. + + Confidences learned only from admitted drafts cannot recover once no draft + is admitted, so after EXPLORE_AFTER_IDLE_STEPS empty budgets one step + verifies a draft per request, then budgeting resumes as before. + """ + manager = make_manager( + np.array([[1e-9, 1e-9]], dtype=np.float32), + np.array([1.0, 2.0, 3.0, 4.0]), + ) + idle = getattr(adaptive_module, "EXPLORE_AFTER_IDLE_STEPS", 32) + budgets = [] + for _ in range(2 * idle): + manager.get_num_tokens({"low": 3}, {"low": [1, 2]}) + budgets.append(manager._batch_budget[2]) + assert budgets == ([0] * (idle - 1) + [1]) * 2 diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py index 46065d0546e7..88b697aa3743 100644 --- a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py +++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py @@ -100,6 +100,9 @@ def _refit_kernel( NUM_SPECULATIVE_STEPS: tl.constexpr, L2: tl.constexpr, DAMPING: tl.constexpr, + MAX_STEP: tl.constexpr, + MAX_SLOPE: tl.constexpr, + MAX_ABS_INTERCEPT: tl.constexpr, BLOCK: tl.constexpr, ): k = tl.arange(0, BLOCK) @@ -136,8 +139,11 @@ def _refit_kernel( step_w *= total_n / (total_n + DAMPING) # Mask out NaNs and steps that would drive the slope negative. step_w = tl.where((step_w == step_w) & (w + step_w >= 0.0), step_w, 0.0) + # Trust region: near-certain predictions carry almost no IRLS weight, so a + # data-poor round's Newton step can be enormous. + step_w = tl.minimum(tl.maximum(step_w, -MAX_STEP), MAX_STEP) # Update the shared slope. - new_w = w + step_w + new_w = tl.minimum(w + step_w, MAX_SLOPE) # Solve for the per-position intercept deltas, Δβₖ, in the same system. step_bias = (g1_k - b_k * step_w) / c_k @@ -145,8 +151,12 @@ def _refit_kernel( step_bias *= n / (n + DAMPING) # Mask out NaNs. step_bias = tl.where(step_bias == step_bias, step_bias, 0.0) - # Update the per-position intercepts. - new_bias = bias + step_bias + step_bias = tl.minimum(tl.maximum(step_bias, -MAX_STEP), MAX_STEP) + # Update the per-position intercepts, kept off saturation: a logistic pinned + # at ~0 admits no drafts, so no graded drafts arrive to move it back. + new_bias = tl.minimum( + tl.maximum(bias + step_bias, -MAX_ABS_INTERCEPT), MAX_ABS_INTERCEPT + ) tl.store(slope_ptr, new_w) tl.store(intercepts_ptr + k, new_bias, mask=mask) @@ -336,6 +346,12 @@ class OnlineAcceptanceEstimator: DAMPING_OBSERVATIONS = 50.0 # Ridge on the Newton solve. L2 = 1e-3 + # Largest change of any coefficient in one refit, and coefficient bounds. + # A single confidently wrong draft otherwise moved an intercept by ~-20, + # predicting ~1e-8 acceptance for everything after it. + MAX_STEP = 1.0 + MAX_SLOPE = 4.0 + MAX_ABS_INTERCEPT = 8.0 def __init__( self, @@ -443,6 +459,9 @@ def step( NUM_SPECULATIVE_STEPS=self.num_speculative_steps, L2=self.L2, DAMPING=self.DAMPING_OBSERVATIONS, + MAX_STEP=self.MAX_STEP, + MAX_SLOPE=self.MAX_SLOPE, + MAX_ABS_INTERCEPT=self.MAX_ABS_INTERCEPT, BLOCK=triton.next_power_of_2(self.num_speculative_steps), ) self._refits += 1 diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index bc2c35cc09f5..af15b9742f14 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -34,6 +34,11 @@ from vllm.v1.worker.utils import AttentionGroup +# Consecutive steps that may schedule drafts yet verify none before one step +# verifies a draft per request anyway (see ``get_num_tokens``). +EXPLORE_AFTER_IDLE_STEPS = 32 + + def resolve_adaptive_cudagraph_mode( mode: CUDAGraphMode, *, piecewise_capture_available: bool ) -> CUDAGraphMode: @@ -451,6 +456,19 @@ def get_num_tokens( } utilities = num_tokens_to_estimated_accepted_tokens / costs draft_budget = int(np.argmax(utilities)) + # The confidences come from an estimator that learns only from admitted + # drafts, so a long run of empty budgets would leave it without the + # evidence to ever admit one again. Periodically verify one draft per + # request to keep it grounded. + idle = getattr(self, "_steps_without_drafts", 0) + if draft_budget == 0 and max_draft_budget > 0: + idle += 1 + if idle >= EXPLORE_AFTER_IDLE_STEPS: + draft_budget = min(num_reqs, max_draft_budget) + idle = 0 + else: + idle = 0 + self._steps_without_drafts = idle logger.debug( "DSpark adaptive verification selected %d/%d draft rows for %d " "requests at cost scale %.3g (selected utility %.4g, full utility " From d7fadbbbac4c917223f70e2f5fa98f111f0554e4 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 16:42:55 -0400 Subject: [PATCH 207/365] perf(scheduler): let prefill-only steps use max_num_batched_tokens max_num_scheduled_tokens bounds a step so that decode streams keep moving while long prompts prefill. When a step has no runnable decode (a single long request, or every running request still prefilling), the cap only shrinks the prefill chunk: more steps and more per-step overhead for the same work. With VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS=1, a step whose decode-aware path finds no eligible decode uses max_num_batched_tokens as its token cap and budget. Steps with runnable decode are unchanged, and the default (0) keeps current scheduling for every deployment. MiMo-V2.6-Flash TP4, batch 4096 / scheduled 2048: standalone prefill at 32K goes from 11.8K to 12.7K tok/s (+8-12% across 8K-128K). Tests: tests/v1/core/test_prefill_compute_share_scheduler.py (both settings). Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .../core/test_prefill_compute_share_scheduler.py | 14 ++++++++++++++ vllm/envs.py | 7 +++++++ vllm/v1/core/sched/scheduler.py | 15 ++++++++++++++- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_prefill_compute_share_scheduler.py b/tests/v1/core/test_prefill_compute_share_scheduler.py index aea462db33da..430957f337f9 100644 --- a/tests/v1/core/test_prefill_compute_share_scheduler.py +++ b/tests/v1/core/test_prefill_compute_share_scheduler.py @@ -1205,3 +1205,17 @@ def test_prefill_fairness_hot_switch_is_atomic(opt_model_path): unchanged = scheduler.get_prefill_fairness() assert unchanged["prefill_compute_share"] == "auto" assert unchanged["prefill_compute_half_life"] == "responsive" + + +@pytest.mark.parametrize("uncap", [False, True]) +def test_prefill_only_step_may_use_the_batched_token_budget( + opt_model_path, monkeypatch, uncap +): + """With no runnable decode, max_num_scheduled_tokens only shrinks the chunk.""" + monkeypatch.setenv("VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS", "1" if uncap else "0") + scheduler = _create_fair_scheduler(opt_model_path, max_num_batched_tokens=64) + scheduler.max_num_scheduled_tokens = 16 + (request,) = create_requests(num_requests=1, num_tokens=100) + scheduler.add_request(request) + output = scheduler.schedule() + assert output.num_scheduled_tokens[request.request_id] == (64 if uncap else 16) diff --git a/vllm/envs.py b/vllm/envs.py index 36efc990f6f7..b7779b454046 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -167,6 +167,7 @@ VLLM_RUST_FRONTEND_PATH: str | None = "auto" VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 + VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS: bool = False VLLM_MLA_DISABLE: bool = False VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" @@ -1457,6 +1458,12 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE": lambda: int( os.getenv("VLLM_V1_OUTPUT_PROC_CHUNK_SIZE", "128") ), + # Let a scheduling step with no runnable decode use the full + # max_num_batched_tokens for prefill instead of max_num_scheduled_tokens, + # which only exists to keep decode streams moving. + "VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS": lambda: ( + os.getenv("VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS", "0") == "1" + ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), # If set, vLLM will pick up the provided Flash Attention MLA diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 785212c3ae6b..b192e92b4a2f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -8,6 +8,7 @@ from dataclasses import replace from typing import Any +import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import KVEventsConfig, VllmConfig from vllm.distributed.ec_transfer.ec_connector.base import ( @@ -860,6 +861,18 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: else 0 ) has_eligible_decode = num_runnable_decodes > 0 + # max_num_scheduled_tokens keeps decode streams moving while a long + # prompt prefills. With no runnable decode it only shrinks the prefill + # chunk, so such a step may use the full batched-token budget. + step_token_cap = self.max_num_scheduled_tokens + if ( + envs.VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS + and needs_decode_count + and not has_eligible_decode + and token_budget > 0 + and step_token_cap < input_budget + ): + step_token_cap = token_budget = input_budget prefill_interleave_step = ( self.prefill_interleave_controller.begin_step( running=self.running, @@ -1892,7 +1905,7 @@ def schedule_running_requests( # Check if the scheduling constraints are satisfied. total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) - assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + assert total_num_scheduled_tokens <= step_token_cap assert token_budget >= 0 assert input_budget >= 0 assert draft_input_budget >= 0 From 6b713161c8974e975b2f800aa49cacd0c1ecbc8d Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 17:38:10 -0400 Subject: [PATCH 208/365] fix(multimodal): serve MiMo-V2 audio without torchaudio MiMo-V2 audio needs two torchaudio transforms: the media parser resamples input audio to 24 kHz (default method "torchaudio"), and MiMoVLProcessor computes a 128-bin mel spectrogram with torchaudio.transforms.MelSpectrogram. Images built on a vendor PyTorch (for example NGC 2.14.0a0) have no matching torchaudio wheel, so every audio request failed ("torchaudio is required for audio"). vllm.multimodal.audio gains plain-PyTorch ports of the default torchaudio.transforms.Resample (Hann-windowed sinc, width 6, rolloff 0.99, kernel built in float64 and stored as float32) and torchaudio.transforms.MelSpectrogram (periodic Hann window, reflect-padded STFT, HTK mel scale, no filterbank normalization). They follow torchaudio's arithmetic step for step, and are used only when torchaudio is not installed: - _get_torchaudio_resampler returns the port (cached per rate pair), so AudioResampler(method="torchaudio") works everywhere with the same output. - MiMoVLProcessor uses the ports for its mel spectrogram and its own resampler. Deployments with torchaudio installed take exactly the same code paths as before. Tests (tests/multimodal/test_audio.py, tests/models/multimodal/ test_mimo_v2_omni.py): - with torchaudio: torch.equal against torchaudio for resampling from 8-48 kHz (mono and stereo) and for MiMo's and torchaudio's default mel settings; - without torchaudio: pinned reference values, the "torchaudio" resampler path, and MiMo's audio preprocessing end to end (frames and token count). On a torchaudio-less image: 75 passed, 13 skipped (the torchaudio comparisons). Against torchaudio 2.11 / torch 2.14: 38/38 cases bitwise identical. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- tests/models/multimodal/test_mimo_v2_omni.py | 39 +++++ tests/multimodal/test_audio.py | 85 +++++++++++ vllm/multimodal/audio.py | 139 +++++++++++++++++- .../processors/mimo_v2_omni.py | 25 +++- 4 files changed, 278 insertions(+), 10 deletions(-) diff --git a/tests/models/multimodal/test_mimo_v2_omni.py b/tests/models/multimodal/test_mimo_v2_omni.py index 8e7419b99db0..58979a56830c 100644 --- a/tests/models/multimodal/test_mimo_v2_omni.py +++ b/tests/models/multimodal/test_mimo_v2_omni.py @@ -178,3 +178,42 @@ def materialize(self, source): ] expected = torch.cat([part for parts in zip(*rank_shards) for part in parts]) torch.testing.assert_close(actual, expected) + + +def _mimo_processor(): + from vllm.transformers_utils.processors.mimo_v2_omni import MiMoVLProcessor + + return MiMoVLProcessor( + tokenizer=None, + patch_size=16, + image_min_pixels=8192, + image_max_pixels=8388608, + video_min_pixels=8192, + video_max_pixels=8388608, + video_total_max_pixels=268435456, + fps=1.0, + ) + + +def test_audio_features_without_torchaudio(monkeypatch): + """Audio preprocessing works without torchaudio and matches it.""" + import vllm.transformers_utils.processors.mimo_v2_omni as processor_module + from vllm.multimodal.audio import MelSpectrogram + + wave = torch.randn(32000, generator=torch.Generator().manual_seed(7)) + audio = (wave, 16000) # resampled to 24 kHz inside the processor + + reference = None + if processor_module._HAS_TORCHAUDIO: + reference = _mimo_processor().preprocess_audio(audio) + + monkeypatch.setattr(processor_module, "_MelSpectrogram", MelSpectrogram) + monkeypatch.setattr(processor_module, "_HAS_TORCHAUDIO", False) + spec, token_len = _mimo_processor().preprocess_audio(audio) + + # 48000 samples at 24 kHz, hop 240: 201 frames of 128 log-mel bins. + assert spec.shape == (201, 128) + assert token_len == 13 + if reference is not None: + assert torch.equal(spec, reference[0]) + assert token_len == reference[1] diff --git a/tests/multimodal/test_audio.py b/tests/multimodal/test_audio.py index 21b39710d4f4..6c7000583c3f 100644 --- a/tests/multimodal/test_audio.py +++ b/tests/multimodal/test_audio.py @@ -14,11 +14,13 @@ AudioResampler, AudioSpec, ChannelReduction, + MelSpectrogram, _get_torchaudio_resampler, normalize_audio, resample_audio_pyav, resample_audio_scipy, resample_audio_torchaudio, + resample_sinc, split_audio, ) @@ -919,3 +921,86 @@ def test_split_audio_with_different_sample_rates(self): ) assert len(chunks_48k) >= 2 + + +# ============================================================ +# torchaudio-compatible ports (used when torchaudio is not installed) +# ============================================================ + + +def _chirp(num_samples: int, sample_rate: int) -> torch.Tensor: + t = torch.arange(num_samples, dtype=torch.float64) / sample_rate + wave = 0.4 * torch.sin(2 * math.pi * (200 * t + 1500 * t * t)) + wave += 0.05 * torch.sin(2 * math.pi * 7000 * t) + return wave.to(torch.float32) + + +# MiMo-V2 audio features (MiMoVLProcessor defaults). +_MIMO_MEL = dict( + sample_rate=24000, + n_fft=960, + hop_length=240, + win_length=960, + f_min=0.0, + f_max=None, + n_mels=128, + power=1.0, + center=True, +) + + +@pytest.mark.parametrize("orig_sr", [8000, 16000, 22050, 44100, 48000]) +@pytest.mark.parametrize("shape", [(24001,), (2, 16007)]) +def test_resample_sinc_matches_torchaudio(orig_sr, shape): + torchaudio = pytest.importorskip("torchaudio") + wave = torch.randn(*shape, generator=torch.Generator().manual_seed(orig_sr)) + expected = torchaudio.transforms.Resample(orig_sr, 24000)(wave) + assert torch.equal(resample_sinc(wave, orig_sr, 24000), expected) + + +@pytest.mark.parametrize("kwargs", [_MIMO_MEL, dict(sample_rate=16000, n_fft=400)]) +def test_mel_spectrogram_matches_torchaudio(kwargs): + torchaudio = pytest.importorskip("torchaudio") + wave = torch.randn(2, 24000 * 3 + 17, generator=torch.Generator().manual_seed(0)) + expected = torchaudio.transforms.MelSpectrogram(**kwargs)(wave) + assert torch.equal(MelSpectrogram(**kwargs)(wave), expected) + + +def test_resample_sinc_reference_values(): + # Reference: torchaudio 2.11 / torch 2.14 on CPU. + out = resample_sinc(_chirp(16000, 16000), 16000, 24000) + assert out.shape == (24000,) + torch.testing.assert_close( + out[[0, 1, 777, 12345, 23999]], + torch.tensor([0.0009301, 0.0465538, 0.0830263, -0.4338826, -0.2898649]), + rtol=0, + atol=2e-6, + ) + assert math.isclose(float((out.double() ** 2).sum()), 1940.9954074553, rel_tol=1e-6) + + +def test_mel_spectrogram_reference_values(): + # Reference: torchaudio 2.11 / torch 2.14 on CPU. + mel = MelSpectrogram(**_MIMO_MEL)(_chirp(24000, 24000)[None]) + assert mel.shape == (1, 128, 101) + torch.testing.assert_close( + mel[0, [0, 10, 90], [0, 5, 100]], + torch.tensor([6.934449, 0.242444, 2.970764]), + rtol=1e-5, + atol=1e-6, + ) + assert math.isclose(float(mel.double().sum()), 29227.087873592, rel_tol=1e-5) + + +def test_resample_audio_torchaudio_without_torchaudio(monkeypatch): + import vllm.multimodal.audio as audio_module + + wave = _chirp(44100, 44100).numpy() + monkeypatch.setattr(audio_module, "_HAS_TORCHAUDIO", False) + _get_torchaudio_resampler.cache_clear() + try: + out = resample_audio_torchaudio(wave, orig_sr=44100, target_sr=24000) + finally: + _get_torchaudio_resampler.cache_clear() + expected = resample_sinc(torch.from_numpy(wave), 44100, 24000).numpy() + np.testing.assert_array_equal(out, expected) diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py index 59f020230372..f2352dfe8c0b 100644 --- a/vllm/multimodal/audio.py +++ b/vllm/multimodal/audio.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from collections.abc import Callable from dataclasses import dataclass from enum import Enum -from functools import lru_cache +from functools import lru_cache, partial from typing import Literal import numpy as np @@ -30,8 +31,11 @@ try: import torchaudio + + _HAS_TORCHAUDIO = True except ImportError: torchaudio = PlaceholderModule("torchaudio") # type: ignore[assignment] + _HAS_TORCHAUDIO = False # ============================================================ @@ -281,13 +285,141 @@ def resample_audio_soxr( return soxr.resample(audio, orig_sr_int, target_sr_int) +# ============================================================ +# torchaudio-compatible transforms in plain PyTorch +# ============================================================ +# Some deployments cannot install torchaudio (for example images built on a +# vendor PyTorch with no matching wheel). These ports reproduce the default +# torchaudio.transforms.Resample and torchaudio.transforms.MelSpectrogram with +# the same arithmetic, so features are identical whether or not torchaudio is +# installed. They are used only when torchaudio is missing. + + +@lru_cache(maxsize=32) +def _sinc_resample_kernel( + orig_freq: int, + new_freq: int, + lowpass_filter_width: int = 6, + rolloff: float = 0.99, +) -> tuple[torch.Tensor, int]: + """The Hann-windowed sinc kernel of ``torchaudio.transforms.Resample``. + + Built in float64 and stored as float32, as torchaudio does when no dtype + is given. + """ + gcd = math.gcd(orig_freq, new_freq) + orig_freq //= gcd + new_freq //= gcd + base_freq = min(orig_freq, new_freq) * rolloff + width = math.ceil(lowpass_filter_width * orig_freq / base_freq) + idx = ( + torch.arange(-width, width + orig_freq, dtype=torch.float64)[None, None] + / orig_freq + ) + t = torch.arange(0, -new_freq, -1).to(torch.float32)[:, None, None] / new_freq + t = (t + idx) * base_freq + t = t.clamp_(-lowpass_filter_width, lowpass_filter_width) + window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2 + t *= math.pi + scale = base_freq / orig_freq + kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t) + kernels *= window * scale + return kernels.to(torch.float32), width + + +def resample_sinc( + waveform: torch.Tensor, orig_freq: int, new_freq: int +) -> torch.Tensor: + """``torchaudio.transforms.Resample(orig_freq, new_freq)(waveform)``. + + Bandlimited sinc interpolation over the trailing axis of a float tensor. + """ + if orig_freq == new_freq: + return waveform + kernel, width = _sinc_resample_kernel(orig_freq, new_freq) + gcd = math.gcd(orig_freq, new_freq) + orig_freq //= gcd + new_freq //= gcd + shape = waveform.size() + waveform = waveform.view(-1, shape[-1]) + num_wavs, length = waveform.shape + waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq)) + resampled = torch.nn.functional.conv1d( + waveform[:, None], kernel.to(waveform.device), stride=orig_freq + ) + resampled = resampled.transpose(1, 2).reshape(num_wavs, -1) + target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long() + resampled = resampled[..., :target_length] + return resampled.view(shape[:-1] + resampled.shape[-1:]) + + +class MelSpectrogram: + """``torchaudio.transforms.MelSpectrogram`` (HTK mel scale, no filterbank + normalization, periodic Hann window) in plain PyTorch.""" + + def __init__( + self, + *, + sample_rate: int, + n_fft: int, + win_length: int | None = None, + hop_length: int | None = None, + f_min: float = 0.0, + f_max: float | None = None, + n_mels: int = 128, + power: float = 2.0, + center: bool = True, + pad_mode: str = "reflect", + ) -> None: + self.n_fft = n_fft + self.win_length = win_length if win_length is not None else n_fft + self.hop_length = hop_length if hop_length is not None else self.win_length // 2 + self.power = power + self.center = center + self.pad_mode = pad_mode + self.window = torch.hann_window(self.win_length) + f_max = f_max if f_max is not None else float(sample_rate // 2) + # torchaudio.functional.melscale_fbanks(norm=None, mel_scale="htk") + all_freqs = torch.linspace(0, sample_rate // 2, n_fft // 2 + 1) + m_min = 2595.0 * math.log10(1.0 + (f_min / 700.0)) + m_max = 2595.0 * math.log10(1.0 + (f_max / 700.0)) + m_pts = torch.linspace(m_min, m_max, n_mels + 2) + f_pts = 700.0 * (10.0 ** (m_pts / 2595.0) - 1.0) + f_diff = f_pts[1:] - f_pts[:-1] + slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1) + down_slopes = (-1.0 * slopes[:, :-2]) / f_diff[:-1] + up_slopes = slopes[:, 2:] / f_diff[1:] + self.fb = torch.max(torch.zeros(1), torch.min(down_slopes, up_slopes)) + + def __call__(self, waveform: torch.Tensor) -> torch.Tensor: + shape = waveform.size() + spec = torch.stft( + waveform.reshape(-1, shape[-1]), + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=self.window.to(waveform.device), + center=self.center, + pad_mode=self.pad_mode, + normalized=False, + onesided=True, + return_complex=True, + ) + spec = spec.reshape(shape[:-1] + spec.shape[-2:]) + spec = spec.abs() if self.power == 1.0 else spec.abs().pow(self.power) + fb = self.fb.to(spec.device) + return torch.matmul(spec.transpose(-1, -2), fb).transpose(-1, -2) + + @lru_cache(maxsize=32) def _get_torchaudio_resampler( orig_sr: int, target_sr: int -) -> "torchaudio.transforms.Resample": +) -> Callable[[torch.Tensor], torch.Tensor]: # `torchaudio.transforms.Resample` precomputes its kernel for a fixed # (orig_sr, target_sr) pair; cache instances so repeated requests at a # common input rate skip the kernel rebuild. + if not _HAS_TORCHAUDIO: + return partial(resample_sinc, orig_freq=orig_sr, new_freq=target_sr) return torchaudio.transforms.Resample(orig_sr, target_sr) @@ -299,6 +431,9 @@ def resample_audio_torchaudio( ) -> npt.NDArray[np.floating]: """Resample audio using torchaudio's bandlimited sinc interpolation. + Without torchaudio installed, the same kernel runs in plain PyTorch + (:func:`resample_sinc`), with identical results. + Unlike the PyAV resampler, this handles any input length without padding and applies the kernel over the trailing axis, so 2D ``(channels, samples)`` input needs no per-channel loop. diff --git a/vllm/transformers_utils/processors/mimo_v2_omni.py b/vllm/transformers_utils/processors/mimo_v2_omni.py index 22349f751364..1d3fc2d09577 100644 --- a/vllm/transformers_utils/processors/mimo_v2_omni.py +++ b/vllm/transformers_utils/processors/mimo_v2_omni.py @@ -12,6 +12,7 @@ from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field +from functools import partial from typing import Any, Literal import numpy as np @@ -22,14 +23,19 @@ from transformers import BatchFeature, TensorType from transformers.processing_utils import ProcessorMixin +from vllm.multimodal.audio import MelSpectrogram as _TorchMelSpectrogram +from vllm.multimodal.audio import resample_sinc + try: import torchaudio from torchaudio.transforms import MelSpectrogram as _MelSpectrogram _HAS_TORCHAUDIO = True except ImportError: + # Same features without torchaudio: plain-PyTorch ports of its + # MelSpectrogram and Resample (vllm.multimodal.audio). torchaudio = None # type: ignore[assignment] - _MelSpectrogram = None # type: ignore[assignment,misc] + _MelSpectrogram = _TorchMelSpectrogram # type: ignore[assignment,misc] _HAS_TORCHAUDIO = False logger = logging.getLogger(__name__) @@ -365,11 +371,6 @@ def __init__( @property def mel_spectrogram(self) -> Any: if self._mel_spectrogram is None: - if _MelSpectrogram is None: - raise RuntimeError( - "torchaudio is required for audio. " - "Install with: pip install torchaudio" - ) self._mel_spectrogram = _MelSpectrogram(**self._mel_spec_kwargs) return self._mel_spectrogram @@ -427,8 +428,16 @@ def preprocess_audio(self, audio: Any) -> tuple[torch.Tensor, int]: if original_sr not in self._resamplers: if len(self._resamplers) >= self._resamplers_max: self._resamplers.popitem(last=False) - self._resamplers[original_sr] = torchaudio.transforms.Resample( - orig_freq=original_sr, new_freq=self.audio_sampling_rate + self._resamplers[original_sr] = ( + torchaudio.transforms.Resample( + orig_freq=original_sr, new_freq=self.audio_sampling_rate + ) + if _HAS_TORCHAUDIO + else partial( + resample_sinc, + orig_freq=int(original_sr), + new_freq=int(self.audio_sampling_rate), + ) ) self._resamplers.move_to_end(original_sr) waveform = self._resamplers[original_sr](waveform) From b04d9bfe2c612021c189c2fe5037545e2d81f775 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 18:59:38 -0400 Subject: [PATCH 209/365] feat(mimo): launch MiMo-V2.6-Flash with DFlash and the b12x serving flags serve-mimo26-flash.sh now starts the configuration the MiMo b12x work was qualified with: - DFlash drafting (7 tokens) from ${MODEL_PATH}/dflash, drafted across the target's TP ranks, with adaptive verification, probabilistic drafts and standard rejection sampling; the drafter keeps an FP8 KV cache (DFLASH_KV_CACHE_DTYPE). - The MiMo b12x switches, each overridable: paged-decode KV layout (VLLM_KV_CACHE_LAYOUT=BLHNC), BF16 activations for the FP4 MoE, FP32 router weights in the W4A16 top-k combine, exact dense split-K (B12X_DENSE_SPLITK_TURBO=0), BF16 GEMV plans, L2 weight prefetch, vocab-parallel drafting, uncapped prefill-only steps, and the PCIe all-reduce size thresholds. - Scheduling: 32 sequences, 4096 batched / 2048 scheduled tokens, prefill compute share 0.8, async scheduling, FULL_AND_PIECEWISE graphs up to 256. - Vision and audio through MiMoV2OmniForCausalLM. - KV cache FP8 by default: twice the BF16 capacity, and what fits the model at TP2. KV_CACHE_DTYPE=bfloat16 selects the exact BF16 cache. - MODEL_REVISION 5711b268: the first Hugging Face revision whose dflash/config.json is valid JSON (fixed in b2674c72); weights unchanged. Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- serve-mimo26-flash.sh | 52 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/serve-mimo26-flash.sh b/serve-mimo26-flash.sh index f05e127af9a7..d8d6687c8686 100755 --- a/serve-mimo26-flash.sh +++ b/serve-mimo26-flash.sh @@ -3,18 +3,22 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" PYTHON_BIN="${PYTHON_BIN:-${SCRIPT_DIR}/.venv/bin/python}" -MODEL_REVISION="${MODEL_REVISION:-3b38d063180c3e4aed9691fdc735f3d10b266ee4}" +MODEL_REVISION="${MODEL_REVISION:-5711b268169967567844e1e560e8a3966da959b1}" hf_cache="${HF_HUB_CACHE:-${HF_HOME:-${XDG_CACHE_HOME:-${HOME}/.cache}/huggingface}/hub}" MODEL_PATH="${MODEL_PATH:-${hf_cache}/models--XiaomiMiMo--MiMo-V2.6-Flash-RL/snapshots/${MODEL_REVISION}}" SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-XiaomiMiMo/MiMo-V2.6-Flash-RL}" HOST="${HOST:-0.0.0.0}" PORT="${PORT:-8000}" DEVICE_IDS="${DEVICE_IDS:-8,9,10,11}" -GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.98}" -MAX_MODEL_LEN="${MAX_MODEL_LEN:-auto}" -MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}" +GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.96}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-1048576}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}" MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-4096}" -NUM_SPECULATIVE_TOKENS="${NUM_SPECULATIVE_TOKENS:-0}" +MAX_NUM_SCHEDULED_TOKENS="${MAX_NUM_SCHEDULED_TOKENS:-2048}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +NUM_SPECULATIVE_TOKENS="${NUM_SPECULATIVE_TOKENS:-7}" +DFLASH_MODEL_PATH="${DFLASH_MODEL_PATH:-${MODEL_PATH}/dflash}" +DFLASH_KV_CACHE_DTYPE="${DFLASH_KV_CACHE_DTYPE:-fp8}" LOAD_FORMAT="${LOAD_FORMAT:-b12x}" if [[ ! -x "${PYTHON_BIN}" ]]; then @@ -53,6 +57,10 @@ if [[ "${CUDA_DEVICE_ORDER:-PCI_BUS_ID}" != PCI_BUS_ID ]]; then echo "CUDA_DEVICE_ORDER must be PCI_BUS_ID when using physical --device-ids" >&2 exit 2 fi +if [[ ! "${MAX_MODEL_LEN}" =~ ^[1-9][0-9]*$ ]]; then + echo "MAX_MODEL_LEN must be a positive integer; got '${MAX_MODEL_LEN}'" >&2 + exit 2 +fi if [[ ! "${NUM_SPECULATIVE_TOKENS}" =~ ^(0|[1-9][0-9]*)$ ]]; then echo "NUM_SPECULATIVE_TOKENS must be a non-negative integer; got '${NUM_SPECULATIVE_TOKENS}'" >&2 exit 2 @@ -73,15 +81,32 @@ export NCCL_P2P_LEVEL="${NCCL_P2P_LEVEL:-SYS}" export NCCL_PROTO="${NCCL_PROTO:-LL,LL128,Simple}" export VLLM_ENABLE_PCIE_ALLREDUCE="${VLLM_ENABLE_PCIE_ALLREDUCE:-1}" export VLLM_PCIE_ALLREDUCE_BACKEND="${VLLM_PCIE_ALLREDUCE_BACKEND:-b12x}" +export VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE="${VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE:-64KB}" +export VLLM_PCIE_TWOSHOT_ALLREDUCE_MAX_SIZE="${VLLM_PCIE_TWOSHOT_ALLREDUCE_MAX_SIZE:-2MB}" +export VLLM_PCIE_DMA_MIN_BYTES="${VLLM_PCIE_DMA_MIN_BYTES:-off}" +export VLLM_USE_V2_MODEL_RUNNER="${VLLM_USE_V2_MODEL_RUNNER:-1}" +export VLLM_KV_CACHE_LAYOUT="${VLLM_KV_CACHE_LAYOUT:-BLHNC}" +export VLLM_B12X_MOE_FP4_FORCE_A16="${VLLM_B12X_MOE_FP4_FORCE_A16:-1}" +export B12X_W4A16_FP32_TOPK_WEIGHTS="${B12X_W4A16_FP32_TOPK_WEIGHTS:-1}" +export B12X_DENSE_SPLITK_TURBO="${B12X_DENSE_SPLITK_TURBO:-0}" +export VLLM_B12X_BF16_GEMV="${VLLM_B12X_BF16_GEMV:-1}" +export VLLM_MIMO_L2_PREFETCH="${VLLM_MIMO_L2_PREFETCH:-1}" +export VLLM_DFLASH_VOCAB_PARALLEL_DRAFT="${VLLM_DFLASH_VOCAB_PARALLEL_DRAFT:-1}" +export VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS="${VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS:-1}" +export VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN="${VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN:-32768}" export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}" export VLLM_WORKER_MULTIPROC_METHOD="${VLLM_WORKER_MULTIPROC_METHOD:-spawn}" export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" speculative_args=() if ((NUM_SPECULATIVE_TOKENS > 0)); then - printf -v speculative_config \ - '{"method":"mtp","num_speculative_tokens":%s}' \ - "${NUM_SPECULATIVE_TOKENS}" + if [[ ! -f "${DFLASH_MODEL_PATH}/config.json" ]]; then + echo "DFlash config not found: ${DFLASH_MODEL_PATH}/config.json" >&2 + exit 1 + fi + speculative_config="$("${PYTHON_BIN}" -c \ + 'import json, sys; print(json.dumps({"method": "dflash", "model": sys.argv[1], "num_speculative_tokens": int(sys.argv[2]), "draft_tensor_parallel_size": int(sys.argv[3]), "max_model_len": int(sys.argv[4]), "enable_adaptive_verification": True, "draft_sample_method": "probabilistic", "rejection_sample_method": "standard", "attention_backend": "B12X", "kv_cache_dtype": sys.argv[5]}))' \ + "${DFLASH_MODEL_PATH}" "${NUM_SPECULATIVE_TOKENS}" "${TP_SIZE}" "${MAX_MODEL_LEN}" "${DFLASH_KV_CACHE_DTYPE}")" speculative_args=(--speculative-config "${speculative_config}") fi @@ -98,8 +123,8 @@ command=( --enable-prefix-caching --enable-chunked-prefill --dtype bfloat16 - --kv-cache-dtype bfloat16 - --block-size 128 + --kv-cache-dtype "${KV_CACHE_DTYPE}" + --block-size 64 --attention-backend B12X --linear-backend b12x --moe-backend b12x @@ -109,6 +134,13 @@ command=( --max-model-len "${MAX_MODEL_LEN}" --max-num-seqs "${MAX_NUM_SEQS}" --max-num-batched-tokens "${MAX_NUM_BATCHED_TOKENS}" + --max-num-scheduled-tokens "${MAX_NUM_SCHEDULED_TOKENS}" + --prefill-compute-share 0.8 + --async-scheduling + --no-scheduler-reserve-full-isl + --max-cudagraph-capture-size 256 + --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' + --hf-overrides '{"architectures":["MiMoV2OmniForCausalLM"]}' --reasoning-parser mimo --tool-call-parser mimo --enable-auto-tool-choice From 45f0e24cac6d6f26f96253e61ca9c9ad597a008d Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:27 -0400 Subject: [PATCH 210/365] docs(release): add vllm-871 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-871.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-871.json diff --git a/.lil/changes/vllm-871.json b/.lil/changes/vllm-871.json new file mode 100644 index 000000000000..fd846782e5d1 --- /dev/null +++ b/.lil/changes/vllm-871.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-871", + "category": "feature", + "summary": "Serve MiMo V2 decode and speculative verify through b12x paged_decode", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "The B12X backend serves DiffKV layers and non-causal FP8-KV drafter layers through b12x paged_decode, including ragged adaptive-verification batches, and appends their KV through the same plan.", + "Other layer geometries keep their existing B12X paths (VLLM_B12X_PAGED_DECODE=auto)." + ], + "pull_requests": [ + 871 + ], + "requires": [ + "b12x-420" + ] +} From b7998ab868497624b87b0d98cf07ea7e72236547 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:28 -0400 Subject: [PATCH 211/365] docs(release): add vllm-872 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-872.json | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 .lil/changes/vllm-872.json diff --git a/.lil/changes/vllm-872.json b/.lil/changes/vllm-872.json new file mode 100644 index 000000000000..5a9f8d69146d --- /dev/null +++ b/.lil/changes/vllm-872.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-872", + "category": "performance", + "summary": "Fuse the b12x PCIe all-reduce with the residual RMSNorm for MiMo V2 and DFlash", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "Fused all-reduce + RMSNorm plans are keyed by the norm weight's storage, so the compiled graph's parameter aliases find them.", + "MiMo V2 and the DFlash drafter declare their fused sites; other models are unchanged." + ], + "pull_requests": [ + 872 + ] +} From 3024840a48f54eaa01010a8f6c99fe3979ce5c76 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:30 -0400 Subject: [PATCH 212/365] docs(release): add vllm-873 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-873.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-873.json diff --git a/.lil/changes/vllm-873.json b/.lil/changes/vllm-873.json new file mode 100644 index 000000000000..fbaa52af1acf --- /dev/null +++ b/.lil/changes/vllm-873.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-873", + "category": "performance", + "summary": "Optionally serve decode-sized BF16 linears through b12x bf16_gemv", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "VLLM_B12X_BF16_GEMV=1 routes BF16 linears with up to eight rows to autotuned b12x GEMV plans under --linear-backend b12x. Default off." + ], + "pull_requests": [ + 873 + ] +} From e8d4aa5eb588c147db0ec334df9490266bab13a7 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:31 -0400 Subject: [PATCH 213/365] docs(release): add vllm-874 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-874.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-874.json diff --git a/.lil/changes/vllm-874.json b/.lil/changes/vllm-874.json new file mode 100644 index 000000000000..0a66f0ec2f0f --- /dev/null +++ b/.lil/changes/vllm-874.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-874", + "category": "fix", + "summary": "Load MiMo V2 fused FP8 QKV exactly below TP4", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "At TP1/TP2 the fused QKV kept the checkpoint's FP8 block scales instead of being re-quantized (which changed 40% of the global layers' K and 65% of their V weights at TP2). TP4 is unchanged." + ], + "pull_requests": [ + 874 + ] +} From 29e1cf451992d35cdee3b6f581d862d9e97fbafe Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:33 -0400 Subject: [PATCH 214/365] docs(release): add vllm-875 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-875.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-875.json diff --git a/.lil/changes/vllm-875.json b/.lil/changes/vllm-875.json new file mode 100644 index 000000000000..8eb577420050 --- /dev/null +++ b/.lil/changes/vllm-875.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-875", + "category": "performance", + "summary": "Optional L2 weight prefetch for MiMo V2 decode", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "VLLM_MIMO_L2_PREFETCH=1 prefetches upcoming decode weights into L2 on a side stream inside FULL CUDA-graph decode. Cache hints only; numerics unchanged. Default off." + ], + "pull_requests": [ + 875 + ] +} From d7b38cad4051b617c21407b485eb9e5e007d4f57 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:34 -0400 Subject: [PATCH 215/365] docs(release): add vllm-876 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-876.json | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 .lil/changes/vllm-876.json diff --git a/.lil/changes/vllm-876.json b/.lil/changes/vllm-876.json new file mode 100644 index 000000000000..e46cc5490616 --- /dev/null +++ b/.lil/changes/vllm-876.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-876", + "category": "fix", + "summary": "MiMo-V2 audio input works without torchaudio, with identical features", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required. Deployments with torchaudio installed are unchanged.", + "details": [ + "Images without a matching torchaudio wheel rejected every MiMo audio request: the media parser's default resampler and the MiMo processor's mel spectrogram both required torchaudio.", + "Plain-PyTorch ports of torchaudio's default Resample and MelSpectrogram now serve when torchaudio is missing; they are bitwise identical to torchaudio (tests compare with torch.equal when torchaudio is present and pin reference values when it is not)." + ], + "pull_requests": [ + 876 + ] +} From 0973d8b2d1bd9a4fc5d9729971f8c721a8d34d24 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:36 -0400 Subject: [PATCH 216/365] docs(release): add vllm-877 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-877.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-877.json diff --git a/.lil/changes/vllm-877.json b/.lil/changes/vllm-877.json new file mode 100644 index 000000000000..fbce227468e0 --- /dev/null +++ b/.lil/changes/vllm-877.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-877", + "category": "performance", + "summary": "Shard the DFlash context projection and optionally draft per vocab shard at TP > 1", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "The DFlash context projection is column-parallel at TP > 1 (bit-identical output, 1/tp of the weight per rank).", + "VLLM_DFLASH_VOCAB_PARALLEL_DRAFT=1 samples probabilistic drafts per vocab shard with identical draws and verification, without the draft-logits all-gather. Default off." + ], + "pull_requests": [ + 877 + ] +} From 2fdc99d83ffe7ad2c0ee5b85f9456c057e4a684a Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:37 -0400 Subject: [PATCH 217/365] docs(release): add vllm-878 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-878.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-878.json diff --git a/.lil/changes/vllm-878.json b/.lil/changes/vllm-878.json new file mode 100644 index 000000000000..252ac0e34e0f --- /dev/null +++ b/.lil/changes/vllm-878.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-878", + "category": "fix", + "summary": "Support block-FP8 DFlash drafters in the fused context projection", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "Online fp8_per_block drafters failed in the fused context K/V projection (BF16 x FP8 matmul); their K/V rows are now dequantized with their block scales for that projection." + ], + "pull_requests": [ + 878 + ] +} From 18992b1c5fef62e6b4a523f2840289cf7d80db9f Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:38 -0400 Subject: [PATCH 218/365] docs(release): add vllm-879 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-879.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-879.json diff --git a/.lil/changes/vllm-879.json b/.lil/changes/vllm-879.json new file mode 100644 index 000000000000..7a87bac126cb --- /dev/null +++ b/.lil/changes/vllm-879.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-879", + "category": "fix", + "summary": "Keep adaptive verification's acceptance estimator from collapsing", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "One confidently wrong draft could drive the estimator's predicted acceptance to ~1e-8, after which adaptive verification admitted no drafts for the rest of the run (DFlash or DSpark drafting stopping mid-generation). Refits are now bounded and an idle-step probe keeps evidence flowing." + ], + "pull_requests": [ + 879 + ] +} From bd4cbd7bce39bf54495ba1cf210c94950af9e083 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:40 -0400 Subject: [PATCH 219/365] docs(release): add vllm-880 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-880.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-880.json diff --git a/.lil/changes/vllm-880.json b/.lil/changes/vllm-880.json new file mode 100644 index 000000000000..bc73e6ebe929 --- /dev/null +++ b/.lil/changes/vllm-880.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-880", + "category": "performance", + "summary": "Optionally let prefill-only scheduling steps use the full batched-token budget", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS=1 lets steps with no runnable decode prefill up to max_num_batched_tokens. Default off." + ], + "pull_requests": [ + 880 + ] +} From 64c563c293900989b5849f04fe194eabee7084b2 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 23 Sep 2026 19:26:41 -0400 Subject: [PATCH 220/365] docs(release): add vllm-881 release fragment Signed-off-by: MadeBy561 Co-authored-by: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-881.json | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 .lil/changes/vllm-881.json diff --git a/.lil/changes/vllm-881.json b/.lil/changes/vllm-881.json new file mode 100644 index 000000000000..cc8bbc45754e --- /dev/null +++ b/.lil/changes/vllm-881.json @@ -0,0 +1,28 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-881", + "category": "feature", + "summary": "serve-mimo26-flash.sh starts MiMo-V2.6-Flash with DFlash and the b12x serving flags", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "serve-mimo26-flash.sh now defaults to DFlash and an FP8 KV cache and pins revision 5711b268; KV_CACHE_DTYPE=bfloat16 and NUM_SPECULATIVE_TOKENS=0 restore the previous cache and no speculation.", + "details": [ + "DFlash drafting (7 tokens) with TP drafting and adaptive verification, FP8 KV cache by default, vision and audio via MiMoV2OmniForCausalLM, and the MiMo b12x switches.", + "MODEL_REVISION moves to 5711b268, the first revision with a valid dflash/config.json." + ], + "pull_requests": [ + 881 + ], + "requires": [ + "b12x-420", + "b12x-423", + "vllm-871", + "vllm-872", + "vllm-873", + "vllm-875", + "vllm-877", + "vllm-880" + ] +} From 4a3a8eceeb7b7eaac27d7c3717897f74b9539a60 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 23:27:13 +0000 Subject: [PATCH 221/365] chore(release): describe #869 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-869.json | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 .lil/changes/vllm-869.json diff --git a/.lil/changes/vllm-869.json b/.lil/changes/vllm-869.json new file mode 100644 index 000000000000..f74ebc4ff1ca --- /dev/null +++ b/.lil/changes/vllm-869.json @@ -0,0 +1,19 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-869", + "category": "fix", + "summary": "Track allocated KV blocks and keep CoW group store ranges aligned in the offloading connector", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "Two OffloadingConnector fixes from #867: newly allocated blocks are tracked before an early return (port of vllm-project/vllm#46960), and skipped align-mode CoW groups keep their store-range slots, which stops an EngineCore assert on the first store of a long prompt." + ], + "pull_requests": [ + 869 + ], + "authors": [ + "@llitz" + ] +} From f278b4a280f68037300733d796d7fef8bc2dfeed Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 23 Sep 2026 23:27:16 +0000 Subject: [PATCH 222/365] chore(release): describe #870 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-870.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-870.json diff --git a/.lil/changes/vllm-870.json b/.lil/changes/vllm-870.json new file mode 100644 index 000000000000..e4650c1855e8 --- /dev/null +++ b/.lil/changes/vllm-870.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-870", + "category": "fix", + "summary": "Serve external prefix hits for align-mode Mamba hybrids through the offloading connector", + "models": [ + "Qwen3.8-Flash-Next" + ], + "compatibility": "No user action required.", + "details": [ + "Boundary states are captured in dedicated blocks and QSA pages transfer their padded selector tail, so OffloadingConnector prefix hits on Qwen3.8 are correct.", + "A restored boundary checkpoint keeps its QSA selector pools instead of treating the slot as fresh; GPU prefix-hit and LMCache restore answers were unchanged in our Qwen TP2 MTP3 checks.", + "Qwen still rejects OffloadingConnector until a connector declares atomic support." + ], + "pull_requests": [ + 870 + ], + "authors": [ + "@llitz" + ] +} From 212f46bfa8261c088826d77a09490fb7ce300f76 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 00:33:17 +0000 Subject: [PATCH 223/365] fix(config): let aligned Qwen offload through the atomic-checkpoint gate #864 requires an atomic request-boundary checkpoint connector for every Qwen KV transfer. That also rejected CACHE_MODE=native: vLLM maps native offload to SimpleCPUOffloadConnector with --recurrent-checkpoint-policy aligned, so every beta since #864 refused to start Qwen with the native CPU cache. Aligned retention does not use request-boundary checkpoints. Accept it when the connector declares supports_aligned_hybrid_transfer (SimpleCPU, and OffloadingConnector with the #869/#870 fixes) and DCP is 1. Other connectors, DCP and the auto/request_boundaries policies keep the atomic requirement. Based on commit 4 of #867, narrowed to connectors that move aligned hybrid state. Co-authored-by: Carlos Augusto Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- tests/v1/kv_connector/unit/test_config.py | 39 +++++++++++++++++++ vllm/config/vllm.py | 22 ++++++++--- .../kv_transfer/kv_connector/v1/base.py | 10 +++++ .../kv_connector/v1/offloading_connector.py | 5 +++ .../v1/simple_cpu_offload_connector.py | 5 +++ 5 files changed, 76 insertions(+), 5 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index b00cca641329..8858ed743a6f 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -105,6 +105,8 @@ def _build_config( enable_sleep_mode: bool = False, enable_cumem_allocator: bool = False, model_type: str | None = None, + recurrent_checkpoint_policy: str | None = None, + decode_context_parallel_size: int = 1, ) -> VllmConfig: """Build a model-free config for KV-transfer compatibility checks. @@ -139,6 +141,13 @@ def _build_config( enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode), hf_text_config=SimpleNamespace(model_type=model_type), ) + if recurrent_checkpoint_policy is not None: + cfg.cache_config = SimpleNamespace( + recurrent_checkpoint_policy=recurrent_checkpoint_policy + ) + cfg.parallel_config = SimpleNamespace( + decode_context_parallel_size=decode_context_parallel_size + ) cfg._verify_kv_transfer_compat() return cfg @@ -166,6 +175,36 @@ def supports_request_boundary_checkpoints(cls, config): _build_config(kv_connector="stub", model_type="qwen3_8_flash_next_text") +@pytest.mark.parametrize( + ("kv_connector", "policy", "dcp", "allowed"), + [ + ("SimpleCPUOffloadConnector", "aligned", 1, True), + ("OffloadingConnector", "aligned", 1, True), + ("SimpleCPUOffloadConnector", "auto", 1, False), + ("SimpleCPUOffloadConnector", "request_boundaries", 1, False), + ("SimpleCPUOffloadConnector", "aligned", 2, False), + ("LMCacheMPConnector", "aligned", 1, False), + ("NixlConnector", "aligned", 1, False), + ], +) +def test_qwen_aligned_retention_needs_an_aligned_hybrid_connector( + kv_connector, policy, dcp, allowed +): + """CACHE_MODE=native (SimpleCPU, aligned) must boot Qwen; other + connectors and DCP keep the atomic checkpoint requirement.""" + build = lambda: _build_config( # noqa: E731 + kv_connector=kv_connector, + model_type="qwen3_8_flash_next_text", + recurrent_checkpoint_policy=policy, + decode_context_parallel_size=dcp, + ) + if allowed: + build() + else: + with pytest.raises(ValueError, match="atomic request-boundary"): + build() + + @pytest.mark.parametrize( "kv_connector", ["NixlConnector", "MooncakeConnectorV1", "SomeOOTConnector"] ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8c61475b917d..af4d8c0c017a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1254,13 +1254,25 @@ def _verify_kv_transfer_compat(self) -> None: KVConnectorFactory, ) - connector = KVConnectorFactory.get_connector_class( - self.kv_transfer_config - ) - if not connector.supports_request_boundary_checkpoints(self): + connector = KVConnectorFactory.get_connector_class(self.kv_transfer_config) + # Explicit aligned retention bypasses request-boundary checkpoints; + # it is correct only with connectors that move hybrid state at + # aligned boundaries, and only validated without DCP. + cache_config = getattr(self, "cache_config", None) + parallel_config = getattr(self, "parallel_config", None) + aligned = ( + getattr(cache_config, "recurrent_checkpoint_policy", None) == "aligned" + and getattr(parallel_config, "decode_context_parallel_size", None) == 1 + and connector.supports_aligned_hybrid_transfer(self) + ) + if not aligned and not connector.supports_request_boundary_checkpoints( + self + ): raise ValueError( "Qwen QSA KV transfer requires an atomic request-boundary " - "checkpoint connector" + "checkpoint connector, or --recurrent-checkpoint-policy " + "aligned without DCP and a connector that transfers aligned " + "hybrid state (SimpleCPUOffloadConnector, OffloadingConnector)" ) # PyTorch's expandable_segments allocator uses CUDA VMM, which can diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index b5666641a8b0..a8a941b56dff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -492,6 +492,16 @@ def supports_request_boundary_checkpoints(cls, config: "VllmConfig") -> bool: """Whether configuration selects a complete external checkpoint adapter.""" return False + @classmethod + def supports_aligned_hybrid_transfer(cls, config: "VllmConfig") -> bool: + """Whether aligned retention transfers hybrid recurrent state correctly. + + Qwen QSA models may use such a connector with + ``--recurrent-checkpoint-policy aligned`` instead of an atomic + request-boundary checkpoint adapter. + """ + return False + def bind_boundary_checkpoint_cache(self, manager: "KVCacheManager") -> None: """Bind the scheduler allocator for atomic external checkpoint imports. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 2eb0f56118b6..b148c5ba7226 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -50,6 +50,11 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): + @classmethod + def supports_aligned_hybrid_transfer(cls, config: "VllmConfig") -> bool: + # Boundary captures and padded QSA page refs serve align-mode hits. + return True + @cached_property def _bounding_group_ids(self) -> tuple[int, ...]: """Prefix-cacheable groups this connector does not offload. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 64bf8146c9e5..f4f7906b9523 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -55,6 +55,11 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): """CPU KV cache offloading with custom kernel transfers and BlockPool LRU.""" + @classmethod + def supports_aligned_hybrid_transfer(cls, config: "VllmConfig") -> bool: + # Aligned Qwen QSA restores are validated with this connector. + return True + @property def requires_kv_delivery(self) -> bool: # Runs as kv_both, but is a best-effort cache: a dropped save is just a From 3c2dd482c496f74e769f8493c5a6a5e84280b0f8 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 00:34:38 +0000 Subject: [PATCH 224/365] chore(release): describe #883 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-883.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-883.json diff --git a/.lil/changes/vllm-883.json b/.lil/changes/vllm-883.json new file mode 100644 index 000000000000..172de113e734 --- /dev/null +++ b/.lil/changes/vllm-883.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-883", + "category": "fix", + "summary": "Start Qwen with the native CPU prefix cache again", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Since #864, Qwen with CACHE_MODE=native stopped at startup with \"Qwen QSA KV transfer requires an atomic request-boundary checkpoint connector\".", + "Aligned retention is accepted again with connectors that transfer aligned hybrid state (SimpleCPU, OffloadingConnector) and DCP 1; other connectors keep the atomic checkpoint requirement." + ], + "pull_requests": [883], + "authors": ["@llitz"] +} From 6cf4fc4ebd4dfc538bcd5acfdbaf33fe5509c8de Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 09:16:35 +0000 Subject: [PATCH 225/365] feat(serve): time requests from HTTP receipt, not only from rendering The request timeline started at the renderer's arrival stamp, so a request held after the server received it but before rendering, for example while its body was still arriving, looked fast. An outermost ASGI middleware now stamps the HTTP receipt and the end of the request body, the timeline reports both spans, and a request that ends without any response after the stall threshold is logged even though it never reached the engine. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-diag-receipt.json | 13 +++ .../serve/middleware/test_request_receipt.py | 99 +++++++++++++++++++ vllm/entrypoints/serve/middleware/register.py | 5 + .../serve/middleware/request_receipt.py | 67 +++++++++++++ vllm/v1/engine/stall_diagnostics.py | 49 +++++++-- 5 files changed, 227 insertions(+), 6 deletions(-) create mode 100644 .lil/changes/vllm-diag-receipt.json create mode 100644 tests/entrypoints/serve/middleware/test_request_receipt.py create mode 100644 vllm/entrypoints/serve/middleware/request_receipt.py diff --git a/.lil/changes/vllm-diag-receipt.json b/.lil/changes/vllm-diag-receipt.json new file mode 100644 index 000000000000..93600b93c689 --- /dev/null +++ b/.lil/changes/vllm-diag-receipt.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-diag-receipt", + "category": "feature", + "summary": "Time requests from HTTP receipt, not only from rendering", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "The request timeline adds HTTP receipt to body read and body read to rendering, so a request held before rendering, for example by a slowly delivered body, shows its wait.", + "A request that ends without any response after VLLM_REQUEST_STALL_WARNING_S logs a warning with the method, path and when its body was read." + ], + "authors": ["@voipmonitor"] +} diff --git a/tests/entrypoints/serve/middleware/test_request_receipt.py b/tests/entrypoints/serve/middleware/test_request_receipt.py new file mode 100644 index 000000000000..0e0725833c46 --- /dev/null +++ b/tests/entrypoints/serve/middleware/test_request_receipt.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The receipt middleware times requests held before rendering.""" + +import asyncio +import time + +import pytest + +from vllm.entrypoints.serve.middleware.request_receipt import RequestReceiptMiddleware +from vllm.v1.engine.stall_diagnostics import REQUEST_RECEIPT, RequestTimeline + +SCOPE = {"type": "http", "method": "POST", "path": "/v1/chat/completions"} + + +def _slow_body(delay: float): + parts = [ + {"type": "http.request", "body": b"{", "more_body": True}, + {"type": "http.request", "body": b"}", "more_body": False}, + ] + + async def receive(): + message = parts.pop(0) + if not message["more_body"]: + await asyncio.sleep(delay) + return message + + return receive + + +def test_timeline_includes_http_receipt_and_body_read(caplog_vllm): + seen = {} + + async def app(scope, receive, send): + while (await receive()).get("more_body"): + pass + seen["timeline"] = RequestTimeline("req-1", time.time(), time.time()) + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + async def send(message): + return None + + asyncio.run(RequestReceiptMiddleware(app)(SCOPE, _slow_body(0.3), send)) + timeline = seen["timeline"] + assert timeline.received is not None and timeline.body_read is not None + assert timeline.body_read - timeline.received >= 0.3 + timeline.close("finished", threshold=0.0, always=True) + assert "HTTP receipt to body read 0.3" in caplog_vllm.text + assert REQUEST_RECEIPT.get() is None + + +def test_request_ending_without_response_is_reported(caplog_vllm): + async def app(scope, receive, send): + while (await receive()).get("more_body"): + pass + await asyncio.sleep(0.3) + raise RuntimeError("client disconnected") + + async def send(message): + return None + + middleware = RequestReceiptMiddleware(app) + middleware.threshold = 0.2 + with pytest.raises(RuntimeError, match="client disconnected"): + asyncio.run(middleware(SCOPE, _slow_body(0.0), send)) + assert "HTTP POST /v1/chat/completions ended after 0.3" in caplog_vllm.text + assert "without a response" in caplog_vllm.text + + +def test_api_server_registers_the_receipt_middleware(): + from argparse import Namespace + + from fastapi import FastAPI, Request + from fastapi.testclient import TestClient + + from vllm.entrypoints.serve.middleware.register import ( + init_entrypoints_middleware, + ) + + app = FastAPI() + + @app.post("/probe") + async def probe(request: Request): + await request.body() + receipt = REQUEST_RECEIPT.get() + return {"received": receipt is not None and receipt.body_read is not None} + + args = Namespace( + allowed_origins=["*"], + allow_credentials=False, + allowed_methods=["*"], + allowed_headers=["*"], + api_key=None, + enable_request_id_headers=False, + middleware=[], + ) + init_entrypoints_middleware(args, app, ()) + assert TestClient(app).post("/probe", json={}).json() == {"received": True} diff --git a/vllm/entrypoints/serve/middleware/register.py b/vllm/entrypoints/serve/middleware/register.py index 3b132ce0886d..f72e5d396c5d 100644 --- a/vllm/entrypoints/serve/middleware/register.py +++ b/vllm/entrypoints/serve/middleware/register.py @@ -73,3 +73,8 @@ def init_entrypoints_middleware( raise ValueError( f"Invalid middleware {middleware}. Must be a function or a class." ) + + # Outermost, so the receipt stamp precedes every other middleware. + from .request_receipt import RequestReceiptMiddleware + + app.add_middleware(RequestReceiptMiddleware) diff --git a/vllm/entrypoints/serve/middleware/request_receipt.py b/vllm/entrypoints/serve/middleware/request_receipt.py new file mode 100644 index 000000000000..bb80f08c0ab0 --- /dev/null +++ b/vllm/entrypoints/serve/middleware/request_receipt.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stamp when each HTTP request arrives and when its body is fully read. + +The renderer stamps a request's arrival only when rendering starts, after the +body has been read and parsed. These stamps let the request timeline show a +wait before that point, such as a slowly delivered request body. +""" + +import time + +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from vllm import envs +from vllm.logger import init_logger +from vllm.v1.engine.stall_diagnostics import REQUEST_RECEIPT, RequestReceipt + +logger = init_logger(__name__) + + +class RequestReceiptMiddleware: + def __init__(self, app: ASGIApp) -> None: + self.app = app + self.threshold = envs.VLLM_REQUEST_STALL_WARNING_S + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + receipt = RequestReceipt(time.time()) + + async def timed_receive() -> Message: + message = await receive() + if ( + receipt.body_read is None + and message["type"] == "http.request" + and not message.get("more_body", False) + ): + receipt.body_read = time.time() + return message + + responded = False + + async def watched_send(message: Message) -> None: + nonlocal responded + if message["type"] == "http.response.start": + responded = True + await send(message) + + token = REQUEST_RECEIPT.set(receipt) + try: + await self.app(scope, timed_receive, watched_send) + finally: + REQUEST_RECEIPT.reset(token) + # A request held before rendering never reaches the engine's + # timeline; report it here if it ended without any response. + elapsed = time.time() - receipt.received + if not responded and 0 < self.threshold < elapsed: + logger.warning( + "HTTP %s %s ended after %.2f s without a response; body read %s", + scope.get("method", ""), + scope.get("path", ""), + elapsed, + "never" + if receipt.body_read is None + else f"after {receipt.body_read - receipt.received:.2f} s", + ) diff --git a/vllm/v1/engine/stall_diagnostics.py b/vllm/v1/engine/stall_diagnostics.py index e50ede7b9e07..e5b65f763453 100644 --- a/vllm/v1/engine/stall_diagnostics.py +++ b/vllm/v1/engine/stall_diagnostics.py @@ -12,6 +12,7 @@ import threading import time import traceback +from contextvars import ContextVar from typing import TYPE_CHECKING, Any from vllm.logger import init_logger @@ -25,6 +26,22 @@ _STACK_FRAMES = 40 +class RequestReceipt: + """When the API server received an HTTP request and finished reading it.""" + + __slots__ = ("received", "body_read") + + def __init__(self, received: float): + self.received = received + self.body_read: float | None = None + + +# Set by the HTTP receipt middleware for the task that serves the request. +REQUEST_RECEIPT: ContextVar[RequestReceipt | None] = ContextVar( + "vllm_request_receipt", default=None +) + + def _clock(timestamp: float) -> str: """Format a wall-clock time like the log line prefix, with milliseconds.""" return time.strftime("%H:%M:%S", time.localtime(timestamp)) + ( @@ -53,6 +70,8 @@ class RequestTimeline: __slots__ = ( "request_id", + "received", + "body_read", "arrival", "started", "processed", @@ -65,6 +84,10 @@ class RequestTimeline: def __init__(self, request_id: str, arrival: float | None, started: float): self.request_id = request_id + receipt = REQUEST_RECEIPT.get() + # HTTP receipt and complete request body, when served over HTTP. + self.received = receipt.received if receipt is not None else None + self.body_read = receipt.body_read if receipt is not None else None self.arrival = arrival self.started = started # Input processing finished and the output collector was created. @@ -90,25 +113,39 @@ def output(self, output: "RequestOutput") -> None: def outside_engine(self) -> float: """Longest stretch spent outside engine work, in seconds. - That is rendering until the generator started, input processing plus - submission, or the handler taking the final output to the client. + That is HTTP receipt (or rendering) until the generator started, input + processing plus submission, or the handler taking the final output to + the client. """ waits = [0.0] - if self.arrival is not None: - waits.append(self.started - self.arrival) + first = self._origin() + waits.append(self.started - first) if self.submitted is not None: waits.append(self.submitted - self.started) if self.final_output is not None and self.closed is not None: waits.append(self.closed - self.final_output) return max(waits) + def _origin(self) -> float: + for stamp in (self.received, self.arrival): + if stamp is not None: + return stamp + return self.started + def describe(self, outcome: str) -> str: """Return a one-line account of every stage the request reached.""" assert self.closed is not None - origin = self.arrival if self.arrival is not None else self.started + origin = self._origin() parts = [ f"Request {self.request_id} {outcome} after " f"{self.closed - origin:.2f} s (arrived {_clock(origin)})", + ] + if self.received is not None: + parts += [ + f"HTTP receipt to body read {_span(self.received, self.body_read)}", + f"body read to rendering {_span(self.body_read, self.arrival)}", + ] + parts += [ f"arrival to generator start {_span(self.arrival, self.started)}", f"input processing {_span(self.started, self.processed)}", f"engine submission {_span(self.processed, self.submitted)}", @@ -153,7 +190,7 @@ def close(self, outcome: str, threshold: float, always: bool) -> None: return if threshold <= 0: return - origin = self.arrival if self.arrival is not None else self.started + origin = self._origin() if self.outside_engine() > threshold or ( self.first_output is None and self.closed - origin > threshold ): From 08615a68f77d66bc7e5ef56c30dd1dea370b5368 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 09:16:51 +0000 Subject: [PATCH 226/365] chore(release): identify the receipt timeline change as #884 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/{vllm-diag-receipt.json => vllm-884.json} | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) rename .lil/changes/{vllm-diag-receipt.json => vllm-884.json} (86%) diff --git a/.lil/changes/vllm-diag-receipt.json b/.lil/changes/vllm-884.json similarity index 86% rename from .lil/changes/vllm-diag-receipt.json rename to .lil/changes/vllm-884.json index 93600b93c689..c2e3d59c41f5 100644 --- a/.lil/changes/vllm-diag-receipt.json +++ b/.lil/changes/vllm-884.json @@ -1,13 +1,17 @@ { "schema": "local-inference-release-change/v1", - "id": "vllm-diag-receipt", + "id": "vllm-884", "category": "feature", "summary": "Time requests from HTTP receipt, not only from rendering", - "models": ["all"], + "models": [ + "all" + ], "compatibility": "No user action required.", "details": [ "The request timeline adds HTTP receipt to body read and body read to rendering, so a request held before rendering, for example by a slowly delivered body, shows its wait.", "A request that ends without any response after VLLM_REQUEST_STALL_WARNING_S logs a warning with the method, path and when its body was read." ], - "authors": ["@voipmonitor"] + "pull_requests": [ + 884 + ] } From bca7b17c6bc69d392fda1bf21e1799e13d2b69a3 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Sun, 20 Sep 2026 18:06:25 -0700 Subject: [PATCH 227/365] [Bugfix][Spec Decode] Stop dummy draft decode steps from writing KV through stale block-table rows (#56734) Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 (cherry picked from commit d2983f2f163249acbe3053c073cb9830290984a0) Signed-off-by: Martin Vit --- tests/v1/worker/test_gpu_block_table.py | 39 +++++++++++++++++++ vllm/v1/worker/gpu/block_table.py | 11 ++++-- .../spec_decode/autoregressive/speculator.py | 1 + .../multi_module_mtp/speculator.py | 1 + vllm/v1/worker/gpu/spec_decode/speculator.py | 14 +++++-- 5 files changed, 59 insertions(+), 7 deletions(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 1aa1da0eafef..1f9628a8aa90 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,6 +5,7 @@ import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.block_table import BlockTables pytestmark = pytest.mark.skipif( @@ -378,3 +379,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() + + +def test_dummy_request_slot_mapping_is_pad(): + """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request. + + A dummy draft decode step must not resolve slots through the persistent + block-table row of a real request slot, which may be stale and point at + blocks that are now in the prefix cache. + """ + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[4], + max_num_reqs=2, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[4], + ) + block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) + block_tables.apply_staged_writes() + query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) + positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) + + real = block_tables.compute_slot_mappings( + torch.tensor([0], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3] + + dummy = block_tables.compute_slot_mappings( + torch.tensor([-1], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert dummy[0].tolist() == [PAD_SLOT_ID] * 3 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 2c40aa28a459..cb9f9eda3669 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -342,7 +342,12 @@ def _compute_slot_mappings_kernel( group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) - num_blocks = tl.load(group_num_blocks_ptr + req_state_idx) + # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns + # no blocks: never read its block-table row and emit PAD for its tokens. + is_real_req = req_state_idx >= 0 + num_blocks = tl.load( + group_num_blocks_ptr + req_state_idx, mask=is_real_req, other=0 + ) start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): @@ -372,7 +377,7 @@ def _compute_slot_mappings_kernel( valid_block = token_mask & (block_indices < num_blocks) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local & valid_block, + mask=is_local & valid_block & is_real_req, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets @@ -380,5 +385,5 @@ def _compute_slot_mappings_kernel( slot_ids = tl.where(is_local, slot_ids, PAD_ID) slot_ids = tl.where(valid_block, slot_ids, PAD_ID) - slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index d7202c83cab8..8acb24869c46 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -299,6 +299,7 @@ def propose( input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) # Get the input ids and last token indices for the speculator. diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 949e71683cd8..f00af0321e17 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -168,6 +168,7 @@ def propose( input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) num_tokens = input_batch.num_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index cc0cf95c4208..ba1ba4fca3bd 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -522,6 +522,7 @@ def _copy_request_inputs( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, + dummy_run: bool = False, ) -> None: # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. # NOTE(woosuk): For draft sampling, we only consider the temperature @@ -531,10 +532,15 @@ def _copy_request_inputs( # affect the output distribution after rejection sampling. self.temperature.copy_(temperature) self.seeds.copy_(seeds) - self.idx_mapping[:num_reqs].copy_(idx_mapping) - # idx_mapping for CG padded requests points to -1, which is ignored - # during sampling to prevent writing stale values to draft logits. - self.idx_mapping[num_reqs:].fill_(-1) + # idx_mapping == -1 marks a row the drafter must not act on: sampling + # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows + # always get it; a dummy batch gets it for every row, since its arange + # idx_mapping names request-state slots with stale block tables. + if dummy_run: + self.idx_mapping.fill_(-1) + else: + self.idx_mapping[:num_reqs].copy_(idx_mapping) + self.idx_mapping[num_reqs:].fill_(-1) def _build_uniform_batch_dp_sync( self, From 5314ef5f118854f912075a6371f90e4061a23d7a Mon Sep 17 00:00:00 2001 From: derek Date: Thu, 24 Sep 2026 06:06:47 -0400 Subject: [PATCH 228/365] fix(qwen): allow full-TP DCP for qwen4_exp checkpoints Qwen3.8-Flash-Next checkpoints are published under two model types that load the same Qwen4Exp implementation: the official ones declare qwen4_exp / qwen4_exp_text, renamed ones declare qwen3_8_flash_next / qwen3_8_flash_next_text. Only Qwen3_8FlashNextTextConfig opted in to supports_full_tp_dcp_with_kv_gather, so ModelConfig rejected TP4/DCP4 (two KV heads) for qwen4_exp checkpoints while accepting it for renamed ones. Declare the capability on Qwen4ExpTextConfig so both spellings inherit it. Only the B12X QSA backend implements DCP; the non-B12X NVIDIA and AMD QSA backends still refuse DCP > 1 when they are constructed. --- tests/models/qwen4_exp/test_config.py | 9 +++++++++ vllm/models/qwen3_8_flash_next/config.py | 1 - vllm/models/qwen4_exp/config.py | 3 +++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index ef2dfbd14804..84a23107dfcf 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -13,6 +13,7 @@ Qwen3_5ForConditionalGenerationConfig, Qwen4ExpForConditionalGenerationConfig, ) +from vllm.models.qwen3_8_flash_next.config import Qwen3_8FlashNextTextConfig from vllm.models.qwen4_exp.config import ( Qwen4ExpConfig, Qwen4ExpTextConfig, @@ -285,3 +286,11 @@ def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None: ) assert second["query_start_loc"].data_ptr() == query_start_loc_ptr assert second["ngram_context"].data_ptr() == ngram_context_ptr + + +@pytest.mark.parametrize("config_cls", [Qwen4ExpTextConfig, Qwen3_8FlashNextTextConfig]) +def test_qwen_text_configs_allow_full_tp_dcp_with_kv_gather(config_cls) -> None: + """Qwen3.8 checkpoints declare either qwen4_exp_text or + qwen3_8_flash_next_text and load the same model. Both must opt in, or + TP4/DCP4 with two KV heads is refused for one spelling only.""" + assert config_cls.supports_full_tp_dcp_with_kv_gather diff --git a/vllm/models/qwen3_8_flash_next/config.py b/vllm/models/qwen3_8_flash_next/config.py index d00d4c59463c..ff22c33b234c 100644 --- a/vllm/models/qwen3_8_flash_next/config.py +++ b/vllm/models/qwen3_8_flash_next/config.py @@ -22,7 +22,6 @@ def __init__(self, **kwargs: Any) -> None: class Qwen3_8FlashNextTextConfig(Qwen4ExpTextConfig): model_type = "qwen3_8_flash_next_text" - supports_full_tp_dcp_with_kv_gather = True def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py index 82a1f5a5f555..6d9c3f8e0a47 100644 --- a/vllm/models/qwen4_exp/config.py +++ b/vllm/models/qwen4_exp/config.py @@ -70,6 +70,9 @@ class Qwen4ExpTextConfig(Qwen3NextConfig): model_type = "qwen4_exp_text" base_config_key = "text_config" keys_to_ignore_at_inference = ["past_key_values"] + # B12X QSA gathers KV across DCP ranks, so DCP may equal TP even with fewer + # KV heads than ranks. Other QSA backends refuse DCP > 1 on their own. + supports_full_tp_dcp_with_kv_gather = True def __init__( self, From d231ddaf9a6fd90d8bf2b993d5ee1c50195ce7c3 Mon Sep 17 00:00:00 2001 From: derek Date: Thu, 24 Sep 2026 06:07:23 -0400 Subject: [PATCH 229/365] docs(release): identify Qwen full-TP DCP fix as PR 886 --- .lil/changes/vllm-886.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 .lil/changes/vllm-886.json diff --git a/.lil/changes/vllm-886.json b/.lil/changes/vllm-886.json new file mode 100644 index 000000000000..62014998b5c1 --- /dev/null +++ b/.lil/changes/vllm-886.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-886", + "category": "fix", + "summary": "Allow decode context parallelism equal to tensor parallelism for official Qwen3.8-Flash-Next checkpoints", + "models": ["Qwen3.8-Flash-Next"], + "compatibility": "No user action required.", + "details": [ + "Checkpoints that declare model_type qwen4_exp, including the official FP8 and NVFP4 releases, now start at TP4/DCP4 instead of failing with \"decode-context-parallel-size (4) exceeds the maximum supported value (2)\".", + "Decode context parallelism still requires the B12X QSA backend." + ], + "pull_requests": [886] +} From c3034cdcf5d93771ca32d5c79890f0e29acc5281 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 10:12:51 +0000 Subject: [PATCH 230/365] chore(release): describe #885 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-885.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-885.json diff --git a/.lil/changes/vllm-885.json b/.lil/changes/vllm-885.json new file mode 100644 index 000000000000..3cc3c1ee4ad9 --- /dev/null +++ b/.lil/changes/vllm-885.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-885", + "category": "fix", + "summary": "Stop idle data-parallel ranks' dummy draft steps from writing KV through stale block tables", + "models": ["all"], + "compatibility": "Affects speculative decoding with data parallelism (DP > 1) on the model-runner v2 path; no user action required.", + "details": [ + "A dummy draft decode step on an idle DP rank could write KV through block-table rows left by finished requests, corrupting cache pages another request reuses.", + "Padded rows now map to the pad slot. Backport of vllm-project/vllm#56734 with the original author." + ], + "pull_requests": [885], + "authors": ["Yifan Qiao"], + "evidence": ["https://github.com/vllm-project/vllm/pull/56734"], + "requires": [] +} From 4bba478e24a710f79c6a0d798159823d42885905 Mon Sep 17 00:00:00 2001 From: derek Date: Thu, 24 Sep 2026 06:28:28 -0400 Subject: [PATCH 231/365] test(qwen): check the TP4/DCP4 gate through ModelConfig Build a ModelConfig from a two-KV-head config.json for each Qwen3.8 model type and run verify_with_parallel_config at TP4/DCP4, instead of only asserting the class attribute. The qwen4_exp case fails without the fix with the reported error. --- tests/models/qwen4_exp/test_config.py | 36 +++++++++++++++++++++------ 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 84a23107dfcf..52c6a1e7362b 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from importlib import import_module from types import SimpleNamespace from unittest.mock import patch @@ -8,12 +9,12 @@ import pytest import torch +from vllm.config import ModelConfig, ParallelConfig from vllm.config.speculative import SpeculativeConfig from vllm.model_executor.models.config import ( Qwen3_5ForConditionalGenerationConfig, Qwen4ExpForConditionalGenerationConfig, ) -from vllm.models.qwen3_8_flash_next.config import Qwen3_8FlashNextTextConfig from vllm.models.qwen4_exp.config import ( Qwen4ExpConfig, Qwen4ExpTextConfig, @@ -288,9 +289,30 @@ def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None: assert second["ngram_context"].data_ptr() == ngram_context_ptr -@pytest.mark.parametrize("config_cls", [Qwen4ExpTextConfig, Qwen3_8FlashNextTextConfig]) -def test_qwen_text_configs_allow_full_tp_dcp_with_kv_gather(config_cls) -> None: - """Qwen3.8 checkpoints declare either qwen4_exp_text or - qwen3_8_flash_next_text and load the same model. Both must opt in, or - TP4/DCP4 with two KV heads is refused for one spelling only.""" - assert config_cls.supports_full_tp_dcp_with_kv_gather +@pytest.mark.parametrize( + ("model_type", "architecture"), + [ + ("qwen4_exp", "Qwen4ExpForConditionalGeneration"), + ("qwen3_8_flash_next", "Qwen3_8FlashNextForConditionalGeneration"), + ], +) +def test_qwen_tp4_allows_dcp4_with_two_kv_heads( + tmp_path, model_type: str, architecture: str +) -> None: + """Qwen3.8 checkpoints declare either qwen4_exp or qwen3_8_flash_next and + load the same model. Both must pass the TP4/DCP4 gate with two KV heads.""" + text_config = _text_config(num_attention_heads=16, num_key_value_heads=2).to_dict() + text_config["model_type"] = f"{model_type}_text" + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": model_type, + "architectures": [architecture], + "text_config": text_config, + } + ) + ) + model_config = ModelConfig(model=str(tmp_path), skip_tokenizer_init=True) + model_config.verify_with_parallel_config( + ParallelConfig(tensor_parallel_size=4, decode_context_parallel_size=4) + ) From 638f90a1fc5464ce5838e0b97c59057ae62b2206 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 11:25:47 +0000 Subject: [PATCH 232/365] fix: DCP LSE correction for non-power-of-two ranks; pad aligned state rows with NULL_BLOCK_ID - correct_attn_out loaded tl.arange(0, N) with N = DCP world size; Triton needs a power-of-two extent, so DCP 3 or 6 did not compile. Use the next power of two and mask the extra lanes to -inf. - get_aligned_state_indices_multi_group_kernel filled rows without tokens (CUDA-graph padding) with -1. Recurrent/conv consumers skip NULL_BLOCK_ID (0), so -1 addressed memory before the state pool. Both found by the GLM-5.3-Flash TP3 port of BoxTops4Coochie (GLM5.3-Flash-TP3-Testing, kraken tp3 experimental/upstream). Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .../test_aligned_state_indices_padding.py | 52 +++++++++++++++++++ .../v1/attention/test_dcp_correct_attn_out.py | 34 ++++++++++++ vllm/v1/attention/ops/dcp.py | 12 ++++- vllm/v1/worker/mamba_utils.py | 5 +- 4 files changed, 100 insertions(+), 3 deletions(-) create mode 100644 tests/kernels/mamba/test_aligned_state_indices_padding.py create mode 100644 tests/v1/attention/test_dcp_correct_attn_out.py diff --git a/tests/kernels/mamba/test_aligned_state_indices_padding.py b/tests/kernels/mamba/test_aligned_state_indices_padding.py new file mode 100644 index 000000000000..f06337fa6fe6 --- /dev/null +++ b/tests/kernels/mamba/test_aligned_state_indices_padding.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Padded rows of aligned recurrent state indices use NULL_BLOCK_ID.""" + +import pytest +import torch + +from vllm.triton_utils import triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.worker.mamba_utils import get_aligned_state_indices_multi_group_kernel + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_padded_rows_point_at_the_null_block(): + block_size, num_groups, num_state_slots = 16, 2, 2 + block_tables = [ + torch.arange(1, 1 + 4 * 8, dtype=torch.int32, device="cuda").view(4, 8) + + 100 * group + for group in range(num_groups) + ] + pointers = torch.tensor( + [bt.data_ptr() for bt in block_tables], dtype=torch.int64, device="cuda" + ) + # Rows 2 and 3 are CUDA-graph padding: no tokens. + seq_lens = torch.tensor([17, 40, 0, 0], dtype=torch.int32, device="cuda") + indices = torch.full( + (num_groups, 4, num_state_slots), 12345, dtype=torch.int32, device="cuda" + ) + get_aligned_state_indices_multi_group_kernel[(1,)]( + pointers, + seq_lens, + indices, + block_tables[0].stride(0), + seq_lens.stride(0), + indices.stride(0), + indices.stride(1), + indices.stride(2), + 4, + CACHE_BLOCK_SIZE=block_size, + NUM_GROUPS=num_groups, + BLOCK_GROUPS=triton.next_power_of_2(num_groups), + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + BLOCK_ROWS=4, + num_warps=1, + ) + for group, bt in enumerate(block_tables): + for row, seq_len in enumerate([17, 40]): + first = (seq_len - 1) // block_size + expected = bt[row, first : first + num_state_slots] + assert torch.equal(indices[group, row], expected) + assert (indices[group, 2:] == NULL_BLOCK_ID).all() diff --git a/tests/v1/attention/test_dcp_correct_attn_out.py b/tests/v1/attention/test_dcp_correct_attn_out.py new file mode 100644 index 000000000000..18ba51919efd --- /dev/null +++ b/tests/v1/attention/test_dcp_correct_attn_out.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""correct_attn_out for every DCP world size, including non-powers of two.""" + +import pytest +import torch + +from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("world_size", [2, 3, 4, 6]) +@pytest.mark.parametrize("is_base_e", [True, False]) +def test_correct_attn_out_matches_reference(world_size, is_base_e): + torch.manual_seed(world_size) + tokens, heads, head_dim = 5, 4, 64 + out = torch.randn(tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16) + lses = torch.randn(world_size, tokens, heads, device="cuda") * 4 + # A rank with no keys for a row contributes -inf. + lses[world_size - 1, 0] = -float("inf") + rank = world_size - 1 + + expected_out = out.float().clone() + scale = 1.0 if is_base_e else torch.log(torch.tensor(2.0)).item() + lse = torch.logsumexp(lses * scale, dim=0) / scale + factor = torch.exp((lses[rank] - lse) * scale) + expected_out *= factor[..., None] + + corrected, final_lse = correct_attn_out( + out.clone(), lses, rank, CPTritonContext(), is_lse_base_on_e=is_base_e + ) + + torch.testing.assert_close(final_lse, lse, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(corrected.float(), expected_out, rtol=2e-2, atol=2e-2) diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 4b37e59c210b..4beff01bf469 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -132,6 +132,7 @@ def kernel( lse_idx, HEAD_DIM: tl.constexpr, N_ROUNDED: tl.constexpr, + BLOCK_N: tl.constexpr, IS_BASE_E: tl.constexpr, ): """ @@ -152,7 +153,9 @@ def kernel( batch_idx = tl.program_id(axis=0).to(tl.int64) head_idx = tl.program_id(axis=1).to(tl.int64) d_offsets = tl.arange(0, HEAD_DIM) - num_n_offsets = tl.arange(0, N_ROUNDED) + # N_ROUNDED is the actual LSE rank extent in this launcher. + # Triton requires a power-of-two vector; masked lanes contribute -inf. + num_n_offsets = tl.arange(0, BLOCK_N) # shape = [N] lse_offsets = ( @@ -162,7 +165,11 @@ def kernel( ) # calc final lse - lse = tl.load(lses_ptr + lse_offsets).to(tl.float32) + lse = tl.load( + lses_ptr + lse_offsets, + mask=num_n_offsets < N_ROUNDED, + other=-float("inf"), + ).to(tl.float32) lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse) lse_max = tl.max(lse, axis=0) lse_max = tl.where(lse_max == -float("inf"), 0, lse_max) @@ -340,6 +347,7 @@ def __call__( lses_stride_H=lses_stride_h, HEAD_DIM=head_dim, N_ROUNDED=n_rounded, + BLOCK_N=triton.next_power_of_2(n_rounded), IS_BASE_E=is_base_e, _runtime_launcher=None if self._warming else ctx.call_kernel, # CPTritonContext caches the non-constexpr positional prefix; derive diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 6b05709a033c..3a9a3ed386ec 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -88,7 +88,10 @@ def get_aligned_state_indices_multi_group_kernel( & (seq_lens[None, :, None] > 0) & valid_state_slot[None, None, :] ), - other=-1, + # Padding must use the recurrent/conv NULL_BLOCK_ID (reserved block 0). + # MTP0 reuses this buffer directly during full CUDA graph replay; + # -1 would be interpreted as a real state address before the pool. + other=0, ) tl.store( state_indices_ptr From 6a35d022a46a3e1db487ebd406a0269595fa0664 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 11:26:12 +0000 Subject: [PATCH 233/365] chore(release): describe #887 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-887.json | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .lil/changes/vllm-887.json diff --git a/.lil/changes/vllm-887.json b/.lil/changes/vllm-887.json new file mode 100644 index 000000000000..03e57dde59cf --- /dev/null +++ b/.lil/changes/vllm-887.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-887", + "category": "fix", + "summary": "DCP with 3 or 6 ranks compiles; padded recurrent state rows use the null block", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "The DCP attention-output correction required a power-of-two rank count, so DCP 3 or 6 failed to compile.", + "CUDA-graph padding rows of aligned recurrent state indices pointed at -1 instead of the reserved null block; consumers that skip the null block could address memory before the state pool.", + "Both were found by BoxTops4Coochie's GLM-5.3-Flash TP3 port." + ], + "pull_requests": [887], + "requires": [] +} From 6b30e92cdab5db4f5c4a40042778f995fc59283f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 11:42:50 +0000 Subject: [PATCH 234/365] Pad GLM-5.3 tensor-parallel axes for TP3 GLM-5.3-Flash is dimensioned in powers of two, so tensor_parallel_size=3 cannot shard its MLA heads, KDA heads or shared-expert width. Add a Glm5NextForCausalLMConfig parallelism hook (target and Glm5NextMTPModel draft) that pads only the axes TP does not divide: - MLA attention/KV heads 64 -> 72 (24 per rank, groups of 8) - KDA heads 64 -> 66 (22 per rank) - shared-expert width 2048 -> 2112 (704 per rank) Checkpoint sizes stay as original_* attributes. Routed experts are not padded, so a non-divisible routed width requires --enable-expert-parallel; a TP-sharded vision tower that cannot be split raises and points at --mm-encoder-tp-mode data / --language-model-only. The vocabulary already pads to 154944 in VocabParallelEmbedding. The hook also runs from verify_and_update_config so hybrid KDA state sizing sees the padded heads. Divisible TP sizes and non-CUDA platforms are exact no-ops. The model loads the logical checkpoint into the padded layouts through the existing allow_tp_padding loaders: MLA q_b/q, kv_b and o_proj, the shared-expert gate_up/down, the KDA input/f_b/g_b/o_proj projections, and padding-aware loaders for KDA dt_bias, A_log and the fused conv1d. Padded heads get zero Q/K/V, zero conv and decay parameters and zero o_proj input columns, so they contribute nothing to the output. CPU tests cover the config hook, zero-filled tails at every rank for column, row, merged, NVFP4-packed and KDA per-head loads, and equivalence of padded MLA-like attention and shared-expert MLP with the unpadded model. Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 (1M context) --- .../models/glm5next/test_glm53_tp_padding.py | 476 ++++++++++++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 82 ++- vllm/model_executor/models/config.py | 146 ++++++ vllm/models/glm5next/nvidia/attention.py | 13 + vllm/models/glm5next/nvidia/model.py | 11 +- 5 files changed, 722 insertions(+), 6 deletions(-) create mode 100644 tests/models/glm5next/test_glm53_tp_padding.py diff --git a/tests/models/glm5next/test_glm53_tp_padding.py b/tests/models/glm5next/test_glm53_tp_padding.py new file mode 100644 index 000000000000..70a928f6d1be --- /dev/null +++ b/tests/models/glm5next/test_glm53_tp_padding.py @@ -0,0 +1,476 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for GLM-5.3 tensor-parallel padding (TP3). + +The config hook pads MLA heads, KDA heads and the shared-expert width; the +model loads the logical checkpoint into the padded layouts with zero tails. +These tests emulate each TP rank on CPU by overriding the rank. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +import vllm.model_executor.layers.linear as linear_module +import vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn as kda_module +import vllm.model_executor.parameter as parameter_module +from vllm.config.model import ModelConfig +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.config import ( + MODELS_CONFIG_MAP, + Glm5NextForCausalLMConfig, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig + +TP = 3 + +# --------------------------------------------------------------------------- +# Config hook +# --------------------------------------------------------------------------- + + +def _glm53_model_config( + architecture: str = "Glm5NextForConditionalGeneration", + multimodal_config=None, +): + hf_config = Glm5NextConfig( + text_config={ + "num_attention_heads": 64, + "num_key_value_heads": 64, + "moe_intermediate_size": 2048, + "n_shared_experts": 1, + "n_routed_experts": 288, + "intermediate_size": 12288, + "first_k_dense_replace": 3, + "vocab_size": 154880, + "linear_attn_config": { + "head_dim": 128, + "num_heads": 64, + "short_conv_kernel_size": 4, + "gate_lower_bound": -5.0, + }, + }, + vision_config={"num_heads": 16}, + ) + model_config = SimpleNamespace( + architecture=architecture, + hf_config=hf_config, + hf_text_config=hf_config.text_config, + multimodal_config=multimodal_config, + model_arch_config=None, + ) + model_config.get_model_arch_config = lambda: SimpleNamespace( + total_num_attention_heads=hf_config.text_config.num_attention_heads + ) + return model_config + + +def _parallel_config(tp_size: int, enable_expert_parallel: bool = True): + return SimpleNamespace( + tensor_parallel_size=tp_size, enable_expert_parallel=enable_expert_parallel + ) + + +@pytest.fixture +def cuda_platform(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + +def test_glm53_architectures_use_the_padding_hook() -> None: + for architecture in ( + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + "Glm5NextMTPModel", + ): + assert MODELS_CONFIG_MAP[architecture] is Glm5NextForCausalLMConfig + + +@pytest.mark.parametrize("architecture", ["Glm5NextForConditionalGeneration"]) +def test_tp3_pads_glm53_geometry(cuda_platform, architecture) -> None: + model_config = _glm53_model_config(architecture) + parallel_config = _parallel_config(TP) + + # Idempotent: the hook runs from both verify paths. + ModelConfig._update_model_config_for_parallelism(model_config, parallel_config) + ModelConfig._update_model_config_for_parallelism(model_config, parallel_config) + + text = model_config.hf_text_config + assert (text.num_attention_heads, text.original_num_attention_heads) == (72, 64) + assert (text.num_key_value_heads, text.original_num_key_value_heads) == (72, 64) + assert (text.linear_num_heads, text.original_linear_num_heads) == (66, 64) + assert text.linear_attn_config["num_heads"] == 66 + assert text.linear_attn_config["original_num_heads"] == 64 + assert text.shared_expert_intermediate_size == 2112 + assert text.original_shared_expert_intermediate_size == 2048 + # Routed experts and the dense MLP are not padded. + assert text.moe_intermediate_size == 2048 + assert text.intermediate_size == 12288 + assert text.vocab_size == 154880 + # The multimodal wrapper forwards to the padded text config. + assert model_config.hf_config.num_attention_heads == 72 + assert model_config.hf_config.linear_num_heads == 66 + assert model_config.model_arch_config.total_num_attention_heads == 72 + # Per-rank sizes. + assert 72 // TP == 24 and 24 % 8 == 0 + assert 66 // TP == 22 + assert 2112 // TP == 704 and 704 % 64 == 0 + + +def test_mtp_draft_config_is_padded(cuda_platform) -> None: + model_config = _glm53_model_config("Glm5NextMTPModel") + ModelConfig._update_model_config_for_parallelism(model_config, _parallel_config(3)) + assert model_config.hf_text_config.num_attention_heads == 72 + assert model_config.hf_text_config.shared_expert_intermediate_size == 2112 + + +def test_verify_and_update_config_pads_before_hybrid_sizing(cuda_platform) -> None: + model_config = _glm53_model_config() + vllm_config = SimpleNamespace( + model_config=model_config, parallel_config=_parallel_config(TP) + ) + Glm5NextForCausalLMConfig.verify_and_update_config(vllm_config) + assert model_config.hf_text_config.linear_num_heads == 66 + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +def test_divisible_tp_sizes_are_noops(cuda_platform, tp_size) -> None: + model_config = _glm53_model_config() + before = model_config.hf_text_config.to_dict() + ModelConfig._update_model_config_for_parallelism( + model_config, _parallel_config(tp_size, enable_expert_parallel=False) + ) + text = model_config.hf_text_config + assert text.to_dict() == before + assert not hasattr(text, "original_num_attention_heads") + assert not hasattr(text, "shared_expert_intermediate_size") + assert model_config.model_arch_config is None + + +def test_tp3_without_expert_parallel_raises(cuda_platform) -> None: + model_config = _glm53_model_config() + with pytest.raises(ValueError, match="--enable-expert-parallel"): + ModelConfig._update_model_config_for_parallelism( + model_config, _parallel_config(TP, enable_expert_parallel=False) + ) + + +def test_tp3_vision_weights_mode_raises(cuda_platform) -> None: + def multimodal_config(mode: str, limit: int): + return SimpleNamespace( + mm_encoder_tp_mode=mode, get_limit_per_prompt=lambda _: limit + ) + + with pytest.raises(ValueError, match="mm-encoder-tp-mode data"): + ModelConfig._update_model_config_for_parallelism( + _glm53_model_config(multimodal_config=multimodal_config("weights", 1)), + _parallel_config(TP), + ) + for mode, limit in (("data", 1), ("weights", 0)): + model_config = _glm53_model_config( + multimodal_config=multimodal_config(mode, limit) + ) + ModelConfig._update_model_config_for_parallelism( + model_config, _parallel_config(TP) + ) + assert model_config.hf_text_config.num_attention_heads == 72 + + +def test_non_cuda_platforms_are_untouched(monkeypatch) -> None: + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + model_config = _glm53_model_config() + ModelConfig._update_model_config_for_parallelism(model_config, _parallel_config(3)) + assert model_config.hf_text_config.num_attention_heads == 64 + + +def test_tp3_vocab_storage_is_padded_with_zero_rows() -> None: + vocab, dim = 154880, 4 + checkpoint = torch.randn(vocab, dim) + shards = [] + for rank in range(TP): + group = SimpleNamespace(rank_in_group=rank, world_size=TP) + embedding = VocabParallelEmbedding(vocab, dim, parallel_group=group) + assert embedding.num_embeddings_padded == 154944 + embedding.weight_loader(embedding.weight, checkpoint) + shards.append(embedding.weight.data) + full = torch.cat(shards) + torch.testing.assert_close(full[:vocab], checkpoint, rtol=0, atol=0) + assert torch.count_nonzero(full[vocab:]) == 0 + + +# --------------------------------------------------------------------------- +# Padded loading +# --------------------------------------------------------------------------- + + +@pytest.fixture +def tp_rank(monkeypatch): + """Emulate a TP rank for layers and parameters built on CPU.""" + state = {"rank": 0} + + def set_rank(rank: int) -> None: + state["rank"] = rank + + for module in (linear_module, parameter_module, kda_module): + if hasattr(module, "get_tensor_model_parallel_rank"): + monkeypatch.setattr( + module, "get_tensor_model_parallel_rank", lambda: state["rank"] + ) + if hasattr(module, "get_tensor_model_parallel_world_size"): + monkeypatch.setattr( + module, "get_tensor_model_parallel_world_size", lambda: TP + ) + return set_rank + + +def _allow_padding(layer: torch.nn.Module) -> torch.nn.Module: + for param in layer.parameters(): + set_weight_attrs(param, {"allow_tp_padding": True}) + return layer + + +def _toy_mla_reference(x, wq, wkv, wo, heads: int, head_dim: int): + """MLA-like attention: head-major Q, interleaved per-head [K, V] rows.""" + tokens = x.shape[0] + q = (x @ wq.T).view(tokens, heads, head_dim).transpose(0, 1) + kv = (x @ wkv.T).view(tokens, heads, 2 * head_dim).transpose(0, 1) + k, v = kv.split(head_dim, dim=-1) + scores = q @ k.transpose(-1, -2) / head_dim**0.5 + out = torch.softmax(scores, dim=-1) @ v + return out.transpose(0, 1).reshape(tokens, heads * head_dim) @ wo.T + + +def test_padded_mla_heads_match_unpadded_attention(tp_rank) -> None: + torch.manual_seed(0) + hidden, head_dim, logical_heads, padded_heads = 16, 4, 4, 6 + wq = torch.randn(logical_heads * head_dim, hidden) + wkv = torch.randn(logical_heads * 2 * head_dim, hidden) + wo = torch.randn(hidden, logical_heads * head_dim) + x = torch.randn(5, hidden) + expected = _toy_mla_reference(x, wq, wkv, wo, logical_heads, head_dim) + + local_heads = padded_heads // TP + partial_sums = [] + padded_wo_columns = [] + for rank in range(TP): + tp_rank(rank) + q_proj = _allow_padding( + ColumnParallelLinear(hidden, padded_heads * head_dim, bias=False) + ) + kv_proj = _allow_padding( + ColumnParallelLinear(hidden, padded_heads * 2 * head_dim, bias=False) + ) + o_proj = _allow_padding( + RowParallelLinear(padded_heads * head_dim, hidden, bias=False) + ) + for layer in (q_proj, kv_proj, o_proj): + layer.weight.data.fill_(float("nan")) + q_proj.weight_loader(q_proj.weight, wq) + kv_proj.weight_loader(kv_proj.weight, wkv) + o_proj.weight_loader(o_proj.weight, wo) + padded_wo_columns.append(o_proj.weight.data) + + partial_sums.append( + _toy_mla_reference( + x, + q_proj.weight.data, + kv_proj.weight.data, + o_proj.weight.data, + local_heads, + head_dim, + ) + ) + # Rank-local shard boundaries follow the unpadded checkpoint. + rows = local_heads * head_dim + available = max(0, min(rows, logical_heads * head_dim - rank * rows)) + torch.testing.assert_close( + q_proj.weight.data[:available], wq[rank * rows : rank * rows + available] + ) + assert torch.count_nonzero(q_proj.weight.data[available:]) == 0 + assert torch.count_nonzero(o_proj.weight.data[:, available:]) == 0 + + # The all-reduce of the per-rank partial sums equals the unpadded model. + torch.testing.assert_close(sum(partial_sums), expected) + full_wo = torch.cat(padded_wo_columns, dim=1) + torch.testing.assert_close(full_wo[:, : wo.shape[1]], wo, rtol=0, atol=0) + assert torch.count_nonzero(full_wo[:, wo.shape[1] :]) == 0 + + +def test_padded_shared_expert_matches_unpadded_mlp(tp_rank) -> None: + torch.manual_seed(0) + hidden, logical, padded = 8, 20, 24 + gate = torch.randn(logical, hidden) + up = torch.randn(logical, hidden) + down = torch.randn(hidden, logical) + x = torch.randn(3, hidden) + expected = (F.silu(x @ gate.T) * (x @ up.T)) @ down.T + + outputs = [] + for rank in range(TP): + tp_rank(rank) + gate_up = _allow_padding( + MergedColumnParallelLinear(hidden, [padded] * 2, bias=False) + ) + down_proj = _allow_padding(RowParallelLinear(padded, hidden, bias=False)) + gate_up.weight.data.fill_(float("nan")) + down_proj.weight.data.fill_(float("nan")) + gate_up.weight_loader(gate_up.weight, gate, 0) + gate_up.weight_loader(gate_up.weight, up, 1) + down_proj.weight_loader(down_proj.weight, down) + local_gate, local_up = gate_up.weight.data.chunk(2) + h = F.silu(x @ local_gate.T) * (x @ local_up.T) + outputs.append(h @ down_proj.weight.data.T) + if rank == TP - 1: + tail = padded // TP - (logical - rank * (padded // TP)) + assert torch.count_nonzero(local_gate[-tail:]) == 0 + assert torch.count_nonzero(local_up[-tail:]) == 0 + assert torch.count_nonzero(down_proj.weight.data[:, -tail:]) == 0 + + torch.testing.assert_close(sum(outputs), expected) + + +def test_nvfp4_packed_row_parallel_weight_and_scales_pad_together(tp_rank) -> None: + """NVFP4 packs two values per byte and has one scale per 16 inputs.""" + out_features, logical, padded, group = 4, 2048, 2112, 16 + packed = torch.randint(1, 256, (out_features, logical // 2), dtype=torch.uint8) + scales = torch.randint(1, 256, (out_features, logical // group), dtype=torch.uint8) + local = padded // TP + + for rank in range(TP): + tp_rank(rank) + weight = PackedvLLMParameter( + data=torch.full((out_features, local // 2), 0xFF, dtype=torch.uint8), + input_dim=1, + output_dim=0, + packed_dim=1, + packed_factor=2, + weight_loader=lambda *_: None, + ) + scale = GroupQuantScaleParameter( + data=torch.full((out_features, local // group), 0xFF, dtype=torch.uint8), + input_dim=1, + output_dim=0, + weight_loader=lambda *_: None, + ) + for param in (weight, scale): + param.allow_tp_padding = True + weight.load_row_parallel_weight(packed) + scale.load_row_parallel_weight(scales) + + available = max(0, min(local, logical - rank * local)) + start = rank * local + torch.testing.assert_close( + weight.data[:, : available // 2], + packed[:, start // 2 : (start + available) // 2], + ) + torch.testing.assert_close( + scale.data[:, : available // group], + scales[:, start // group : (start + available) // group], + ) + # A zero byte is two FP4 zeros; a zero scale also zeroes the group. + assert torch.count_nonzero(weight.data[:, available // 2 :]) == 0 + assert torch.count_nonzero(scale.data[:, available // group :]) == 0 + + +def test_kda_per_head_loaders_zero_fill_padded_heads(tp_rank) -> None: + head_dim, logical_heads, padded_heads = 2, 4, 6 + local_heads = padded_heads // TP + projection = padded_heads * head_dim + dt_bias = torch.randn(logical_heads * head_dim) + a_log = torch.randn(logical_heads) + old_a_log = a_log.view(1, 1, logical_heads, 1) + conv = torch.randn(logical_heads * head_dim, 1, 4) + + for rank in range(TP): + tp_rank(rank) + heads_available = max(0, min(local_heads, logical_heads - rank * local_heads)) + rows_available = heads_available * head_dim + + dt_param = torch.full((local_heads * head_dim,), float("nan")) + kda_module._tp_padded_weight_loader(0)(dt_param, dt_bias) + start = rank * local_heads * head_dim + torch.testing.assert_close( + dt_param[:rows_available], dt_bias[start : start + rows_available] + ) + assert torch.count_nonzero(dt_param[rows_available:]) == 0 + + for source in (a_log, old_a_log): + a_param = torch.full((local_heads,), float("nan")) + kda_module.a_log_weight_loader(0, True)(a_param, source) + head_start = rank * local_heads + torch.testing.assert_close( + a_param[:heads_available], + a_log[head_start : head_start + heads_available], + ) + assert torch.count_nonzero(a_param[heads_available:]) == 0 + + conv_param = torch.full((3 * local_heads * head_dim, 1, 4), float("nan")) + loader = kda_module._make_fused_conv1d_weight_loader( + [projection] * 3, TP, rank, allow_tp_padding=True + ) + for shard_id in range(3): + loader(conv_param, conv + shard_id, shard_id) + for shard_id, shard in enumerate(conv_param.chunk(3)): + torch.testing.assert_close( + shard[:rows_available], + (conv + shard_id)[start : start + rows_available], + ) + assert torch.count_nonzero(shard[rows_available:]) == 0 + + +def test_kda_input_projection_pads_heads_and_replicates_f_a(tp_rank) -> None: + hidden, head_dim, logical_heads, padded_heads = 8, 2, 4, 6 + local_heads = padded_heads // TP + local_proj = local_heads * head_dim + checkpoint = { + "q": torch.randn(logical_heads * head_dim, hidden), + "k": torch.randn(logical_heads * head_dim, hidden), + "v": torch.randn(logical_heads * head_dim, hidden), + "b": torch.randn(logical_heads, hidden), + "f_a": torch.randn(head_dim, hidden), + } + + for rank in range(TP): + tp_rank(rank) + layer = kda_module._KimiGDNMergedColumnParallelLinear( + hidden, + [padded_heads * head_dim] * 3 + [padded_heads, head_dim], + replicated_shard_id=4, + tp_size=TP, + bias=False, + ) + _allow_padding(layer) + layer.weight.data.fill_(float("nan")) + for shard_id, name in enumerate(checkpoint): + layer.weight_loader(layer.weight, checkpoint[name], shard_id) + q, k, v, b, f_a = layer.weight.data.split( + [local_proj] * 3 + [local_heads, head_dim] + ) + heads_available = max(0, min(local_heads, logical_heads - rank * local_heads)) + rows = heads_available * head_dim + for local, name in ((q, "q"), (k, "k"), (v, "v")): + torch.testing.assert_close( + local[:rows], + checkpoint[name][rank * local_proj : rank * local_proj + rows], + ) + assert torch.count_nonzero(local[rows:]) == 0 + torch.testing.assert_close( + b[:heads_available], + checkpoint["b"][rank * local_heads : rank * local_heads + heads_available], + ) + assert torch.count_nonzero(b[heads_available:]) == 0 + torch.testing.assert_close(f_a, checkpoint["f_a"]) diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 4d409c6b2167..156dd0694d8f 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -21,7 +21,10 @@ default_weight_loader, sharded_weight_loader, ) -from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.parameter import ( + BasevLLMParameter, + copy_tensor_parallel_shard, +) from vllm.model_executor.utils import set_weight_attrs from vllm.model_executor.weight_transfer import copy_weight from vllm.platforms import current_platform @@ -247,8 +250,12 @@ def resolve_kda_prefill_backend( def a_log_weight_loader( shard_axis: int, + allow_tp_padding: bool = False, ) -> Callable[[torch.Tensor, torch.Tensor], None]: - """Load KDA A_log stored as either old 4D or current 1D weights.""" + """Load KDA A_log stored as either old 4D or current 1D weights. + + With ``allow_tp_padding``, heads past the checkpoint are zero-filled. + """ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: tp_rank = get_tensor_model_parallel_rank() @@ -264,16 +271,46 @@ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: ) loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + if allow_tp_padding: + copy_tensor_parallel_shard( + param.data, + loaded_weight, + shard_axis, + start_idx, + shard_size, + allow_padding=True, + ) + return None loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) return default_weight_loader(param, loaded_weight) return loader +def _tp_padded_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load a TP shard, zero-filling rows past the end of the checkpoint.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + shard_size = param.data.shape[shard_axis] + copy_tensor_parallel_shard( + param.data, + loaded_weight, + shard_axis, + get_tensor_model_parallel_rank() * shard_size, + shard_size, + allow_padding=True, + ) + + return loader + + def _make_fused_conv1d_weight_loader( dims: list[int], tp_size: int, tp_rank: int, + allow_tp_padding: bool = False, ) -> Callable[..., None]: sharded_dims = [dim // tp_size for dim in dims] @@ -287,8 +324,19 @@ def weight_loader( shard_size = sharded_dims[loaded_shard_id] source_start = tp_rank * shard_size target_start = sum(sharded_dims[:loaded_shard_id]) + destination = param.data[target_start : target_start + shard_size] + if allow_tp_padding: + copy_tensor_parallel_shard( + destination, + loaded_weight, + 0, + source_start, + shard_size, + allow_padding=True, + ) + return loaded_shard = loaded_weight[source_start : source_start + shard_size] - copy_weight(param.data[target_start : target_start + shard_size], loaded_shard) + copy_weight(destination, loaded_shard) return weight_loader @@ -432,6 +480,12 @@ def __init__( self.head_dim = kda_config["head_dim"] self.num_heads = kda_config["num_heads"] assert self.num_heads % self.tp_size == 0 + # TP padding (e.g. GLM-5.3 at TP3) appends zero heads after the + # checkpoint heads. Their projections, conv, decay and output columns + # load as zeros, so they keep a zero state and contribute nothing. + self.tp_padded = ( + kda_config.get("original_num_heads", self.num_heads) != self.num_heads + ) self.local_num_heads = divide(self.num_heads, self.tp_size) self.projection_size = self.head_dim * self.num_heads @@ -483,7 +537,14 @@ def __init__( torch.empty(self.local_projection_size, dtype=torch.float32) ) - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + set_weight_attrs( + self.dt_bias, + { + "weight_loader": _tp_padded_weight_loader(0) + if self.tp_padded + else sharded_weight_loader(0) + }, + ) # One packed parameter and cache let decode run a single conv update. # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. @@ -503,6 +564,7 @@ def __init__( [self.projection_size] * 3, self.tp_size, self.tp_rank, + allow_tp_padding=self.tp_padded, ) }, ) @@ -510,7 +572,10 @@ def __init__( self.A_log = nn.Parameter( torch.empty(self.local_num_heads, dtype=torch.float32) ) - set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + set_weight_attrs( + self.A_log, + {"weight_loader": a_log_weight_loader(0, self.tp_padded)}, + ) self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) if self.gate_lower_bound is not None: @@ -586,6 +651,13 @@ def __init__( quant_config=self.quant_config, prefix=f"{prefix}.o_proj", ) + if self.tp_padded: + linears = [self.in_proj_qkvgfab, self.f_b_proj, self.o_proj] + if not self.use_full_rank_gate: + linears.append(self.g_b_proj) + for linear in linears: + for param in linear.parameters(): + set_weight_attrs(param, {"allow_tp_padding": True}) self._b12x_preparation_prefix = prefix if ( self._b12x_kda_api is not None or self._b12x_prefill_api is not None diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index e6797dee9a21..0a6ff1fb5cef 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -451,6 +451,149 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: vllm_config.attention_config.backend = backend +class Glm5NextForCausalLMConfig(VerifyAndUpdateConfig): + """Pad GLM-5.3 axes that the tensor-parallel size does not divide. + + GLM-5.3 is dimensioned in powers of two, so TP3 needs physical padding: + MLA heads 64 -> 72 (whole groups of eight local heads), KDA heads + 64 -> 66 and the shared-expert width 2048 -> 2112. The checkpoint sizes + stay available as ``original_*`` attributes; the model loads them into the + padded layouts with zero tails, so padded heads and channels contribute + nothing. Routed experts are not padded and need expert parallelism. The + vocabulary needs no config change: ``VocabParallelEmbedding`` already pads + its storage to a multiple of the TP size. Divisible TP sizes are no-ops. + """ + + _MLA_LOCAL_HEAD_ALIGNMENT = 8 + _SHARED_EXPERT_LOCAL_ALIGNMENT = 64 + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + # Hybrid KDA state sizing runs before verify_with_parallel_config and + # reads the KDA head count, so pad here as well. The hook is idempotent. + if vllm_config.model_config is not None: + Glm5NextForCausalLMConfig.update_model_config_for_parallelism( + vllm_config.model_config, vllm_config.parallel_config + ) + + @staticmethod + def update_model_config_for_parallelism( + model_config: "ModelConfig", parallel_config: "ParallelConfig" + ) -> None: + from vllm.platforms import current_platform + + # Only the NVIDIA GLM-5.3 implementation loads padded layouts. + if not current_platform.is_cuda(): + return + cls = Glm5NextForCausalLMConfig + text_config: Any = model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + + def logical(name: str) -> int: + return getattr(text_config, f"original_{name}", getattr(text_config, name)) + + heads = logical("num_attention_heads") + kv_heads = logical("num_key_value_heads") + kda_heads = logical("linear_num_heads") + n_shared_experts = getattr(text_config, "n_shared_experts", None) + shared_size = getattr( + text_config, + "original_shared_expert_intermediate_size", + text_config.moe_intermediate_size * n_shared_experts + if n_shared_experts + else None, + ) + + padded_heads = ( + heads + if heads % tp_size == 0 + else round_up(heads, tp_size * cls._MLA_LOCAL_HEAD_ALIGNMENT) + ) + padded_kda_heads = round_up(kda_heads, tp_size) + padded_shared_size = shared_size + if shared_size is not None and shared_size % tp_size: + padded_shared_size = round_up( + shared_size, tp_size * cls._SHARED_EXPERT_LOCAL_ALIGNMENT + ) + if (padded_heads, padded_kda_heads, padded_shared_size) == ( + heads, + kda_heads, + shared_size, + ): + return + + if kv_heads != heads: + raise ValueError( + "GLM-5.3 TP padding expects num_key_value_heads == " + f"num_attention_heads for MLA, got {kv_heads} and {heads}." + ) + if ( + text_config.is_moe + and text_config.moe_intermediate_size % tp_size + and not parallel_config.enable_expert_parallel + ): + raise ValueError( + f"GLM-5.3 at tensor_parallel_size={tp_size} requires " + "--enable-expert-parallel: the routed-expert width " + f"({text_config.moe_intermediate_size}) is not divisible by " + f"{tp_size} and routed experts are not padded." + ) + if "dense" in text_config.mlp_layer_types and ( + text_config.intermediate_size % tp_size + ): + raise ValueError( + f"GLM-5.3 dense MLP width ({text_config.intermediate_size}) is " + f"not divisible by tensor_parallel_size={tp_size}." + ) + vision_config = getattr(model_config.hf_config, "vision_config", None) + multimodal_config = getattr(model_config, "multimodal_config", None) + if ( + vision_config is not None + and multimodal_config is not None + and multimodal_config.mm_encoder_tp_mode != "data" + and any( + multimodal_config.get_limit_per_prompt(modality) > 0 + for modality in ("image", "video") + ) + and vision_config.num_heads % tp_size + ): + raise ValueError( + f"The GLM-5.3 vision tower ({vision_config.num_heads} heads) " + f"cannot be sharded across tensor_parallel_size={tp_size}. Use " + "--mm-encoder-tp-mode data or --language-model-only." + ) + + text_config.original_num_attention_heads = heads + text_config.num_attention_heads = padded_heads + text_config.original_num_key_value_heads = kv_heads + text_config.num_key_value_heads = padded_heads + text_config.original_linear_num_heads = kda_heads + text_config.linear_num_heads = padded_kda_heads + linear_attn_config = getattr(text_config, "linear_attn_config", None) + if isinstance(linear_attn_config, dict): + text_config.linear_attn_config = { + **linear_attn_config, + "num_heads": padded_kda_heads, + "original_num_heads": kda_heads, + } + if shared_size is not None: + text_config.original_shared_expert_intermediate_size = shared_size + text_config.shared_expert_intermediate_size = padded_shared_size + + model_config.model_arch_config = model_config.get_model_arch_config() + logger.warning( + "Padded GLM-5.3 for TP%d: MLA heads %d -> %d, KDA heads %d -> %d, " + "shared-expert width %s -> %s.", + tp_size, + heads, + padded_heads, + kda_heads, + padded_kda_heads, + shared_size, + padded_shared_size, + ) + + class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig): """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface. @@ -1139,6 +1282,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "Gemma4ForCausalLM": Gemma4Config, "Gemma4ForConditionalGeneration": Gemma4Config, "Gemma4UnifiedForConditionalGeneration": Gemma4Config, + "Glm5NextForCausalLM": Glm5NextForCausalLMConfig, + "Glm5NextForConditionalGeneration": Glm5NextForCausalLMConfig, + "Glm5NextMTPModel": Glm5NextForCausalLMConfig, "GlmMoeDsaForCausalLM": GlmMoeDsaForCausalLM, "GptOssForCausalLM": GptOssForCausalLMConfig, "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig, diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py index fe8714b7c0e6..0cffb4720f1b 100644 --- a/vllm/models/glm5next/nvidia/attention.py +++ b/vllm/models/glm5next/nvidia/attention.py @@ -19,6 +19,7 @@ DeepSeekV2FusedQkvAProjLinear, yarn_get_mscale, ) +from vllm.model_executor.utils import set_weight_attrs from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -196,6 +197,18 @@ def __init__( is_sparse=self.is_sparse, topk_indices_buffer=topk_indices_buffer, ) + # TP padding (see Glm5NextForCausalLMConfig) appends zero heads after + # the checkpoint heads: zero Q and KV-up rows and zero o_proj input + # columns, so padded heads contribute nothing to the output. + if getattr(config, "original_num_attention_heads", num_heads) != num_heads: + for linear in ( + self.q_b_proj if q_lora_rank is not None else self.q_proj, + self.kv_b_proj, + self.o_proj, + ): + for param in linear.parameters(): + set_weight_attrs(param, {"allow_tp_padding": True}) + self.mla_attn = MultiHeadLatentAttentionWrapper( hidden_size, self.num_local_heads, diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index a662f6b4c904..e7acf6b728b4 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -77,6 +77,7 @@ maybe_prefix, sequence_parallel_chunk, ) +from vllm.model_executor.utils import set_weight_attrs from vllm.model_executor.weight_transfer import materialize_weight from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, @@ -259,10 +260,15 @@ def __init__( self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts + # TP padding (see Glm5NextForCausalLMConfig) widens the shared + # expert with zero gate/up rows and zero down_proj input columns. + padded_size = getattr( + config, "shared_expert_intermediate_size", intermediate_size + ) self.shared_experts = Glm5NextMLP( hidden_size=config.hidden_size, - intermediate_size=intermediate_size, + intermediate_size=padded_size, hidden_act=config.hidden_act, quant_config=quant_config, is_sequence_parallel=self.is_sequence_parallel, @@ -270,6 +276,9 @@ def __init__( prefix=f"{prefix}.shared_experts", swiglu_limit=swiglu_limit, ) + if padded_size != intermediate_size: + for param in self.shared_experts.parameters(): + set_weight_attrs(param, {"allow_tp_padding": True}) self.experts = FusedMoEFactory( shared_experts=self.shared_experts, From ed3e6b39ff7d4b75dd37cc70a1bf67130c873086 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 12:10:20 +0000 Subject: [PATCH 235/365] Run the GLM-5.3 vision tower data-parallel when TP cannot shard it The GLM-5.3 vision tower (16 heads) is constructed even with --language-model-only, so TP3 hit the "16 is not divisible by 3" assertion unless --mm-encoder-tp-mode data was passed. The TP padding hook now switches mm_encoder_tp_mode to "data" with a warning when the TP size does not divide the vision heads, instead of raising. Divisible TP sizes keep the configured mode. Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 (1M context) --- .../models/glm5next/test_glm53_tp_padding.py | 39 ++++++++++--------- vllm/model_executor/models/config.py | 37 +++++++++--------- 2 files changed, 40 insertions(+), 36 deletions(-) diff --git a/tests/models/glm5next/test_glm53_tp_padding.py b/tests/models/glm5next/test_glm53_tp_padding.py index 70a928f6d1be..28e2ce95a90d 100644 --- a/tests/models/glm5next/test_glm53_tp_padding.py +++ b/tests/models/glm5next/test_glm53_tp_padding.py @@ -169,25 +169,28 @@ def test_tp3_without_expert_parallel_raises(cuda_platform) -> None: ) -def test_tp3_vision_weights_mode_raises(cuda_platform) -> None: - def multimodal_config(mode: str, limit: int): - return SimpleNamespace( - mm_encoder_tp_mode=mode, get_limit_per_prompt=lambda _: limit - ) +@pytest.mark.parametrize("language_model_only", [False, True]) +def test_tp3_switches_unsplittable_vision_tower_to_data_parallel( + cuda_platform, language_model_only +) -> None: + # The tower is built even with --language-model-only, so switch regardless. + multimodal_config = SimpleNamespace( + mm_encoder_tp_mode="weights", language_model_only=language_model_only + ) + model_config = _glm53_model_config(multimodal_config=multimodal_config) + ModelConfig._update_model_config_for_parallelism(model_config, _parallel_config(TP)) + assert multimodal_config.mm_encoder_tp_mode == "data" + assert model_config.hf_text_config.num_attention_heads == 72 - with pytest.raises(ValueError, match="mm-encoder-tp-mode data"): - ModelConfig._update_model_config_for_parallelism( - _glm53_model_config(multimodal_config=multimodal_config("weights", 1)), - _parallel_config(TP), - ) - for mode, limit in (("data", 1), ("weights", 0)): - model_config = _glm53_model_config( - multimodal_config=multimodal_config(mode, limit) - ) - ModelConfig._update_model_config_for_parallelism( - model_config, _parallel_config(TP) - ) - assert model_config.hf_text_config.num_attention_heads == 72 + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_splittable_vision_tower_keeps_weights_mode(cuda_platform, tp_size) -> None: + multimodal_config = SimpleNamespace(mm_encoder_tp_mode="weights") + ModelConfig._update_model_config_for_parallelism( + _glm53_model_config(multimodal_config=multimodal_config), + _parallel_config(tp_size), + ) + assert multimodal_config.mm_encoder_tp_mode == "weights" def test_non_cuda_platforms_are_untouched(monkeypatch) -> None: diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 0a6ff1fb5cef..8c5ac89c6de8 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -492,6 +492,25 @@ def update_model_config_for_parallelism( def logical(name: str) -> int: return getattr(text_config, f"original_{name}", getattr(text_config, name)) + # The vision tower is built even with --language-model-only; when TP + # cannot shard its heads, run it data-parallel (replicated) instead. + vision_config = getattr(model_config.hf_config, "vision_config", None) + multimodal_config = getattr(model_config, "multimodal_config", None) + if ( + vision_config is not None + and multimodal_config is not None + and multimodal_config.mm_encoder_tp_mode != "data" + and vision_config.num_heads % tp_size + ): + logger.warning( + "The GLM-5.3 vision tower (%d heads) cannot be sharded across " + "TP%d; switching mm_encoder_tp_mode from %r to 'data'.", + vision_config.num_heads, + tp_size, + multimodal_config.mm_encoder_tp_mode, + ) + multimodal_config.mm_encoder_tp_mode = "data" + heads = logical("num_attention_heads") kv_heads = logical("num_key_value_heads") kda_heads = logical("linear_num_heads") @@ -545,24 +564,6 @@ def logical(name: str) -> int: f"GLM-5.3 dense MLP width ({text_config.intermediate_size}) is " f"not divisible by tensor_parallel_size={tp_size}." ) - vision_config = getattr(model_config.hf_config, "vision_config", None) - multimodal_config = getattr(model_config, "multimodal_config", None) - if ( - vision_config is not None - and multimodal_config is not None - and multimodal_config.mm_encoder_tp_mode != "data" - and any( - multimodal_config.get_limit_per_prompt(modality) > 0 - for modality in ("image", "video") - ) - and vision_config.num_heads % tp_size - ): - raise ValueError( - f"The GLM-5.3 vision tower ({vision_config.num_heads} heads) " - f"cannot be sharded across tensor_parallel_size={tp_size}. Use " - "--mm-encoder-tp-mode data or --language-model-only." - ) - text_config.original_num_attention_heads = heads text_config.num_attention_heads = padded_heads text_config.original_num_key_value_heads = kv_heads From 920d8c5f0c3a4e2033b0d39e65f3d733c66b404b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 12:19:23 +0000 Subject: [PATCH 236/365] Add GLM-5.3 TP3 FP8 decode and low-latency GEMM selection Port the community TP3 decode work onto the padded TP3 geometry: glm53_fp8_dense.py (VLLM_GLM53_FP8_DENSE=1, default off): after loading, quantize the large BF16 projections -- KDA in_proj/o_proj, DSA q_b/o_proj, the dense FFN (VLLM_GLM53_FP8_FFN), the LM head (VLLM_GLM53_FP8_LM_HEAD) and the MTP draft's DSA projections (VLLM_GLM53_FP8_MTP) -- to FP8 E4M3 with per-channel scales and run batches of at most VLLM_GLM53_FP8_DENSE_MAX_M (32) rows on Marlin W8A16. Larger batches keep the BF16 GEMM, or with VLLM_GLM53_FP8_PREFILL=w8a8 run CUTLASS FP8 on the same FP8 weights and release the BF16 copy. glm53_low_latency_gemm.py (VLLM_GLM53_LOW_LATENCY_GEMM, on by default at TP3 only): route the measured (N, K, M<=8) BF16 decode shapes to the CuTe shape-dynamic skinny GEMM on SM120, with configs validated against the kernel's shape checks and compiled in the CuTeDSL warmup before capture. The MTP eh_proj entry is dropped: it is a replicated nn.Linear here. Both tables are TP3 per-rank shapes and collide with other projections at TP2/TP4 (for example (4096, 4096) is the TP3 dense FFN down_proj but the TP4 DSA o_proj), so both paths only engage at tensor_parallel_size == 3. The LM head wrapper skips runtime-quantized heads, and neither wrapper accepts pre-processed weights. Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 (1M context) --- .../models/glm5next/test_glm53_decode_gemm.py | 315 +++++++++++++++ .../models/glm5next/nvidia/glm53_fp8_dense.py | 376 ++++++++++++++++++ .../glm5next/nvidia/glm53_low_latency_gemm.py | 154 +++++++ vllm/models/glm5next/nvidia/model.py | 6 + vllm/models/glm5next/nvidia/mtp.py | 5 + 5 files changed, 856 insertions(+) create mode 100644 tests/models/glm5next/test_glm53_decode_gemm.py create mode 100644 vllm/models/glm5next/nvidia/glm53_fp8_dense.py create mode 100644 vllm/models/glm5next/nvidia/glm53_low_latency_gemm.py diff --git a/tests/models/glm5next/test_glm53_decode_gemm.py b/tests/models/glm5next/test_glm53_decode_gemm.py new file mode 100644 index 000000000000..34c34672fd73 --- /dev/null +++ b/tests/models/glm5next/test_glm53_decode_gemm.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for the GLM-5.3 TP3 FP8 decode and low-latency GEMM selection. + +Layers are built on the meta device; only flag parsing, TP/shape guards and +method selection are exercised. The kernels themselves need a GPU. +""" + +import pytest +import torch +from torch import nn + +import vllm.model_executor.layers.linear as linear_module +import vllm.model_executor.layers.vocab_parallel_embedding as vocab_module +import vllm.model_executor.parameter as parameter_module +import vllm.models.glm5next.nvidia.glm53_fp8_dense as fp8 +import vllm.models.glm5next.nvidia.glm53_low_latency_gemm as llg +from vllm.model_executor.layers.linear import ( + ReplicatedLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform + +FP8_ENVS = ( + "VLLM_GLM53_FP8_DENSE", + "VLLM_GLM53_FP8_LM_HEAD", + "VLLM_GLM53_FP8_FFN", + "VLLM_GLM53_FP8_MTP", + "VLLM_GLM53_FP8_PREFILL", + "VLLM_GLM53_FP8_DENSE_MAX_M", + "VLLM_GLM53_LOW_LATENCY_GEMM", +) + +# Per-rank BF16 projection shapes (N, K) of one GLM-5.3 decoder slice. +TP3_SHAPES = { + "kda_in_proj": (8598, 4096), + "kda_o_proj": (4096, 2816), + "kda_g_a_proj": (128, 4096), + "dsa_fused_qkv_a": (2048, 4096), + "dsa_q_b_proj": (6144, 1536), + "dsa_o_proj": (4096, 6144), + "ffn_gate_up": (8192, 4096), + "ffn_down": (4096, 4096), + "shared_down": (4096, 704), + "indexer_weights_proj": (32, 4096), +} +# At TP4 the DSA o_proj is (4096, 4096), the TP3 dense FFN down_proj shape. +TP4_SHAPES = { + "kda_in_proj": (6288, 4096), + "dsa_q_b_proj": (4096, 1536), + "dsa_o_proj": (4096, 4096), + "ffn_gate_up": (6144, 4096), + "dsa_fused_qkv_a": (2048, 4096), + "indexer_weights_proj": (32, 4096), +} + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + for name in FP8_ENVS: + monkeypatch.delenv(name, raising=False) + # Replicated layers built on CPU still ask for the TP rank. + for module in (linear_module, parameter_module): + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + + +@pytest.fixture +def tp_size(monkeypatch): + def set_tp(size: int) -> None: + for module in (fp8, llg): + monkeypatch.setattr( + module, "get_tensor_model_parallel_world_size", lambda: size + ) + + set_tp(3) + return set_tp + + +def _toy_model(shapes: dict[str, tuple[int, int]]) -> nn.Module: + model = nn.Module() + with torch.device("meta"): + for name, (n, k) in shapes.items(): + setattr( + model, + name, + ReplicatedLinear(k, n, bias=False, params_dtype=torch.bfloat16), + ) + return model + + +def _methods(model: nn.Module) -> dict[str, type]: + return { + name: type(child.quant_method) + for name, child in model.named_children() + if hasattr(child, "quant_method") + } + + +def _all_plain(model: nn.Module) -> bool: + return all(cls is UnquantizedLinearMethod for cls in _methods(model).values()) + + +# --------------------------------------------------------------------------- +# FP8 dense decode +# --------------------------------------------------------------------------- + + +def test_fp8_dense_is_off_by_default(tp_size) -> None: + model = _toy_model(TP3_SHAPES) + assert not fp8.enabled() + assert fp8.enable_glm53_fp8_dense(model) == 0 + assert fp8.enable_glm53_fp8_dense(model, draft=True) == 0 + assert _all_plain(model) + + +@pytest.mark.parametrize("size", [1, 2, 4, 8]) +def test_fp8_dense_skips_non_tp3(monkeypatch, tp_size, size) -> None: + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE", "1") + tp_size(size) + model = _toy_model(TP4_SHAPES) + assert fp8.enable_glm53_fp8_dense(model) == 0 + assert fp8.enable_glm53_fp8_dense(model, draft=True) == 0 + assert _all_plain(model) + + +def test_fp8_dense_wraps_tp3_projections(monkeypatch, tp_size) -> None: + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE", "1") + model = _toy_model(TP3_SHAPES) + assert fp8.enable_glm53_fp8_dense(model) == 6 + methods = _methods(model) + wrapped = { + name for name, cls in methods.items() if cls is fp8.GLM53Fp8DecodeLinearMethod + } + assert wrapped == { + "kda_in_proj", + "kda_o_proj", + "dsa_q_b_proj", + "dsa_o_proj", + "ffn_gate_up", + "ffn_down", + } + # Small, shared-expert and fused_qkv_a projections stay BF16. + for name in ("kda_g_a_proj", "dsa_fused_qkv_a", "shared_down"): + assert methods[name] is UnquantizedLinearMethod + # Wrapping twice is a no-op. + assert fp8.enable_glm53_fp8_dense(model) == 0 + + +def test_fp8_dense_sub_flags(monkeypatch, tp_size) -> None: + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE", "1") + monkeypatch.setenv("VLLM_GLM53_FP8_FFN", "0") + model = _toy_model(TP3_SHAPES) + assert fp8.enable_glm53_fp8_dense(model) == 4 + assert _methods(model)["ffn_down"] is UnquantizedLinearMethod + + monkeypatch.setenv("VLLM_GLM53_FP8_MTP", "0") + draft = _toy_model({"dsa_o_proj": TP3_SHAPES["dsa_o_proj"]}) + assert fp8.enable_glm53_fp8_dense(draft, draft=True) == 0 + monkeypatch.setenv("VLLM_GLM53_FP8_MTP", "1") + assert fp8.enable_glm53_fp8_dense(draft, draft=True) == 1 + + +def test_fp8_prefill_and_batch_flags(monkeypatch) -> None: + assert fp8.prefill_mode() == "bf16" + assert fp8.max_m() == 32 + monkeypatch.setenv("VLLM_GLM53_FP8_PREFILL", " W8A8 ") + assert fp8.prefill_mode() == "w8a8" + monkeypatch.setenv("VLLM_GLM53_FP8_PREFILL", "fp8") + with pytest.raises(ValueError, match="bf16 or w8a8"): + fp8.prefill_mode() + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE_MAX_M", "0") + with pytest.raises(ValueError): + fp8.max_m() + + +def test_fp8_decode_eligibility() -> None: + x = torch.empty(32, 8, dtype=torch.bfloat16) + assert fp8._decode_eligible(x, None, 32) + assert not fp8._decode_eligible(torch.empty(33, 8, dtype=torch.bfloat16), None, 32) + assert not fp8._decode_eligible(x.float(), None, 32) + assert not fp8._decode_eligible(x, torch.empty(8), 32) + assert not fp8._decode_eligible(torch.empty(0, 8, dtype=torch.bfloat16), None, 32) + + +def _lm_head(monkeypatch, size: int) -> ParallelLMHead: + monkeypatch.setattr(vocab_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + vocab_module, "get_tensor_model_parallel_world_size", lambda: size + ) + with torch.device("meta"): + return ParallelLMHead(154880, 4096, params_dtype=torch.bfloat16) + + +def test_fp8_lm_head_only_at_tp3(monkeypatch, tp_size) -> None: + head = _lm_head(monkeypatch, 3) + assert tuple(head.weight.shape) == fp8.GLM53_TP3_LM_HEAD_SHAPE + assert not fp8.enable_glm53_fp8_lm_head(head) + + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE", "1") + monkeypatch.setenv("VLLM_GLM53_FP8_LM_HEAD", "0") + assert not fp8.enable_glm53_fp8_lm_head(head) + monkeypatch.setenv("VLLM_GLM53_FP8_LM_HEAD", "1") + assert fp8.enable_glm53_fp8_lm_head(head) + method = head.quant_method + assert isinstance(method, fp8.GLM53Fp8DecodeLMHeadMethod) + # LogitsProcessor must route through apply(), not its raw-weight branch. + assert not isinstance(method, UnquantizedEmbeddingMethod | UnquantizedLinearMethod) + # The BF16 weight stays for the MTP draft head copy and large batches. + assert head.weight.dtype == torch.bfloat16 + + tp_size(4) + head4 = _lm_head(monkeypatch, 4) + assert not fp8.enable_glm53_fp8_lm_head(head4) + assert type(head4.quant_method) is UnquantizedEmbeddingMethod + + +# --------------------------------------------------------------------------- +# Low-latency GEMM plan +# --------------------------------------------------------------------------- + + +@pytest.fixture +def sm120(monkeypatch): + requested: list = [] + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability_family", lambda family: family == 120 + ) + monkeypatch.setattr(llg.shape_dynamic_skinny_gemm, "is_available", lambda: True) + monkeypatch.setattr( + llg.shape_dynamic_skinny_gemm, + "request_warmup_configs", + lambda dtype, configs: requested.append((dtype, set(configs))), + ) + return requested + + +def test_low_latency_plan_configs_fit_the_kernel() -> None: + for (n, k), ms in llg.GLM53_TP3_PLAN.items(): + plan = llg.plan_for(n, k) + assert set(plan) == set(ms), (n, k) + for m, config in plan.items(): + assert llg._config_fits(config, m, n, k) + + +def test_low_latency_gemm_swaps_tp3_shapes(tp_size, sm120) -> None: + model = _toy_model(TP3_SHAPES) + assert llg.enable_glm53_low_latency_gemm(model, torch.bfloat16) == 7 + methods = _methods(model) + for name in ("ffn_gate_up", "ffn_down", "shared_down"): + assert methods[name] is UnquantizedLinearMethod + assert methods["kda_in_proj"] is llg.GLM53LowLatencyLinearMethod + assert len(sm120) == 1 and sm120[0][0] is torch.bfloat16 + + +@pytest.mark.parametrize("size", [1, 2, 4, 8]) +def test_low_latency_gemm_is_tp3_only(monkeypatch, tp_size, sm120, size) -> None: + monkeypatch.setenv("VLLM_GLM53_LOW_LATENCY_GEMM", "1") + tp_size(size) + model = _toy_model(TP4_SHAPES) + assert llg.enable_glm53_low_latency_gemm(model, torch.bfloat16) == 0 + assert _all_plain(model) + assert not sm120 + + +def test_low_latency_gemm_opt_out_and_guards(monkeypatch, tp_size, sm120) -> None: + model = _toy_model(TP3_SHAPES) + assert llg.enable_glm53_low_latency_gemm(model, torch.float16) == 0 + monkeypatch.setenv("VLLM_GLM53_LOW_LATENCY_GEMM", "0") + assert llg.enable_glm53_low_latency_gemm(model, torch.bfloat16) == 0 + monkeypatch.delenv("VLLM_GLM53_LOW_LATENCY_GEMM") + monkeypatch.setattr( + current_platform, "is_device_capability_family", lambda family: False + ) + assert llg.enable_glm53_low_latency_gemm(model, torch.bfloat16) == 0 + assert _all_plain(model) + + +def test_low_latency_dispatch_uses_measured_batch_sizes(monkeypatch) -> None: + method = llg.GLM53LowLatencyLinearMethod(llg.plan_for(2048, 4096)) + calls: list[int] = [] + + def fake_skinny_gemm(x, weight, config): + calls.append(config.num_rows) + return "skinny" + + monkeypatch.setattr(llg, "shape_dynamic_skinny_gemm", fake_skinny_gemm) + monkeypatch.setattr( + UnquantizedLinearMethod, "apply", lambda self, layer, x, bias=None: "cublas" + ) + layer = nn.Module() + layer.weight = nn.Parameter(torch.zeros(2048, 4096, dtype=torch.bfloat16)) + for m, expected in ((1, "skinny"), (2, "skinny"), (3, "cublas"), (4, "cublas")): + x = torch.zeros(m, 4096, dtype=torch.bfloat16) + assert method.apply(layer, x) == expected + x = torch.zeros(1, 4096, dtype=torch.bfloat16) + assert method.apply(layer, x, torch.zeros(2048)) == "cublas" + assert calls == [1, 2] + + +def test_fp8_wraps_low_latency_projections(monkeypatch, tp_size, sm120) -> None: + monkeypatch.setenv("VLLM_GLM53_FP8_DENSE", "1") + model = _toy_model(TP3_SHAPES) + llg.enable_glm53_low_latency_gemm(model, torch.bfloat16) + fp8.enable_glm53_fp8_dense(model) + method = model.kda_in_proj.quant_method + assert isinstance(method, fp8.GLM53Fp8DecodeLinearMethod) + # Batches above the FP8 threshold fall back to the wrapped method. + assert isinstance(method._inner, llg.GLM53LowLatencyLinearMethod) + assert type(model.dsa_fused_qkv_a.quant_method) is llg.GLM53LowLatencyLinearMethod diff --git a/vllm/models/glm5next/nvidia/glm53_fp8_dense.py b/vllm/models/glm5next/nvidia/glm53_fp8_dense.py new file mode 100644 index 000000000000..f9ceafb17347 --- /dev/null +++ b/vllm/models/glm5next/nvidia/glm53_fp8_dense.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash TP3 FP8 weight-only decode for the large BF16 projections. + +``modelopt_mixed`` keeps the attention and KDA projections, the dense FFN and +the LM head in BF16. At decode batch sizes those GEMMs stream their weights, +so halving the bytes nearly halves their time. Measured per TP3 rank at M=4 +(RTX PRO 6000, weights from HBM): KDA ``in_proj`` 44.4 -> 26.9 us, DSA +``o_proj`` 31.8 -> 19.3 us, KDA ``o_proj`` 16.5 -> 10.7 us, LM head +261 -> 137 us. Small projections are slower under Marlin and stay BF16, as +do the shared-expert projections and the DSA ``fused_qkv_a_proj``. + +After loading, each selected weight is quantized to FP8 E4M3 with one scale per +output channel and packed for the Marlin W8A16 kernel (activations stay BF16). +Batches of at most ``VLLM_GLM53_FP8_DENSE_MAX_M`` rows (default 32, every +decode and verify step up to 8 sequences x 4 positions) use it. Larger batches +use the original BF16 GEMM, or with ``VLLM_GLM53_FP8_PREFILL=w8a8`` CUTLASS FP8 +on a row-major copy of the same FP8 weights with per-token activation scales +(the BF16 weight is then released). TP padding rows and columns are zero, so +they quantize to exact zeros. + +The decoder projections expose the packed FP8 tensor as ``layer.weight`` and +keep the BF16 original in a private attribute, so the GLM-5.3 L2 prefetcher +(which skips private attributes) streams the bytes decode reads. The LM head +keeps its BF16 ``weight`` (the MTP draft head copies it) and holds the FP8 +copy in its quant method. Its wrapper is not an unquantized method, so +``LogitsProcessor`` calls ``apply()`` instead of the B12X vocabulary +projection; batches above the threshold use the BF16 cuBLAS projection. A +runtime-quantized LM head (``VLLM_MXFP8_LM_HEAD`` / ``VLLM_MTP_NVFP4_LM_HEAD``, +the only heads ``VLLM_LM_HEAD_A16`` affects) is left untouched. + +The shape tables are per-rank TP3 shapes, and several of them name different +projections at TP2/TP4 (for example (4096, 4096) is the dense FFN down_proj at +TP3 but the DSA o_proj at TP4 and the KDA o_proj at TP2), so the path only +engages at tensor_parallel_size == 3 and is a no-op elsewhere. + +Flags (all read at model construction): + +- ``VLLM_GLM53_FP8_DENSE=1`` enables the path (default off). +- ``VLLM_GLM53_FP8_LM_HEAD``, ``VLLM_GLM53_FP8_FFN`` (the three dense FFN + layers) and ``VLLM_GLM53_FP8_MTP`` (the MTP draft's DSA projections) default + on once enabled; set them to 0 to leave those BF16. +- ``VLLM_GLM53_FP8_PREFILL=bf16|w8a8`` (default ``bf16``). +- ``VLLM_GLM53_FP8_DENSE_MAX_M`` (default 32). + +This changes decode numerics of layers the checkpoint left in BF16 and needs +quality qualification before production use. +""" + +from __future__ import annotations + +import os + +import torch +from torch import nn + +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, +) + +logger = init_logger(__name__) + +# The per-rank shapes below are TP3 physical shapes (see the module docstring). +SUPPORTED_TP_SIZE = 3 + +# (N, K) per TP3 rank of the BF16 weights to convert. +GLM53_TP3_FP8_SHAPES: dict[tuple[int, int], str] = { + (8598, 4096): "KDA in_proj_qkvgfab", + (4096, 2816): "KDA o_proj", + (4096, 6144): "DSA o_proj", + (6144, 1536): "DSA q_b_proj", +} +# The three dense FFN layers (first_k_dense_replace=3, intermediate 12288/3). +GLM53_TP3_FP8_FFN_SHAPES: dict[tuple[int, int], str] = { + (8192, 4096): "dense FFN gate_up", + (4096, 4096): "dense FFN down", +} +# Padded vocabulary 154944 / 3 by hidden size 4096. +GLM53_TP3_LM_HEAD_SHAPE = (51648, 4096) + + +def _flag(name: str, default: str) -> bool: + return os.getenv(name, default).strip() not in ("", "0") + + +def enabled() -> bool: + return _flag("VLLM_GLM53_FP8_DENSE", "0") + + +def prefill_mode() -> str: + """How batches above the decode threshold run: ``bf16`` or ``w8a8``.""" + mode = os.getenv("VLLM_GLM53_FP8_PREFILL", "bf16").strip().lower() + if mode not in ("bf16", "w8a8"): + raise ValueError(f"VLLM_GLM53_FP8_PREFILL must be bf16 or w8a8; got {mode!r}") + return mode + + +def max_m() -> int: + value = int(os.getenv("VLLM_GLM53_FP8_DENSE_MAX_M", "32")) + if value <= 0: + raise ValueError("VLLM_GLM53_FP8_DENSE_MAX_M must be positive") + return value + + +def _tp_supported(what: str) -> bool: + tp_size = get_tensor_model_parallel_world_size() + if tp_size == SUPPORTED_TP_SIZE: + return True + logger.warning_once( + "GLM-5.3 FP8 decode (%s) is measured and keyed on TP%d shapes; " + "tensor_parallel_size=%d keeps BF16.", + what, + SUPPORTED_TP_SIZE, + tp_size, + ) + return False + + +def fp8_dense_shapes() -> dict[tuple[int, int], str]: + shapes = dict(GLM53_TP3_FP8_SHAPES) + if _flag("VLLM_GLM53_FP8_FFN", "1"): + shapes.update(GLM53_TP3_FP8_FFN_SHAPES) + return shapes + + +def _pack_marlin_fp8(weight: torch.Tensor, keep_rowmajor: bool = False) -> nn.Module: + """Quantize an (N, K) BF16 weight per output channel and pack it for Marlin. + + With ``keep_rowmajor`` the same FP8 values are also kept row-major with an + (N, 1) scale for the CUTLASS prefill path, so decode and prefill read + identical quantized weights. + """ + from vllm import _custom_ops as ops + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + prepare_fp8_layer_for_marlin, + ) + + n, k = weight.shape + qweight, scale = ops.scaled_fp8_quant( + weight.contiguous(), use_per_token_if_dynamic=True + ) + packed = nn.Module() + packed.rowmajor = qweight if keep_rowmajor else None + packed.scale_col = scale.reshape(n, 1).to(torch.float32) if keep_rowmajor else None + packed.weight = nn.Parameter(qweight.t().contiguous(), requires_grad=False) + packed.weight_scale = nn.Parameter( + scale.reshape(1, n).to(torch.float32), requires_grad=False + ) + packed.input_size_per_partition = k + packed.output_size_per_partition = n + packed.orig_dtype = weight.dtype + prepare_fp8_layer_for_marlin(packed, size_k_first=True) + return packed + + +def _decode_eligible(x: torch.Tensor, bias: torch.Tensor | None, limit: int) -> bool: + return ( + bias is None + and x.dim() == 2 + and x.dtype == torch.bfloat16 + and 0 < x.shape[0] <= limit + ) + + +class GLM53Fp8DecodeLinearMethod(UnquantizedLinearMethod): + """Marlin FP8 W8A16 for decode-sized batches, the original method otherwise.""" + + # The packed weight is derived from the BF16 checkpoint after loading. + supports_pre_processed_weights = False + + def __init__(self, inner: UnquantizedLinearMethod, label: str) -> None: + super().__init__() + self._inner = inner + self._label = label + self._max_m = max_m() + self._prefill = prefill_mode() + + def process_weights_after_loading(self, layer: nn.Module) -> None: + self._inner.process_weights_after_loading(layer) + bf16 = layer.weight.data + n, k = bf16.shape + w8a8 = self._prefill == "w8a8" + packed = _pack_marlin_fp8(bf16, keep_rowmajor=w8a8) + if w8a8: + # Prefill reads the row-major FP8 copy; the BF16 weight is released. + layer._glm53_bf16_weight = None + layer._glm53_fp8_rowmajor = packed.rowmajor + layer._glm53_fp8_scale_col = packed.scale_col + else: + layer._glm53_bf16_weight = bf16 + del bf16 + layer._glm53_fp8_scale = packed.weight_scale + layer._glm53_fp8_workspace = packed.workspace + layer._glm53_fp8_nk = (n, k) + # Decode reads the packed tensor; expose it as the public weight the + # L2 prefetcher discovers. + layer.weight = packed.weight + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + ) + + if _decode_eligible(x, bias, self._max_m): + n, k = layer._glm53_fp8_nk + return apply_fp8_marlin_linear( + x, + layer.weight, + layer._glm53_fp8_scale, + layer._glm53_fp8_workspace, + n, + k, + None, + ) + if self._prefill == "w8a8": + from vllm import _custom_ops as ops + + shape = x.shape + x2 = x.reshape(-1, shape[-1]) + xq, xs = ops.scaled_fp8_quant(x2, use_per_token_if_dynamic=True) + out = ops.cutlass_scaled_mm( + xq, + layer._glm53_fp8_rowmajor.t(), + scale_a=xs, + scale_b=layer._glm53_fp8_scale_col, + out_dtype=x.dtype, + bias=bias, + ) + return out.reshape(*shape[:-1], out.shape[-1]) + return self._inner._gemm_impl(layer, x, layer._glm53_bf16_weight, bias) + + +class GLM53Fp8DecodeLMHeadMethod(QuantizeMethodBase): + """Wrap the LM head's method: keep its BF16 weight and add an FP8 copy. + + It must be a ``QuantizeMethodBase`` so the loader calls + ``process_weights_after_loading``, and must not be an + ``UnquantizedEmbeddingMethod``, or ``LogitsProcessor`` projects the raw + BF16 weight itself and never calls ``apply()``. + """ + + supports_pre_processed_weights = False + + def __init__(self, inner: UnquantizedEmbeddingMethod) -> None: + super().__init__() + self._inner = inner + self._max_m = max_m() + self._packed: nn.Module | None = None + + def __getattr__(self, name: str): + if name.startswith("__") or name in ("_inner", "_max_m", "_packed"): + raise AttributeError(name) + return getattr(self._inner, name) + + def create_weights(self, *args, **kwargs): + return self._inner.create_weights(*args, **kwargs) + + def embedding(self, layer: nn.Module, *args, **kwargs) -> torch.Tensor: + return self._inner.embedding(layer, *args, **kwargs) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + self._inner.process_weights_after_loading(layer) + weight = layer.weight.data + if ( + tuple(weight.shape) != GLM53_TP3_LM_HEAD_SHAPE + or weight.dtype != torch.bfloat16 + ): + logger.warning( + "GLM-5.3 FP8 decode: LM head is %s %s, not converting", + tuple(weight.shape), + weight.dtype, + ) + return + self._packed = _pack_marlin_fp8(weight) + logger.info_once("GLM-5.3 FP8 decode: LM head packed to Marlin FP8") + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + ) + + packed = self._packed + if packed is not None and _decode_eligible(x, bias, self._max_m): + return apply_fp8_marlin_linear( + x, + packed.weight, + packed.weight_scale, + packed.workspace, + packed.output_size_per_partition, + packed.input_size_per_partition, + None, + ) + return self._inner.apply(layer, x, bias) + + +def enable_glm53_fp8_dense(module: nn.Module, *, draft: bool = False) -> int: + """Wrap the selected BF16 projections; conversion happens after loading. + + ``draft=True`` is the MTP draft module (its DSA projections share the + target's shapes); ``VLLM_GLM53_FP8_MTP=0`` leaves it BF16. Returns the + number of wrapped projections. + """ + if not enabled() or (draft and not _flag("VLLM_GLM53_FP8_MTP", "1")): + return 0 + if not _tp_supported("MTP draft" if draft else "decoder"): + return 0 + shapes = fp8_dense_shapes() + mode = prefill_mode() + counts: dict[str, int] = {} + for child in module.modules(): + if not isinstance(child, LinearBase) or not isinstance( + child.quant_method, UnquantizedLinearMethod + ): + continue + if isinstance(child.quant_method, GLM53Fp8DecodeLinearMethod): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + continue + label = shapes.get((int(weight.shape[0]), int(weight.shape[1]))) + if label is None: + continue + child.quant_method = GLM53Fp8DecodeLinearMethod(child.quant_method, label) + counts[label] = counts.get(label, 0) + 1 + wrapped = sum(counts.values()) + summary = ", ".join(f"{k} x{v}" for k, v in sorted(counts.items())) + logger.info_once( + "GLM-5.3 FP8 decode%s: %d BF16 projections use Marlin FP8 W8A16 for " + "batches of at most %d rows, prefill %s [%s]", + " (MTP draft)" if draft else "", + wrapped, + max_m(), + mode, + summary, + ) + return wrapped + + +def enable_glm53_fp8_lm_head(lm_head: nn.Module) -> bool: + """Give the target LM head an FP8 decode copy. Returns whether it did.""" + if not enabled() or not _flag("VLLM_GLM53_FP8_LM_HEAD", "1"): + return False + if not _tp_supported("LM head"): + return False + weight = getattr(lm_head, "weight", None) + if ( + getattr(lm_head, "runtime_lm_head_quantization", None) is not None + or type(getattr(lm_head, "quant_method", None)) + is not UnquantizedEmbeddingMethod + or weight is None + or tuple(weight.shape) != GLM53_TP3_LM_HEAD_SHAPE + or weight.dtype != torch.bfloat16 + ): + logger.warning( + "GLM-5.3 FP8 decode: LM head is not an unquantized BF16 %s shard; " + "keeping it as is.", + GLM53_TP3_LM_HEAD_SHAPE, + ) + return False + lm_head.quant_method = GLM53Fp8DecodeLMHeadMethod(lm_head.quant_method) + logger.info_once("GLM-5.3 FP8 decode: LM head uses Marlin FP8 W8A16 for decode") + return True diff --git a/vllm/models/glm5next/nvidia/glm53_low_latency_gemm.py b/vllm/models/glm5next/nvidia/glm53_low_latency_gemm.py new file mode 100644 index 000000000000..445a2c309c07 --- /dev/null +++ b/vllm/models/glm5next/nvidia/glm53_low_latency_gemm.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash TP3 decode GEMM selection for unquantized BF16 on SM120. + +Mirrors ``deepseek_v32/nvidia/glm52_low_latency_gemm.py``. cuBLAS serves the +small per-rank BF16 projections at decode batch sizes with a split-K pair +(GEMM plus ``splitKreduce``) or an SM80 WMMA fallback. For the measured +(N, K, M) combinations where the CuTe shape-dynamic skinny GEMM beat cuBLAS by +more than 3% on RTX PRO 6000 at TP3 (+1.3-1.7% decode), route them to it; +every other shape and batch size keeps cuBLAS. + +The plan is keyed on TP3 per-rank shapes, and the replicated projections +(``fused_qkv_a_proj``, indexer ``wk``/``wq_b``/``weights_proj``) have the same +shape at every TP size, so the swap only happens at tensor_parallel_size == 3. +``VLLM_GLM53_LOW_LATENCY_GEMM=0`` disables it there; it is never enabled at +other TP sizes. +""" + +from __future__ import annotations + +import os + +import torch +from torch import nn + +import vllm.envs as envs +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +SUPPORTED_TP_SIZE = 3 + +# (N, K) per TP3 rank -> decode M values where the skinny GEMM wins. +GLM53_TP3_PLAN: dict[tuple[int, int], frozenset[int]] = { + (8598, 4096): frozenset({1, 2, 4, 8}), # KDA in_proj_qkvgfab + (128, 4096): frozenset({1, 2, 4}), # KDA g_a_proj, indexer wk + (4096, 2816): frozenset({1, 2}), # KDA o_proj + (2048, 4096): frozenset({1, 2}), # DSA fused_qkv_a_proj + (6144, 1536): frozenset({1}), # DSA q_b_proj + (4096, 6144): frozenset({1, 2, 4}), # DSA o_proj + (4096, 1536): frozenset({1, 2, 4}), # indexer wq_b + (32, 4096): frozenset({1, 2, 4, 8}), # indexer weights_proj +} + + +def _config_fits(config: SkinnyGemmConfig, m: int, n: int, k: int) -> bool: + """Whether the skinny GEMM accepts this shape (mirrors its call checks).""" + return ( + 1 <= m <= 16 + and config.num_rows == m + and n % config.outputs_per_block == 0 + and k % (config.block_size * config.vector_width) == 0 + and (config.static_k is None or config.static_k == k) + ) + + +def plan_for(n: int, k: int) -> dict[int, SkinnyGemmConfig]: + """Measured M -> config for a TP3 (N, K) shape, dropping unusable Ms.""" + plan: dict[int, SkinnyGemmConfig] = {} + for m in sorted(GLM53_TP3_PLAN.get((n, k), ())): + config = shape_dynamic_skinny_gemm._config(m, n, k) + if _config_fits(config, m, n, k): + plan[m] = config + return plan + + +def requested() -> bool: + """``VLLM_GLM53_LOW_LATENCY_GEMM``: unset or nonzero means on (at TP3).""" + return os.getenv("VLLM_GLM53_LOW_LATENCY_GEMM", "1").strip() not in ("", "0") + + +def _enabled(dtype: torch.dtype | None) -> bool: + # A pure-performance path: any unexpected configuration keeps cuBLAS + # rather than aborting model construction. + if not requested() or dtype != torch.bfloat16: + return False + if get_tensor_model_parallel_world_size() != SUPPORTED_TP_SIZE: + return False + if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(120) + ): + return False + return shape_dynamic_skinny_gemm.is_available() + + +def _row_major(t: torch.Tensor) -> bool: + return t.dim() == 2 and t.stride() == (t.shape[1], 1) + + +class GLM53LowLatencyLinearMethod(UnquantizedLinearMethod): + def __init__(self, plan: dict[int, SkinnyGemmConfig]) -> None: + super().__init__() + self._plan = plan + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + weight = layer.weight + config = self._plan.get(x.shape[0]) if x.dim() == 2 else None + if ( + config is not None + and bias is None + and not envs.VLLM_BATCH_INVARIANT + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and _row_major(x) + and _row_major(weight) + and x.shape[1] == weight.shape[1] + ): + return shape_dynamic_skinny_gemm(x, weight, config) + return super().apply(layer, x, bias) + + +def enable_glm53_low_latency_gemm(module: nn.Module, dtype: torch.dtype | None) -> int: + """Swap measured TP3 projections to the skinny GEMM; returns the count.""" + if not _enabled(dtype): + return 0 + warmup: set[SkinnyGemmConfig] = set() + swapped = 0 + for child in module.modules(): + if ( + not isinstance(child, LinearBase) + or type(child.quant_method) is not UnquantizedLinearMethod + ): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + continue + plan = plan_for(int(weight.shape[0]), int(weight.shape[1])) + if not plan: + continue + child.quant_method = GLM53LowLatencyLinearMethod(plan) + warmup.update(plan.values()) + swapped += 1 + if warmup: + # Compile before CUDA graph capture. + shape_dynamic_skinny_gemm.request_warmup_configs(torch.bfloat16, warmup) + logger.info_once( + "GLM-5.3 low-latency GEMM: %d BF16 projections routed to the CuTe skinny " + "GEMM for their measured decode batch sizes (%d warmup configs)", + swapped, + len(warmup), + ) + return swapped diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index e7acf6b728b4..2e3472b8bf07 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -97,6 +97,8 @@ from . import l2_prefetch as _l2pf from .attention import Glm5NextMLAAttention +from .glm53_fp8_dense import enable_glm53_fp8_dense, enable_glm53_fp8_lm_head +from .glm53_low_latency_gemm import enable_glm53_low_latency_gemm from .kda import Glm5NextLinearAttention from .pooled_indexer import Glm5NextIndexerScratch, Glm5NextPooledIndexer @@ -1341,6 +1343,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = Glm5NextModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + # TP3-only decode GEMM selection; both are no-ops at other TP sizes. + enable_glm53_low_latency_gemm(self.model, self.model_config.dtype) + enable_glm53_fp8_dense(self.model) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( self.config.vocab_size, @@ -1348,6 +1353,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head"), ) + enable_glm53_fp8_lm_head(self.lm_head) else: self.lm_head = PPMissingLayer() logit_scale = getattr(self.config, "logit_scale", 1.0) diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index ab119cb0fb34..b45d3ca2de44 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -26,6 +26,8 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from .glm53_fp8_dense import enable_glm53_fp8_dense +from .glm53_low_latency_gemm import enable_glm53_low_latency_gemm from .model import ( GLM5NEXT_PACKED_MODULES_MAPPING, Glm5NextDecoderLayer, @@ -279,6 +281,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = Glm5NextMultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + # TP3-only decode GEMM selection; both are no-ops at other TP sizes. + enable_glm53_low_latency_gemm(self.model, vllm_config.model_config.dtype) + enable_glm53_fp8_dense(self.model, draft=True) head = self.model._mtp_layers[0].shared_head.head self.has_own_lm_head = head.runtime_lm_head_quantization == "nvfp4" self.checkpoint_weight_name_prefixes = self._checkpoint_weight_name_prefixes() From 8cef00eeaea45e63447e19f1b2ea2d090b6a473b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 12:23:12 +0000 Subject: [PATCH 237/365] Save every rank's FlashInfer autotune entries and load them consistently FlashInfer MoE autotune keys include the expert-parallel rank, but only the leader saved its cache. On a warm start with EP > 1, every rank loaded the leader's file: rank 0 hit its keys and skipped profiling, while its peers missed theirs and entered synchronized (all-reduced) profiling, so startup hung with "No available shared memory broadcast block found". For multi-rank worlds: - Save the union of every rank's entries to a new versioned file, autotune_configs.allranks-v1.json, tagged with the world size. Old leader-only files keep the base name and are never read. Shared keys that disagree keep the lowest rank's entry and are logged. - Load collectively: the leader validates the union marker and broadcasts the cache; each rank loads it from a rank-local file, and unless every rank accepted it, all ranks clear their tuner and retune. A cache hit skips a profile's all-reduce, so all ranks must make the same hit/miss decisions. Single-rank behavior is unchanged. Ported from the community GLM-5.3 TP3 fix, generalized to any multi-rank world. Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 (1M context) --- .../test_flashinfer_autotune_rank_union.py | 233 ++++++++++++++++++ .../warmup/flashinfer_autotune_cache.py | 163 +++++++++++- vllm/model_executor/warmup/kernel_warmup.py | 27 +- 3 files changed, 409 insertions(+), 14 deletions(-) create mode 100644 tests/model_executor/test_flashinfer_autotune_rank_union.py diff --git a/tests/model_executor/test_flashinfer_autotune_rank_union.py b/tests/model_executor/test_flashinfer_autotune_rank_union.py new file mode 100644 index 000000000000..a06f57570267 --- /dev/null +++ b/tests/model_executor/test_flashinfer_autotune_rank_union.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multi-rank FlashInfer autotune cache: rank union and consistent loading. + +FlashInfer MoE cache keys include the EP rank. A leader-only cache let rank 0 +hit its keys and skip synchronized profiling while its peers missed and +entered the all-reduce, hanging the warm start. These tests run the +collective helpers on threads that stand in for ranks. +""" + +import json +import threading +from pathlib import Path +from typing import Any + +import pytest + +from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + RANK_UNION_CACHE_NAME, + RANK_UNION_MARKER_KEY, + load_autotune_cache_on_all_ranks, + merge_rank_autotune_configs, + rank_union_cache_path, + read_rank_union_cache, + save_rank_union_autotune_cache, +) + +pytestmark = pytest.mark.cpu_test + +WORLD = 3 +META = {"gpu": "RTX PRO 6000", "flashinfer": "0.6.18"} + + +class _Exchange: + def __init__(self, world_size: int) -> None: + self.world_size = world_size + self.barrier = threading.Barrier(world_size, timeout=10) + self.slot: Any = None + + +class _FakeGroup: + """GroupCoordinator stand-in: broadcast_object and barrier over threads.""" + + def __init__(self, exchange: _Exchange, rank: int) -> None: + self._exchange = exchange + self.rank_in_group = rank + self.world_size = exchange.world_size + + def broadcast_object(self, obj: Any = None, src: int = 0) -> Any: + if self.rank_in_group == src: + self._exchange.slot = obj + self._exchange.barrier.wait() + value = self._exchange.slot + self._exchange.barrier.wait() + return value + + def barrier(self) -> None: + self._exchange.barrier.wait() + + +class _FakeTuner: + """AutoTuner stand-in with FlashInfer's load/save/clear contract.""" + + def __init__(self, metadata: dict | None = None) -> None: + self.metadata = metadata or META + self.file_configs: dict[str, Any] = {} + self.profiled: dict[str, Any] = {} + + def load_configs(self, path: str) -> bool: + configs = json.loads(Path(path).read_text()) + assert RANK_UNION_MARKER_KEY not in configs + saved_meta = configs.pop("_metadata", None) + configs.pop("_generation", None) + if saved_meta is not None and saved_meta != self.metadata: + return False + self.file_configs.update(configs) + return True + + def save_configs(self, path: str) -> None: + configs = {"_metadata": self.metadata, **self.file_configs, **self.profiled} + Path(path).write_text(json.dumps(configs)) + + def clear_cache(self) -> None: + self.file_configs.clear() + self.profiled.clear() + + def tune(self, key: str, value: Any) -> bool: + """Return whether ``key`` needs synchronized profiling (a cache miss).""" + if key in self.file_configs or key in self.profiled: + return False + self.profiled[key] = value + return True + + +def _run_ranks(fn, world_size: int = WORLD) -> list[Any]: + exchange = _Exchange(world_size) + results: list[Any] = [None] * world_size + errors: list[BaseException] = [] + + def target(rank: int) -> None: + try: + results[rank] = fn(_FakeGroup(exchange, rank), rank) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + exchange.barrier.abort() + + threads = [threading.Thread(target=target, args=(r,)) for r in range(world_size)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + if errors: + raise errors[0] + return results + + +def _rank_keys(rank: int) -> list[str]: + # One shared GEMM key and one EP-rank-specific MoE key per rank. + return ["gemm_m128", f"moe_ep{rank}_m128"] + + +def _autotune_pass(cache_path: Path, tuners: list[_FakeTuner]) -> list[list[str]]: + """One start: load, tune every rank's keys, save. Returns misses per rank.""" + + def run(group: _FakeGroup, rank: int) -> list[str]: + tuner = tuners[rank] + load_autotune_cache_on_all_ranks(cache_path, tuner, group) + misses = [key for key in _rank_keys(rank) if tuner.tune(key, ["Runner", 1])] + save_rank_union_autotune_cache(cache_path, tuner, group) + return misses + + return _run_ranks(run) + + +def test_merge_unions_rank_specific_keys() -> None: + configs = [ + {"_metadata": META, "_generation": "a", "gemm": ["R", 1], "moe_ep0": ["M", 2]}, + {"_metadata": {"other": 1}, "gemm": ["R", 1], "moe_ep1": ["M", 3]}, + {"gemm": ["R", 1], "moe_ep2": ["M", 4], "_records": {"ns": {"x": 1}}}, + ] + contents, conflicts = merge_rank_autotune_configs(configs, WORLD) + merged = json.loads(contents) + assert conflicts == [] + assert merged["_metadata"] == META + assert {k for k in merged if not k.startswith("_")} == { + "gemm", + "moe_ep0", + "moe_ep1", + "moe_ep2", + } + assert merged["_records"] == {"ns": {"x": 1}} + assert merged[RANK_UNION_MARKER_KEY] == {"version": 1, "world_size": WORLD} + assert merged["_generation"] != "a" + # Deterministic output for identical inputs. + assert merge_rank_autotune_configs(configs, WORLD)[0] == contents + + +def test_merge_keeps_lowest_rank_on_disagreement() -> None: + contents, conflicts = merge_rank_autotune_configs( + [ + {"gemm": ["R", 1], "_records": {"ns": {"x": 1}}}, + {"gemm": ["R", 2], "_records": {"ns": {"x": 2}}}, + ], + 2, + ) + merged = json.loads(contents) + assert merged["gemm"] == ["R", 1] + assert merged["_records"] == {"ns": {"x": 1}} + assert conflicts == ["_records.ns", "gemm"] + + +def test_union_cache_validation() -> None: + contents, _ = merge_rank_autotune_configs([{"gemm": ["R", 1]}], WORLD) + stripped = json.loads(read_rank_union_cache(contents, WORLD)) + assert RANK_UNION_MARKER_KEY not in stripped + assert stripped["gemm"] == ["R", 1] + assert read_rank_union_cache(contents, 2) is None + # Leader-only files (no marker) and garbage are rejected. + assert read_rank_union_cache(json.dumps({"gemm": ["R", 1]}).encode(), 3) is None + assert read_rank_union_cache(b"not json", 3) is None + assert rank_union_cache_path(Path("/c/autotune_configs.json")) == Path( + "/c", RANK_UNION_CACHE_NAME + ) + + +def test_warm_start_hits_on_every_rank(tmp_path) -> None: + cache_path = tmp_path / RANK_UNION_CACHE_NAME + cold = _autotune_pass(cache_path, [_FakeTuner() for _ in range(WORLD)]) + assert cold == [_rank_keys(rank) for rank in range(WORLD)] + saved = json.loads(cache_path.read_text()) + for rank in range(WORLD): + assert f"moe_ep{rank}_m128" in saved + + # Fresh processes: every rank finds its own EP key, so none profiles. + warm = _autotune_pass(cache_path, [_FakeTuner() for _ in range(WORLD)]) + assert warm == [[] for _ in range(WORLD)] + + +def test_leader_only_cache_is_ignored_by_all_ranks(tmp_path) -> None: + cache_path = tmp_path / RANK_UNION_CACHE_NAME + # The pre-fix format: rank 0's keys only, no union marker. + cache_path.write_text( + json.dumps({"_metadata": META, "gemm_m128": ["R", 1], "moe_ep0_m128": ["M", 1]}) + ) + misses = _autotune_pass(cache_path, [_FakeTuner() for _ in range(WORLD)]) + # Every rank retunes everything, so the synchronized profiles line up. + assert misses == [_rank_keys(rank) for rank in range(WORLD)] + + +def test_cache_rejected_on_one_rank_is_dropped_everywhere(tmp_path) -> None: + cache_path = tmp_path / RANK_UNION_CACHE_NAME + _autotune_pass(cache_path, [_FakeTuner() for _ in range(WORLD)]) + + tuners = [_FakeTuner() for _ in range(WORLD)] + tuners[2].metadata = {"gpu": "different"} + + def run(group: _FakeGroup, rank: int) -> bool: + return load_autotune_cache_on_all_ranks(cache_path, tuners[rank], group) + + assert _run_ranks(run) == [False] * WORLD + assert all(not tuner.file_configs for tuner in tuners) + + +def test_single_rank_loads_plain_cache(tmp_path) -> None: + cache_path = tmp_path / "autotune_configs.json" + cache_path.write_text(json.dumps({"_metadata": META, "gemm": ["R", 1]})) + tuner = _FakeTuner() + + def run(group: _FakeGroup, rank: int) -> bool: + return load_autotune_cache_on_all_ranks(cache_path, tuner, group) + + assert _run_ranks(run, world_size=1) == [True] + assert tuner.file_configs == {"gemm": ["R", 1]} diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index afc3ef3ac014..1f630c113620 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -3,18 +3,22 @@ """FlashInfer autotune cache helpers.""" import hashlib +import json import os import tempfile from contextlib import suppress from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import vllm.envs as envs +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +logger = init_logger(__name__) + def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: config_hash = runner.vllm_config.compute_hash(include_version=False) @@ -92,3 +96,160 @@ def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None: with suppress(OSError): os.unlink(tmp_path) raise + + +# Multi-rank caches hold the union of every rank's entries. FlashInfer MoE +# cache keys include the expert-parallel rank, so a leader-only file misses +# the other ranks' keys: on a warm start the leader hits its cache and skips +# profiling while its peers miss and enter synchronized (all-reduced) +# profiling, which deadlocks. Older leader-only files use the base name and +# are never read for a multi-rank world. +RANK_UNION_CACHE_NAME = "autotune_configs.allranks-v1.json" +RANK_UNION_MARKER_KEY = "_vllm_rank_union" +_METADATA_KEY = "_metadata" +_GENERATION_KEY = "_generation" + + +def rank_union_cache_path(cache_path: Path) -> Path: + return cache_path.with_name(RANK_UNION_CACHE_NAME) + + +def gather_objects(group: "GroupCoordinator", obj: Any) -> list[Any]: + """Every rank's ``obj``, in rank order, on every rank.""" + return [ + group.broadcast_object(obj if group.rank_in_group == rank else None, src=rank) + for rank in range(group.world_size) + ] + + +def merge_rank_autotune_configs( + configs: list[dict[str, Any]], world_size: int +) -> tuple[bytes, list[str]]: + """Union per-rank FlashInfer autotune configs into one cache file. + + Returns the file contents and the keys on which ranks disagreed. Tuning is + synchronized across ranks, so shared keys should agree; on a disagreement + the lowest rank's entry is kept. Reserved ``_``-prefixed dict sections + (FlashInfer's namespaced records) are merged per entry the same way; the + lowest rank's ``_metadata`` is kept and ``_generation`` is recomputed. + """ + merged: dict[str, Any] = {} + conflicts: list[str] = [] + for config in configs: + for key, value in config.items(): + if key in (_GENERATION_KEY, RANK_UNION_MARKER_KEY): + continue + if key == _METADATA_KEY: + merged.setdefault(key, value) + continue + if key.startswith("_") and isinstance(value, dict): + section = merged.setdefault(key, {}) + for name, entry in value.items(): + if name in section and section[name] != entry: + conflicts.append(f"{key}.{name}") + continue + section[name] = entry + continue + if key in merged and merged[key] != value: + conflicts.append(key) + continue + merged[key] = value + ordered: dict[str, Any] = {} + if _METADATA_KEY in merged: + ordered[_METADATA_KEY] = merged.pop(_METADATA_KEY) + for key in sorted(merged): + ordered[key] = merged[key] + ordered[_GENERATION_KEY] = hashlib.sha256( + json.dumps(ordered, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + ordered[RANK_UNION_MARKER_KEY] = {"version": 1, "world_size": world_size} + return json.dumps(ordered, indent=2).encode(), sorted(set(conflicts)) + + +def read_rank_union_cache(contents: bytes, world_size: int) -> bytes | None: + """Validate a rank-union cache and strip the vLLM marker for FlashInfer. + + Returns None when the file is not a union of ``world_size`` ranks (for + example a leader-only file), so it must not be loaded by any rank. + """ + try: + config = json.loads(contents) + except ValueError: + return None + if not isinstance(config, dict): + return None + marker = config.pop(RANK_UNION_MARKER_KEY, None) + if not isinstance(marker, dict) or marker.get("world_size") != world_size: + return None + return json.dumps(config, indent=2).encode() + + +def save_rank_union_autotune_cache( + cache_path: Path, tuner: Any, group: "GroupCoordinator" +) -> None: + """Persist the union of every rank's autotune entries (collective).""" + with tempfile.TemporaryDirectory() as directory: + local_path = Path(directory) / "local.json" + tuner.save_configs(str(local_path)) + local = json.loads(local_path.read_text()) if local_path.exists() else {} + configs = gather_objects(group, local) + if group.rank_in_group == 0: + contents, conflicts = merge_rank_autotune_configs(configs, group.world_size) + if conflicts: + logger.warning( + "FlashInfer autotune ranks disagree on %d cache entries; keeping " + "the lowest rank's (first: %s).", + len(conflicts), + conflicts[0], + ) + write_flashinfer_autotune_cache(cache_path, contents) + group.barrier() + + +def load_autotune_cache_on_all_ranks( + cache_path: Path, tuner: Any, group: "GroupCoordinator" +) -> bool: + """Load the leader's cache on every rank so all make the same decisions. + + Collective. Within the synchronized tuning pass, a cache hit skips a + profile's all-reduce, so the ranks must hold identical entries: either + every rank loads the same validated cache or every rank starts empty. + A multi-rank world only accepts a rank-union file for its world size. + """ + world_size = group.world_size + contents: bytes | None = None + if group.rank_in_group == 0 and cache_path.exists(): + contents = cache_path.read_bytes() + if world_size > 1: + validated = read_rank_union_cache(contents, world_size) + if validated is None: + logger.warning( + "Ignoring FlashInfer autotune cache %s: not a %d-rank union.", + cache_path, + world_size, + ) + contents = validated + contents = group.broadcast_object(contents, src=0) + if contents is None: + return False + + try: + with tempfile.TemporaryDirectory() as directory: + local_path = Path(directory) / "autotune_configs.json" + local_path.write_bytes(contents) + loaded = bool(tuner.load_configs(str(local_path))) + except Exception as exc: + logger.warning("Failed to load FlashInfer autotune cache: %s", exc) + loaded = False + if world_size > 1 and not all(gather_objects(group, loaded)): + # Some rank rejected the cache (e.g. an environment mismatch): discard + # it everywhere so no rank skips a profile its peers run. + tuner.clear_cache() + if group.rank_in_group == 0: + logger.warning( + "FlashInfer autotune cache %s was not accepted by every rank; " + "retuning on all ranks.", + cache_path, + ) + return False + return loaded diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c765e12b7aea..c9f9de64c715 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -17,8 +17,10 @@ from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + load_autotune_cache_on_all_ranks, + rank_union_cache_path, resolve_flashinfer_autotune_file, - write_flashinfer_autotune_cache, + save_rank_union_autotune_cache, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( autotune_hisparse_flashinfer_attention, @@ -395,6 +397,11 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: autotune_kwargs["skip_ops"] = skip_ops cache_path = resolve_flashinfer_autotune_file(runner) + multi_rank = world.world_size > 1 + if multi_rank: + # Keys can be rank-specific (FlashInfer MoE keys include the EP rank), + # so multi-rank caches hold every rank's entries under their own name. + cache_path = rank_union_cache_path(cache_path) if is_leader: logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) @@ -403,16 +410,9 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: # which lead to some EP ranks receiving no tokens and skipping their # MoE kernel entirely, and cause hang due to all-reduce collective # during synchronized autotuning. - # Read cached autotune results and broadcast to all ranks. - cached_results: bytes | None = None - if is_leader and cache_path.exists(): - with open(cache_path, "rb") as f: - cached_results = f.read() - cached_results = world.broadcast_object(cached_results, src=0) - if cached_results is not None: - write_flashinfer_autotune_cache(cache_path, cached_results) - world.barrier() - tuner.load_configs(str(cache_path)) + # Load the leader's cache on every rank, or on none: a cache hit skips a + # profile's all-reduce, so all ranks must make the same hit/miss decisions. + load_autotune_cache_on_all_ranks(cache_path, tuner, world) group = world.cpu_group if world.world_size > 1 else None set_autotune_process_group(group) @@ -434,7 +434,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: finally: set_autotune_process_group(None) - if world.world_size > 1: + if multi_rank: world.barrier() - if is_leader: + save_rank_union_autotune_cache(cache_path, tuner, world) + else: tuner.save_configs(str(cache_path)) From 54be92b9d70f86e86d0e166867a38166e93fdfb9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 12:44:55 +0000 Subject: [PATCH 238/365] chore(release): describe #889 Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-889.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-889.json diff --git a/.lil/changes/vllm-889.json b/.lil/changes/vllm-889.json new file mode 100644 index 000000000000..43508baec843 --- /dev/null +++ b/.lil/changes/vllm-889.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-889", + "category": "feature", + "summary": "GLM-5.3-Flash on three GPUs (TP3)", + "models": ["GLM-5.3-Flash"], + "compatibility": "Opt-in: TP=3 with expert parallelism; other TP sizes are unchanged. The vision tower runs data-parallel at TP3.", + "details": [ + "At TP3, MLA heads are padded 64->72, KDA heads 64->66 and the shared expert 2048->2112 with zero weights, so outputs match the unpadded model; routed experts use expert parallelism.", + "Optional FP8 weight-only decode for the large BF16 projections (VLLM_GLM53_FP8_DENSE=1) and a low-latency GEMM plan for TP3 decode shapes.", + "FlashInfer autotune caches now hold every rank's entries and all ranks load them consistently; with expert parallelism a warm start could hang.", + "Based on BoxTops4Coochie's GLM-5.3-Flash TP3 port." + ], + "pull_requests": [889], + "requires": [887] +} From 9db80626a7922d717fd4794cd4d5b38e4f571438 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 13:18:41 +0000 Subject: [PATCH 239/365] Revert #885: dummy draft steps marked every row as padding With FlashInfer autotune on, warmup dummy runs sent -1 idx_mapping rows into the GLM-5.3 pooled indexer (restore_speculative_interval_starts), tripping index_select's bounds assert: GLM TP3 failed at startup, and DS4 profiles enable autotune by default. The fix only mattered for idle data-parallel ranks, which we do not deploy. Reverted merge 33db77e879d; re-land once consumers accept padded rows. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-885.json | 16 -------- .lil/changes/vllm-890.json | 14 +++++++ tests/v1/worker/test_gpu_block_table.py | 39 ------------------- vllm/v1/worker/gpu/block_table.py | 11 ++---- .../spec_decode/autoregressive/speculator.py | 1 - .../multi_module_mtp/speculator.py | 1 - vllm/v1/worker/gpu/spec_decode/speculator.py | 14 ++----- 7 files changed, 21 insertions(+), 75 deletions(-) delete mode 100644 .lil/changes/vllm-885.json create mode 100644 .lil/changes/vllm-890.json diff --git a/.lil/changes/vllm-885.json b/.lil/changes/vllm-885.json deleted file mode 100644 index 3cc3c1ee4ad9..000000000000 --- a/.lil/changes/vllm-885.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "schema": "local-inference-release-change/v1", - "id": "vllm-885", - "category": "fix", - "summary": "Stop idle data-parallel ranks' dummy draft steps from writing KV through stale block tables", - "models": ["all"], - "compatibility": "Affects speculative decoding with data parallelism (DP > 1) on the model-runner v2 path; no user action required.", - "details": [ - "A dummy draft decode step on an idle DP rank could write KV through block-table rows left by finished requests, corrupting cache pages another request reuses.", - "Padded rows now map to the pad slot. Backport of vllm-project/vllm#56734 with the original author." - ], - "pull_requests": [885], - "authors": ["Yifan Qiao"], - "evidence": ["https://github.com/vllm-project/vllm/pull/56734"], - "requires": [] -} diff --git a/.lil/changes/vllm-890.json b/.lil/changes/vllm-890.json new file mode 100644 index 000000000000..1818dda3a504 --- /dev/null +++ b/.lil/changes/vllm-890.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-890", + "category": "fix", + "summary": "Revert the dummy-draft block-table change (#885): startup crashed with FlashInfer autotune", + "models": ["all"], + "compatibility": "No user action required. Beta 3e670ca5 could fail at startup with a CUDA device-side assert when FlashInfer autotune is on (DS4 profiles by default, GLM TP3).", + "details": [ + "#885 marked every row of a dummy draft batch as padding. Warmup dummy runs then passed -1 indices to the GLM-5.3 pooled indexer, which index_select rejects.", + "The data-parallel fix it carried will return once the consumers handle padded rows." + ], + "pull_requests": [890], + "requires": [] +} diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 1f9628a8aa90..1aa1da0eafef 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,7 +5,6 @@ import torch from vllm.platforms import current_platform -from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.block_table import BlockTables pytestmark = pytest.mark.skipif( @@ -379,41 +378,3 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() - - -def test_dummy_request_slot_mapping_is_pad(): - """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request. - - A dummy draft decode step must not resolve slots through the persistent - block-table row of a real request slot, which may be stale and point at - blocks that are now in the prefix cache. - """ - device = torch.device("cuda") - block_tables = BlockTables( - block_sizes=[4], - max_num_reqs=2, - max_num_batched_tokens=16, - max_num_blocks_per_group=[4], - device=device, - kernel_block_sizes=[4], - ) - block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) - block_tables.apply_staged_writes() - query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) - positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) - - real = block_tables.compute_slot_mappings( - torch.tensor([0], dtype=torch.int32, device=device), - query_start_loc, - positions, - num_tokens_padded=3, - ) - assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3] - - dummy = block_tables.compute_slot_mappings( - torch.tensor([-1], dtype=torch.int32, device=device), - query_start_loc, - positions, - num_tokens_padded=3, - ) - assert dummy[0].tolist() == [PAD_SLOT_ID] * 3 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index cb9f9eda3669..2c40aa28a459 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -342,12 +342,7 @@ def _compute_slot_mappings_kernel( group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) - # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns - # no blocks: never read its block-table row and emit PAD for its tokens. - is_real_req = req_state_idx >= 0 - num_blocks = tl.load( - group_num_blocks_ptr + req_state_idx, mask=is_real_req, other=0 - ) + num_blocks = tl.load(group_num_blocks_ptr + req_state_idx) start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): @@ -377,7 +372,7 @@ def _compute_slot_mappings_kernel( valid_block = token_mask & (block_indices < num_blocks) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local & valid_block & is_real_req, + mask=is_local & valid_block, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets @@ -385,5 +380,5 @@ def _compute_slot_mappings_kernel( slot_ids = tl.where(is_local, slot_ids, PAD_ID) slot_ids = tl.where(valid_block, slot_ids, PAD_ID) - slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 8acb24869c46..d7202c83cab8 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -299,7 +299,6 @@ def propose( input_batch.idx_mapping, temperature, seeds, - dummy_run=dummy_run, ) # Get the input ids and last token indices for the speculator. diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index f00af0321e17..949e71683cd8 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -168,7 +168,6 @@ def propose( input_batch.idx_mapping, temperature, seeds, - dummy_run=dummy_run, ) num_tokens = input_batch.num_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index ba1ba4fca3bd..cc0cf95c4208 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -522,7 +522,6 @@ def _copy_request_inputs( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dummy_run: bool = False, ) -> None: # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. # NOTE(woosuk): For draft sampling, we only consider the temperature @@ -532,15 +531,10 @@ def _copy_request_inputs( # affect the output distribution after rejection sampling. self.temperature.copy_(temperature) self.seeds.copy_(seeds) - # idx_mapping == -1 marks a row the drafter must not act on: sampling - # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows - # always get it; a dummy batch gets it for every row, since its arange - # idx_mapping names request-state slots with stale block tables. - if dummy_run: - self.idx_mapping.fill_(-1) - else: - self.idx_mapping[:num_reqs].copy_(idx_mapping) - self.idx_mapping[num_reqs:].fill_(-1) + self.idx_mapping[:num_reqs].copy_(idx_mapping) + # idx_mapping for CG padded requests points to -1, which is ignored + # during sampling to prevent writing stale values to draft logits. + self.idx_mapping[num_reqs:].fill_(-1) def _build_uniform_batch_dp_sync( self, From 42e19d06777f918af080462a7bf0f730085ba414 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 14:24:38 +0000 Subject: [PATCH 240/365] chore(release): reference vllm-887 by fragment id in #889 The release compiler requires fragment ids as strings; the integer broke the beta publisher. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-889.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-889.json b/.lil/changes/vllm-889.json index 43508baec843..80916b931b90 100644 --- a/.lil/changes/vllm-889.json +++ b/.lil/changes/vllm-889.json @@ -12,5 +12,5 @@ "Based on BoxTops4Coochie's GLM-5.3-Flash TP3 port." ], "pull_requests": [889], - "requires": [887] + "requires": ["vllm-887"] } From 736ef874693a195ad7adc644d08768fb00748b58 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 14:35:27 +0000 Subject: [PATCH 241/365] chore(release): restore the published vllm-885 fragment Published change fragments are immutable; the #885 revert (#890) deleted it and the beta publisher refused the tree. #890's own fragment records the revert. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Martin Vit --- .lil/changes/vllm-885.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-885.json diff --git a/.lil/changes/vllm-885.json b/.lil/changes/vllm-885.json new file mode 100644 index 000000000000..3cc3c1ee4ad9 --- /dev/null +++ b/.lil/changes/vllm-885.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-885", + "category": "fix", + "summary": "Stop idle data-parallel ranks' dummy draft steps from writing KV through stale block tables", + "models": ["all"], + "compatibility": "Affects speculative decoding with data parallelism (DP > 1) on the model-runner v2 path; no user action required.", + "details": [ + "A dummy draft decode step on an idle DP rank could write KV through block-table rows left by finished requests, corrupting cache pages another request reuses.", + "Padded rows now map to the pad slot. Backport of vllm-project/vllm#56734 with the original author." + ], + "pull_requests": [885], + "authors": ["Yifan Qiao"], + "evidence": ["https://github.com/vllm-project/vllm/pull/56734"], + "requires": [] +} From 4b57a8a67fe60b18a5d65f9d8a25b856f184333a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 17:01:13 +0000 Subject: [PATCH 242/365] chore(release): describe the 2026-09-24 canonical sync Co-Authored-By: Claude Opus 5.5 Signed-off-by: Martin Vit --- .../vllm-karmic-canonical-sync-20260924.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-karmic-canonical-sync-20260924.json diff --git a/.lil/changes/vllm-karmic-canonical-sync-20260924.json b/.lil/changes/vllm-karmic-canonical-sync-20260924.json new file mode 100644 index 000000000000..c920d496fb1b --- /dev/null +++ b/.lil/changes/vllm-karmic-canonical-sync-20260924.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-karmic-canonical-sync-20260924", + "category": "compatibility", + "summary": "Carry the current canonical Karmic runtime into beta: DeepSeek V4.1 decode prefetch and compact IQ2/Q8_0 checkpoints", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "DeepSeek V4.1 decode prefetches the next weights into L2 and reads disk Engram rows for both Engram layers concurrently; the NVFP4 Markov head stays opt-in (VLLM_DS41_MARKOV_NVFP4=1).", + "ModelOpt IQ2_XS, IQ2_XXS and Q8_0 checkpoints load through compact B12X dense, expert, embedding and LM-head paths." + ], + "authors": ["lukealonso"], + "requires": ["b12x-master-sync-20260924"] +} From b1e0a401f4f5c3f1311f6462b3298753b701b390 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 24 Sep 2026 19:47:09 +0000 Subject: [PATCH 243/365] fix(worker): save startup plans under the pre-profiling fingerprint maybe_save_startup_plan recomputed the fingerprint after profiling and graph capture. DeepSeek-V4-Flash rewrites cudagraph_mode from FULL_AND_PIECEWISE to FULL_DECODE_ONLY in between, so each boot saved a plan under a key that the next boot, which fingerprints before profiling, never looks up. VLLM_ENABLE_STARTUP_PLAN therefore never skipped the CUDA-graph memory estimation pass for DS4. Remember the fingerprint computed at apply time and save under it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Martin Vit --- .lil/changes/vllm-895.json | 13 +++++++++++++ tests/v1/worker/test_gpu_worker.py | 17 +++++++++++++++++ vllm/v1/worker/gpu_worker.py | 2 ++ vllm/v1/worker/startup_plan.py | 12 +++++++++--- 4 files changed, 41 insertions(+), 3 deletions(-) create mode 100644 .lil/changes/vllm-895.json diff --git a/.lil/changes/vllm-895.json b/.lil/changes/vllm-895.json new file mode 100644 index 000000000000..8f5f5fc53a0e --- /dev/null +++ b/.lil/changes/vllm-895.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-895", + "category": "fix", + "summary": "Persisted startup plans are found again after profiling rewrites the configuration", + "models": ["DeepSeek-V4-Flash"], + "compatibility": "Takes effect only with VLLM_ENABLE_STARTUP_PLAN=1.", + "details": [ + "The startup plan was saved under a fingerprint computed after profiling. DeepSeek-V4-Flash rewrites its CUDA-graph mode from FULL_AND_PIECEWISE to FULL_DECODE_ONLY during profiling, so the next boot looked for a different fingerprint and never applied the plan.", + "The plan is now saved under the fingerprint computed before profiling. A matching plan skips the CUDA-graph memory estimation pass, which re-prepares every b12x attention plan." + ], + "pull_requests": [895] +} diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 8fbcce51b9d2..c17eec5e6cf1 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -180,6 +180,23 @@ def test_startup_plan_apply_gate(plan_env): assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes +def test_startup_plan_survives_profiling_config_rewrites(plan_env): + """A boot saves under the key computed before profiling, so the next + boot finds the plan although profiling rewrote the config (DS4 falls + back from FULL_AND_PIECEWISE to FULL_DECODE_ONLY graphs).""" + config_hash = ["before-profiling"] + first = _plan_worker() + first.vllm_config = SimpleNamespace(compute_hash=lambda: config_hash[0]) + maybe_apply_startup_plan(first) + assert first.cache_config.kv_cache_memory_bytes is None + config_hash[0] = "after-capture" + maybe_save_startup_plan(first, 50 * GiB_bytes) + + second = _plan_worker(config_hash="before-profiling") + maybe_apply_startup_plan(second) + assert second.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes + + # Memory accounting of the profiling run (Worker.determine_available_memory). diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index eb10ea6f28bd..3222d7ef2df6 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -258,6 +258,8 @@ def __init__( self._b12x_profile_batch: B12xPreparedBatch | None = None self._b12x_tuning_batch: B12xPreparedBatch | None = None self._b12x_tuning_cache = False + # Startup-plan key computed before profiling can rewrite the config. + self.startup_plan_fingerprint: str | None = None # Resolved lazily on first sleep/wake; persists worker-process state. self._sleep_mode_backend: SleepModeBackend | None = None diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py index 2c4942070705..fa3d8adaae39 100644 --- a/vllm/v1/worker/startup_plan.py +++ b/vllm/v1/worker/startup_plan.py @@ -143,6 +143,10 @@ def maybe_apply_startup_plan(worker: "Worker") -> None: fingerprint = compute_plan_fingerprint( worker.vllm_config, worker.rank, worker.parallel_config.world_size ) + # Profiling and capture can still rewrite the config, for example a + # CUDA-graph mode the loaded model cannot capture. Save under the key + # that the next boot computes here, before those rewrites. + worker.startup_plan_fingerprint = fingerprint plan = _load_plan(fingerprint) if plan is None: return @@ -170,9 +174,11 @@ def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> Non never raised.""" if not envs.VLLM_ENABLE_STARTUP_PLAN: return - fingerprint = compute_plan_fingerprint( - worker.vllm_config, worker.rank, worker.parallel_config.world_size - ) + fingerprint = getattr(worker, "startup_plan_fingerprint", None) + if fingerprint is None: + fingerprint = compute_plan_fingerprint( + worker.vllm_config, worker.rank, worker.parallel_config.world_size + ) path = _plan_path(fingerprint) try: os.makedirs(os.path.dirname(path), exist_ok=True) From b39ca56600623a46ddd3d614e2052375e88953f9 Mon Sep 17 00:00:00 2001 From: derek Date: Thu, 24 Sep 2026 16:15:49 -0400 Subject: [PATCH 244/365] [Bugfix] Stop workers when an EngineCore that failed to start exits Workers are non-daemon processes, so multiprocessing's exit handler joins them. When EngineCore construction fails after the executor starts its workers, nothing calls executor.shutdown(): the process joins workers that wait on its death pipe, and the front end waits in wait_for_engine_startup forever. A failed start (e.g. a wrong CP_KV_CACHE_INTERLEAVE_SIZE) left the container up with no error exit. Register executor shutdown as a multiprocessing exit finalizer, which runs before the join. shutdown() is idempotent. --- .../executor/test_multiproc_executor_exit.py | 53 +++++++++++++++++++ vllm/v1/executor/multiproc_executor.py | 15 +++++- 2 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 tests/v1/executor/test_multiproc_executor_exit.py diff --git a/tests/v1/executor/test_multiproc_executor_exit.py b/tests/v1/executor/test_multiproc_executor_exit.py new file mode 100644 index 000000000000..c18e8b2cb54d --- /dev/null +++ b/tests/v1/executor/test_multiproc_executor_exit.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A process whose engine failed to start must not wait forever for workers.""" + +import contextlib +import multiprocessing +from multiprocessing.connection import Connection + +from vllm.v1.executor.multiproc_executor import MultiprocExecutor, WorkerProcHandle + +CONTEXT = multiprocessing.get_context("fork") + + +def _worker(death_reader: Connection, death_writer: Connection) -> None: + """Stand-in worker: runs until its parent closes the death pipe.""" + death_writer.close() # as real workers close their inherited fds + with contextlib.suppress(EOFError): + death_reader.recv() + + +def _engine_that_fails_after_starting_workers() -> None: + death_reader, death_writer = CONTEXT.Pipe(duplex=False) + # Not a daemon, like the real workers. + worker = CONTEXT.Process( + target=_worker, args=(death_reader, death_writer), daemon=False + ) + worker.start() + death_reader.close() + executor = MultiprocExecutor.__new__(MultiprocExecutor) + executor.workers = [ + WorkerProcHandle( + proc=worker, + rank=0, + worker_response_mq=None, + peer_worker_response_mqs=[], + death_writer=death_writer, + ) + ] + executor._register_exit_shutdown() + # EngineCore construction fails; nothing calls executor.shutdown(). + raise RuntimeError("engine startup failed") + + +def test_failed_engine_start_exits_and_stops_its_workers(): + engine = CONTEXT.Process(target=_engine_that_fails_after_starting_workers) + engine.start() + engine.join(timeout=60) + hung = engine.is_alive() + if hung: + engine.kill() # closes the death pipe, so the worker exits too + engine.join() + assert not hung, "the engine process waited forever for its workers" + assert engine.exitcode == 1 diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index dcd9b6b79240..598121f9e6af 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -19,6 +19,7 @@ from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess from multiprocessing.synchronize import Lock as LockType +from multiprocessing.util import Finalize from threading import Thread from typing import Any, cast @@ -115,10 +116,22 @@ def __init__(self, vllm_config: VllmConfig, monitor_workers: bool = True): self.monitor_workers = monitor_workers super().__init__(vllm_config) - def _init_executor(self) -> None: + def _register_exit_shutdown(self) -> None: + """Shut the workers down when this process exits, however it exits.""" # Call self.shutdown at exit to clean up # and ensure workers will be terminated. self._finalizer = weakref.finalize(self, self.shutdown) + # Workers are not daemons (their startup compilers may fork), so the + # multiprocessing exit handler joins them. If this process exits + # without shutting the executor down, e.g. an EngineCore whose + # startup failed after the workers started, that join waits forever + # on workers that are waiting for this process. Finalizers with an + # exit priority run before the join; the weakref finalizer above only + # runs at interpreter exit, which a multiprocessing child skips. + Finalize(self, self.shutdown, exitpriority=0) + + def _init_executor(self) -> None: + self._register_exit_shutdown() self.is_failed = False self.failure_callback: FailureCallback | None = None From 27fda44f176e1cb7e620a6704ddc6cf047cae5fd Mon Sep 17 00:00:00 2001 From: derek Date: Fri, 25 Sep 2026 07:45:13 -0400 Subject: [PATCH 245/365] docs(release): identify PR 897 --- .lil/changes/vllm-897.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 .lil/changes/vllm-897.json diff --git a/.lil/changes/vllm-897.json b/.lil/changes/vllm-897.json new file mode 100644 index 000000000000..2c1e1348e36f --- /dev/null +++ b/.lil/changes/vllm-897.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-897", + "category": "fix", + "summary": "A server whose engine fails to start now exits instead of hanging", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "When engine startup fails after the GPU workers have started (for example an invalid CP_KV_CACHE_INTERLEAVE_SIZE), the workers are stopped and the container exits with an error within seconds, so restart policies take effect.", + "Previously the engine process waited for its workers forever and the container stayed up without serving." + ], + "pull_requests": [897] +} From e83cbce41028f4b5f7967b6dfc3784da61a5cdf7 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 13:31:29 +0000 Subject: [PATCH 246/365] feat(glm53): optional online MXFP8 for the vision tower VLLM_GLM53_VISION_MXFP8=1 quantizes the vision tower's linear layers (blocks and merger) to MXFP8 while loading their BF16 weights. Convolutions and norms stay BF16. Saves 0.24 GiB per GPU at TP2. Co-Authored-By: Claude Opus 5.5 --- vllm/envs.py | 6 ++++++ vllm/models/glm5next/nvidia/model.py | 18 ++++++++++++++++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 2fe1250ba56d..ff0a25298f51 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -50,6 +50,7 @@ VLLM_TRACE_FUNCTION: int = 0 VLLM_USE_FLASHINFER_SAMPLER: bool = True VLLM_GLM53_MTP_DRAFT_HEAD: Literal["bf16", "nvfp4"] = "bf16" + VLLM_GLM53_VISION_MXFP8: bool = False VLLM_PP_LAYER_PARTITION: str | None = None VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" @@ -915,6 +916,11 @@ def _resolve_rust_cli_path() -> str | None: ["bf16", "nvfp4"], case_sensitive=False, ), + # Quantize the GLM-5.3 vision tower's linear layers to MXFP8 while loading + # its BF16 weights. Convolutions and norms stay BF16. + "VLLM_GLM53_VISION_MXFP8": lambda: bool( + int(os.getenv("VLLM_GLM53_VISION_MXFP8", "0")) + ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), # (CPU backend only) CPU key-value cache space. diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 2e3472b8bf07..967c332efc10 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -7,6 +7,7 @@ import torch from torch import nn +import vllm.envs as envs from vllm.config import ParallelConfig, VllmConfig from vllm.distributed import ( get_ep_group, @@ -1541,8 +1542,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # the latter onto text_config (1e-5), silently ignoring the # vision tower's own (1e-6) rms_norm_eps. norm_eps=config.vision_config.rms_norm_eps, - # The vision tower ships BF16 weights and is not quantized. - quant_config=None, + # The vision tower ships BF16 weights; it is quantized only + # on request (VLLM_GLM53_VISION_MXFP8). + quant_config=_vision_quant_config(), prefix=maybe_prefix(prefix, "visual"), ) @@ -1681,6 +1683,18 @@ def _try_load_fp8_attn_proj( return True +def _vision_quant_config() -> QuantizationConfig | None: + """Online MXFP8 for the vision tower's linear layers, when requested.""" + if not envs.VLLM_GLM53_VISION_MXFP8: + return None + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + + return OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) + + def _try_load_mxfp8_bf16_attn_proj( name, tensor, From bd1bad8455df5bcf83733786f9ad0a449401d6c5 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 14:11:01 +0000 Subject: [PATCH 247/365] feat(glm53): optional input embedding table in pinned host RAM VLLM_GLM53_EMBED_HOST=1 keeps the vocabulary table of the target and the MTP draft in pinned host memory; the GPU reads the rows of the current tokens through a UVA view. Frees 0.59 GiB per GPU at TP2 (Spark checkpoint); decode steps/s unchanged, prefill about 1% slower. Co-Authored-By: Claude Opus 5.5 --- vllm/envs.py | 6 ++++++ vllm/models/glm5next/nvidia/model.py | 23 +++++++++++++++++++++++ vllm/models/glm5next/nvidia/mtp.py | 2 ++ 3 files changed, 31 insertions(+) diff --git a/vllm/envs.py b/vllm/envs.py index ff0a25298f51..de946938e076 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -51,6 +51,7 @@ VLLM_USE_FLASHINFER_SAMPLER: bool = True VLLM_GLM53_MTP_DRAFT_HEAD: Literal["bf16", "nvfp4"] = "bf16" VLLM_GLM53_VISION_MXFP8: bool = False + VLLM_GLM53_EMBED_HOST: bool = False VLLM_PP_LAYER_PARTITION: str | None = None VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" @@ -921,6 +922,11 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_GLM53_VISION_MXFP8": lambda: bool( int(os.getenv("VLLM_GLM53_VISION_MXFP8", "0")) ), + # Keep the GLM-5.3 input embedding table in pinned host RAM; the GPU reads + # the rows of the current tokens through a UVA view. + "VLLM_GLM53_EMBED_HOST": lambda: bool( + int(os.getenv("VLLM_GLM53_EMBED_HOST", "0")) + ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), # (CPU backend only) CPU key-value cache space. diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 967c332efc10..e4ec78697ea2 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -982,6 +982,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.hidden_size, prefix=f"{prefix}.embed_tokens", ) + host_embedding_if_requested(self.embed_tokens) else: self.embed_tokens = PPMissingLayer() @@ -1683,6 +1684,28 @@ def _try_load_fp8_attn_proj( return True +def host_embedding_if_requested(embed: VocabParallelEmbedding) -> None: + """Move the vocabulary table to pinned host RAM (VLLM_GLM53_EMBED_HOST). + + A lookup reads one row per token, so the GPU reads the rows it needs + through a UVA view instead of keeping the table in device memory. + Checkpoint loading writes through the same view. + """ + if not envs.VLLM_GLM53_EMBED_HOST: + return + from vllm.utils.platform_utils import is_uva_available + from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + + if not is_uva_available(): + raise RuntimeError("VLLM_GLM53_EMBED_HOST requires UVA and pinned memory") + weight = embed.weight + host = torch.empty( + weight.shape, dtype=weight.dtype, device="cpu", pin_memory=True + ) + weight.data = get_accelerator_view_from_cpu_tensor(host) + weight._vllm_is_uva_offloaded = True + + def _vision_quant_config() -> QuantizationConfig | None: """Online MXFP8 for the vision tower's linear layers, when requested.""" if not envs.VLLM_GLM53_VISION_MXFP8: diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index b45d3ca2de44..3e16a18f59a4 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -35,6 +35,7 @@ _try_load_fp8_attn_proj, _try_load_mxfp8_bf16_attn_proj, get_spec_layer_idx_from_weight_name, + host_embedding_if_requested, ) from .mtp_draft_head import QuantizedDraftHead, make_quantized_draft_head from .pooled_indexer import Glm5NextPooledIndexer @@ -137,6 +138,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.hidden_size, prefix=maybe_prefix(prefix, "embed_tokens"), ) + host_embedding_if_requested(self.embed_tokens) # Plain list for the per-propose lookup: ModuleDict[str(...)] builds a # string and hashes it on every draft step. self._mtp_layers = list(self.layers.values()) From bfeab715137832fb2196e17b6e05ec7a94f4955b Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 15:30:51 +0000 Subject: [PATCH 248/365] feat(distributed): share one PyNCCL communicator across groups over the same ranks VLLM_SHARE_PYNCCL_COMMS=1 lets groups whose ranks and device match (TP, DCP and EP at TP2+DCP2) reuse one PyNCCL communicator, released when its last owner is destroyed. Their device collectives are issued from the current stream, so the order on the shared communicator matches on every rank. Saves 74 MiB per GPU of NCCL channel and P2P buffers at TP2+DCP2; decode and prefill speed unchanged. Co-Authored-By: Claude Opus 5.5 --- .../device_communicators/cuda_communicator.py | 57 +++++++++++++++++-- vllm/envs.py | 5 ++ 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index a2e60d34025b..f0969847faaa 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -26,6 +26,52 @@ logger = init_logger(__name__) +# PyNCCL communicators shared by groups over the same ranks and device +# (VLLM_SHARE_PYNCCL_COMMS): key -> [communicator, owners]. +_SHARED_PYNCCL: dict[tuple, list] = {} + + +def _acquire_pynccl(group, device: torch.device): + """Return (communicator, share key, created). + + With VLLM_SHARE_PYNCCL_COMMS, groups over the same ranks (for example TP, + DCP and EP at TP2+DCP2) reuse one NCCL communicator instead of each + allocating channel and P2P buffers. Every rank creates its groups in the + same order, so each rank makes the same reuse decision. The groups' device + collectives are issued from the current stream, so their order on the + shared communicator matches on every rank. + """ + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + + if not envs.VLLM_SHARE_PYNCCL_COMMS or isinstance(group, StatelessProcessGroup): + return PyNcclCommunicator(group=group, device=device), None, True + import torch.distributed as dist + + key = (tuple(dist.get_process_group_ranks(group)), str(device)) + entry = _SHARED_PYNCCL.get(key) + if entry is not None: + entry[1] += 1 + logger.info("Reusing the PyNCCL communicator of ranks %s", key[0]) + return entry[0], key, False + comm = PyNcclCommunicator(group=group, device=device) + _SHARED_PYNCCL[key] = [comm, 1] + return comm, key, True + + +def _release_pynccl(comm, key: tuple | None) -> None: + if key is None: + comm.destroy() + return + entry = _SHARED_PYNCCL.get(key) + if entry is None or entry[0] is not comm: + comm.destroy() + return + entry[1] -= 1 + if entry[1] == 0: + del _SHARED_PYNCCL[key] + comm.destroy() + + class CudaCommunicator(DeviceCommunicatorBase): def __init__( self, @@ -109,12 +155,13 @@ def __init__( from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator self.pynccl_comm: PyNcclCommunicator | None = None + self._pynccl_share_key: tuple | None = None if self.world_size > 1: - self.pynccl_comm = PyNcclCommunicator( - group=self.cpu_group if tcp_store_group is None else tcp_store_group, - device=self.device, + group = self.cpu_group if tcp_store_group is None else tcp_store_group + self.pynccl_comm, self._pynccl_share_key, created = _acquire_pynccl( + group, self.device ) - if is_symmetric_memory_enabled(): + if created and is_symmetric_memory_enabled(): register_nccl_symmetric_ops(self.pynccl_comm) self.ca_comm: CustomAllreduce | None = None @@ -722,7 +769,7 @@ def destroy(self): self.b12x_ar_comm.close() self.b12x_ar_comm = None if self.pynccl_comm is not None: - self.pynccl_comm.destroy() + _release_pynccl(self.pynccl_comm, self._pynccl_share_key) self.pynccl_comm = None if self.ca_comm is not None: self.ca_comm = None diff --git a/vllm/envs.py b/vllm/envs.py index de946938e076..3f7c05734003 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -52,6 +52,7 @@ VLLM_GLM53_MTP_DRAFT_HEAD: Literal["bf16", "nvfp4"] = "bf16" VLLM_GLM53_VISION_MXFP8: bool = False VLLM_GLM53_EMBED_HOST: bool = False + VLLM_SHARE_PYNCCL_COMMS: bool = False VLLM_PP_LAYER_PARTITION: str | None = None VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" @@ -927,6 +928,10 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_GLM53_EMBED_HOST": lambda: bool( int(os.getenv("VLLM_GLM53_EMBED_HOST", "0")) ), + # Groups over the same ranks share one PyNCCL communicator. + "VLLM_SHARE_PYNCCL_COMMS": lambda: bool( + int(os.getenv("VLLM_SHARE_PYNCCL_COMMS", "0")) + ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), # (CPU backend only) CPU key-value cache space. From f4cd9fa63bbd395e55e542ba7d9020ffd80ead11 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 16:55:25 +0000 Subject: [PATCH 249/365] test(distributed): cover shared PyNCCL communicators; format Co-Authored-By: Claude Opus 5.5 --- tests/distributed/test_shared_pynccl.py | 59 +++++++++++++++++++++++++ vllm/envs.py | 4 +- vllm/models/glm5next/nvidia/model.py | 4 +- 3 files changed, 61 insertions(+), 6 deletions(-) create mode 100644 tests/distributed/test_shared_pynccl.py diff --git a/tests/distributed/test_shared_pynccl.py b/tests/distributed/test_shared_pynccl.py new file mode 100644 index 000000000000..e48c3e8d3f01 --- /dev/null +++ b/tests/distributed/test_shared_pynccl.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Groups over the same ranks share one PyNCCL communicator on request.""" + +import pytest +import torch + +from vllm.distributed.device_communicators import cuda_communicator as cc + + +class _FakeComm: + def __init__(self, group, device): + self.group = group + self.device = device + self.destroyed = 0 + + def destroy(self): + self.destroyed += 1 + + +@pytest.fixture +def fake_pynccl(monkeypatch): + import torch.distributed as dist + + from vllm.distributed.device_communicators import pynccl + + monkeypatch.setattr(pynccl, "PyNcclCommunicator", _FakeComm) + monkeypatch.setattr(dist, "get_process_group_ranks", lambda group: group) + monkeypatch.setattr(cc, "_SHARED_PYNCCL", {}) + + +def test_groups_over_the_same_ranks_share_one_communicator(fake_pynccl, monkeypatch): + monkeypatch.setenv("VLLM_SHARE_PYNCCL_COMMS", "1") + device = torch.device("cuda:0") + tp, tp_key, tp_created = cc._acquire_pynccl([0, 1], device) + dcp, dcp_key, dcp_created = cc._acquire_pynccl([0, 1], device) + other, _, other_created = cc._acquire_pynccl([0, 2], device) + + assert tp is dcp and tp_key == dcp_key + assert tp_created and not dcp_created + assert other is not tp and other_created + + cc._release_pynccl(dcp, dcp_key) + assert tp.destroyed == 0 + cc._release_pynccl(tp, tp_key) + assert tp.destroyed == 1 + assert tp_key not in cc._SHARED_PYNCCL + + +def test_communicators_are_not_shared_by_default(fake_pynccl, monkeypatch): + monkeypatch.delenv("VLLM_SHARE_PYNCCL_COMMS", raising=False) + device = torch.device("cuda:0") + first, first_key, _ = cc._acquire_pynccl([0, 1], device) + second, second_key, _ = cc._acquire_pynccl([0, 1], device) + + assert first is not second + assert first_key is None and second_key is None + cc._release_pynccl(first, first_key) + assert first.destroyed == 1 and second.destroyed == 0 diff --git a/vllm/envs.py b/vllm/envs.py index 3f7c05734003..dd87b8b883a2 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -925,9 +925,7 @@ def _resolve_rust_cli_path() -> str | None: ), # Keep the GLM-5.3 input embedding table in pinned host RAM; the GPU reads # the rows of the current tokens through a UVA view. - "VLLM_GLM53_EMBED_HOST": lambda: bool( - int(os.getenv("VLLM_GLM53_EMBED_HOST", "0")) - ), + "VLLM_GLM53_EMBED_HOST": lambda: bool(int(os.getenv("VLLM_GLM53_EMBED_HOST", "0"))), # Groups over the same ranks share one PyNCCL communicator. "VLLM_SHARE_PYNCCL_COMMS": lambda: bool( int(os.getenv("VLLM_SHARE_PYNCCL_COMMS", "0")) diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index e4ec78697ea2..aa36dab1c009 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -1699,9 +1699,7 @@ def host_embedding_if_requested(embed: VocabParallelEmbedding) -> None: if not is_uva_available(): raise RuntimeError("VLLM_GLM53_EMBED_HOST requires UVA and pinned memory") weight = embed.weight - host = torch.empty( - weight.shape, dtype=weight.dtype, device="cpu", pin_memory=True - ) + host = torch.empty(weight.shape, dtype=weight.dtype, device="cpu", pin_memory=True) weight.data = get_accelerator_view_from_cpu_tensor(host) weight._vllm_is_uva_offloaded = True From cfd58db76905f5e19c57749e1493236300aa4046 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 16:55:37 +0000 Subject: [PATCH 250/365] chore: release fragment vllm-898 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-898.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-898.json diff --git a/.lil/changes/vllm-898.json b/.lil/changes/vllm-898.json new file mode 100644 index 000000000000..a039b21c3905 --- /dev/null +++ b/.lil/changes/vllm-898.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-898", + "category": "performance", + "summary": "GLM-5.3-Flash can free about 0.9 GiB per GPU at TP2 for the KV cache", + "models": ["GLM-5.3-Flash"], + "compatibility": "Opt-in through VLLM_GLM53_EMBED_HOST=1, VLLM_GLM53_VISION_MXFP8=1 and VLLM_SHARE_PYNCCL_COMMS=1; all default to off.", + "details": [ + "VLLM_GLM53_EMBED_HOST keeps the input embedding table in pinned host RAM and reads the rows of the current tokens through UVA (0.59 GiB per GPU, lossless). Decode speed is unchanged; prefill is 1-3% slower.", + "VLLM_GLM53_VISION_MXFP8 quantizes the vision tower's linear layers to MXFP8 at load (0.24 GiB per GPU). ChartQA 79.8% vs 79.5% and OCR error 0.13% vs 0.07% against BF16.", + "VLLM_SHARE_PYNCCL_COMMS lets TP, DCP and EP groups over the same GPUs share one NCCL communicator (74 MiB per GPU at TP2+DCP2)." + ], + "pull_requests": [898] +} From d7301dd9700750689db26231428673032832bff2 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 21:31:03 +0000 Subject: [PATCH 251/365] fix(engine): bound the boundary-checkpoint copy wait like execute_model wait_for_boundary_checkpoint_copies was the only step-path RPC without a timeout; a wedged copy stream would stall the engine loop forever. It now uses VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, so the engine fails and /health reports it instead. The worker still synchronizes on the event, so the healthy path is unchanged. Co-Authored-By: Claude Opus 5.5 --- vllm/v1/engine/core.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ce12bb2744b9..316e5fd22961 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -706,7 +706,12 @@ def _wait_for_boundary_checkpoint_copies( self, model_output: ModelRunnerOutput ) -> None: if model_output.boundary_checkpoint_tokens is not None: - self.model_executor.collective_rpc("wait_for_boundary_checkpoint_copies") + # Bounded like execute_model: a wedged copy stream ends the engine + # instead of stalling the step loop forever. + self.model_executor.collective_rpc( + "wait_for_boundary_checkpoint_copies", + timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, + ) def step_with_batch_queue( self, From 21080615ed1dee55e50e1bef54269110fb2a4c8b Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 21:31:14 +0000 Subject: [PATCH 252/365] chore: release fragment vllm-903 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-903.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .lil/changes/vllm-903.json diff --git a/.lil/changes/vllm-903.json b/.lil/changes/vllm-903.json new file mode 100644 index 000000000000..736597c25431 --- /dev/null +++ b/.lil/changes/vllm-903.json @@ -0,0 +1,12 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-903", + "category": "fix", + "summary": "The boundary-checkpoint copy wait has the same timeout as model execution", + "models": ["GLM-5.3-Flash", "Qwen3.8-Flash-Next"], + "compatibility": "Uses VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (default 300 s); the healthy path is unchanged.", + "details": [ + "wait_for_boundary_checkpoint_copies was the only RPC on the engine step path without a timeout, so a wedged copy stream would stall the engine forever without an error. It now fails the engine after the execute-model timeout, and /health reports it." + ], + "pull_requests": [903] +} From 72847c56d4332aefcddef9bb1dd9b69d72c112c3 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 21:51:58 +0000 Subject: [PATCH 253/365] feat(core): boundary checkpoints for image and video requests Requests with image or video inputs were excluded from boundary checkpoints, so after eviction from GPU memory an external cache could not restore them and one image at the start of a conversation forced a full re-prefill every turn. content_token_ids() names each multimodal placeholder span by its content hash: positions in the span become pseudo-IDs in [2**30, 2**31) derived from the item identifier and the position. Checkpoint tails use this view, so an image wholly after the last full hash block is also part of the identity; the hash-block roots already include multimodal keys. Text requests are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../test_boundary_checkpoint_multimodal.py | 110 ++++++++++++++++++ vllm/v1/core/boundary_checkpoint.py | 61 +++++++++- 2 files changed, 165 insertions(+), 6 deletions(-) create mode 100644 tests/v1/core/test_boundary_checkpoint_multimodal.py diff --git a/tests/v1/core/test_boundary_checkpoint_multimodal.py b/tests/v1/core/test_boundary_checkpoint_multimodal.py new file mode 100644 index 000000000000..cd8dd60aa936 --- /dev/null +++ b/tests/v1/core/test_boundary_checkpoint_multimodal.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Boundary checkpoints for requests with image or video placeholders.""" + +from types import SimpleNamespace + +import pytest + +from tests.v1.core.test_boundary_admission import drain +from tests.v1.core.test_boundary_admission import ( + initialize_hash as initialize_hash, # noqa: F401 +) +from tests.v1.core.test_boundary_admission import manager as make_manager +from tests.v1.core.test_prefix_caching import make_request +from vllm.multimodal.inputs import PlaceholderRange +from vllm.utils.hashing import sha256 +from vllm.v1.core.boundary_checkpoint import ( + BoundaryCheckpointCache, + content_token_ids, +) +from vllm.v1.request import RequestStatus + + +def image_request(name, image, offset, length=140, span=40): + result = make_request( + name, + list(range(length)), + 16, + sha256, + mm_positions=[PlaceholderRange(offset=offset, length=span)], + mm_hashes=[image], + ) + result.max_tokens = result.sampling_params.max_tokens = 256 + return result + + +def publish(cache, producer, *ends): + """Compute the producer's prompt and publish checkpoints at ``ends``. + + All but the last end are instruction boundaries; the last is the prompt. + """ + if len(ends) > 1: + producer.recurrent_instruction_boundary = ends[0] + cache.get_computed_blocks(producer) + computed = 0 + for index, end in enumerate(ends): + assert cache.allocate_slots(producer, end - computed, num_lookahead_tokens=3) + producer.num_computed_tokens = computed = end + drain(cache) + kind = "prompt" if index == len(ends) - 1 else "instruction" + assert cache.publish_boundary_checkpoint(producer, end, kind=kind) + producer.status = RequestStatus.FINISHED_STOPPED + cache.free(producer) + + +def hit(cache, request): + return cache.get_computed_blocks(request)[1] + + +def test_content_tokens_name_placeholder_spans_by_content(): + def fake(identifier): + return SimpleNamespace( + all_token_ids=list(range(20)), + mm_features=[ + SimpleNamespace( + identifier=identifier, + mm_position=PlaceholderRange(offset=5, length=10), + ) + ], + ) + + text = SimpleNamespace(all_token_ids=list(range(20)), mm_features=None) + assert content_token_ids(text, 3, 12) == tuple(range(3, 12)) + a, b = fake("image-a"), fake("image-b") + assert content_token_ids(a, 0, 5) == tuple(range(5)) + assert content_token_ids(a, 15, 20) == tuple(range(15, 20)) + assert content_token_ids(a, 0, 20) != content_token_ids(b, 0, 20) + assert content_token_ids(a, 0, 20) == content_token_ids(fake("image-a"), 0, 20) + # A window inside the span matches the same positions of the full view. + assert content_token_ids(a, 7, 9) == content_token_ids(a, 0, 20)[7:9] + # Salt (for example encoder precision) separates otherwise equal content. + assert content_token_ids(a, 0, 20, b"x") != content_token_ids(a, 0, 20) + assert all(1 << 30 <= value < 1 << 31 for value in content_token_ids(a, 5, 15)) + + +@pytest.mark.parametrize( + "modality,supported", [("image", True), ("video", True), ("audio", False)] +) +def test_supported_modalities(modality, supported): + request = image_request("r", "image-a", 20) + request.mm_features[0].modality = modality + assert BoundaryCheckpointCache.supports_request(request) is supported + + +def test_image_checkpoint_restores_only_for_the_same_image(): + manager = make_manager() + publish(manager, image_request("producer", "image-a", 20), 10, 140) + # Same text and image: the whole prompt restores. + assert hit(manager, image_request("same", "image-a", 20)) == 140 + # Different image: only the checkpoint that ends before the image applies. + assert hit(manager, image_request("other", "image-b", 20)) == 10 + + +def test_image_inside_the_final_partial_block_is_part_of_the_identity(): + manager = make_manager() + # The image lies wholly after the last full hash block (128), so the block + # hashes of both requests agree and only the tail comparison differs. + publish(manager, image_request("producer", "image-a", 130, span=8), 140) + assert hit(manager, image_request("same", "image-a", 130, span=8)) == 140 + assert hit(manager, image_request("other", "image-b", 130, span=8)) == 0 diff --git a/vllm/v1/core/boundary_checkpoint.py b/vllm/v1/core/boundary_checkpoint.py index cf5595009ce4..13e55bf6787f 100644 --- a/vllm/v1/core/boundary_checkpoint.py +++ b/vllm/v1/core/boundary_checkpoint.py @@ -3,6 +3,7 @@ """Exact request-boundary lookup and ownership for recurrent prefix caches.""" +import hashlib from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal @@ -14,6 +15,53 @@ MAX_BOUNDARY_STOP_TOKENS = 128 +# Modalities whose placeholder spans are fully identified by their content hash. +CHECKPOINT_MODALITIES = frozenset({"image", "video"}) +# Placeholder positions map into [2**30, 2**31): above every vocabulary and +# inside LMCache's token wire domain. +_CONTENT_TOKEN_BASE = 1 << 30 +_CONTENT_TOKEN_STEP = 0x9E3779B1 + + +def content_token_ids( + request: Request, start: int, end: int, salt: bytes = b"" +) -> tuple[int, ...]: + """Token IDs of ``[start, end)`` with each multimodal span named by content. + + Placeholder tokens are identical for every image, so a checkpoint keyed by + token IDs alone could restore one image's state for another. Positions in + a multimodal span are replaced by pseudo-IDs derived from the item's + content hash and the position inside the span. Text requests are unchanged. + + Args: + request: Request whose ``all_token_ids`` and ``mm_features`` are read. + start: First position, inclusive. + end: Last position, exclusive. + salt: Extra identity (for example encoder precision) mixed into the IDs. + """ + tokens = request.all_token_ids[start:end] + features = request.mm_features + if not features: + return tuple(tokens) + tokens = list(tokens) + for feature in features: + offset = feature.mm_position.offset + length = feature.mm_position.length + lo, hi = max(offset, start), min(offset + length, end) + if lo >= hi: + continue + digest = hashlib.sha256( + b"vllm-mm-content-tokens-v1\0" + salt + b"\0" + feature.identifier.encode() + ).digest() + base = int.from_bytes(digest[:8], "little") + for position in range(lo, hi): + value = ( + base + (position - offset) * _CONTENT_TOKEN_STEP + ) % _CONTENT_TOKEN_BASE + tokens[position - start] = _CONTENT_TOKEN_BASE + value + return tuple(tokens) + + # Optional checkpoints have fixed slots; absent slots contain no KV blocks. PROMPT_CHECKPOINT_SLOT = 0 RESPONSE_CHECKPOINT_SLOT = 1 @@ -187,7 +235,10 @@ def supports_request(request: Request) -> bool: return ( request.prompt_token_ids is not None and request.prompt_embeds is None - and not request.mm_features + and all( + feature.modality in CHECKPOINT_MODALITIES and feature.identifier + for feature in request.mm_features or () + ) and not request.resumable and request.sampling_params is not None and len(request.sampling_params.stop_token_ids or ()) @@ -222,7 +273,7 @@ def stage( pending = _PendingCheckpoint( checkpoint, self._root_key(request, start), - tuple(request.all_token_ids[start : checkpoint.num_tokens]), + content_token_ids(request, start, checkpoint.num_tokens), num_ranks, ) self.block_pool.touch(blocks) @@ -289,10 +340,8 @@ def find(self, request: Request, max_length: int) -> BoundaryCheckpoint | None: root = self._roots.get(self._root_key(request, offset)) if root is None: continue - tokens = tuple( - request.all_token_ids[ - offset : min(offset + self.hash_block_size, max_length) - ] + tokens = content_token_ids( + request, offset, min(offset + self.hash_block_size, max_length) ) checkpoint_id = root.find(tokens) if checkpoint_id is not None: From 7c4fea8375ceebd78dfcda8d4255c70854fdfe42 Mon Sep 17 00:00:00 2001 From: Martin Date: Fri, 25 Sep 2026 22:06:20 +0000 Subject: [PATCH 254/365] chore: release fragment vllm-904 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-904.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 .lil/changes/vllm-904.json diff --git a/.lil/changes/vllm-904.json b/.lil/changes/vllm-904.json new file mode 100644 index 000000000000..250725601d49 --- /dev/null +++ b/.lil/changes/vllm-904.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-904", + "category": "feature", + "summary": "Image and video conversations get boundary checkpoints, so an external cache can restore them", + "models": ["GLM-5.3-Flash"], + "compatibility": "No user action required. Text requests are unchanged.", + "details": [ + "Requests with images or video were excluded from boundary checkpoints, so after eviction from GPU memory or a restart they were computed from scratch, including all text after the image.", + "Placeholder positions are now named by the image's content hash, so a checkpoint of one image is never restored for another. In GPU memory the next turn also reuses the exact request boundary instead of the last block boundary (15,331 instead of 12,288 of 15,410 tokens in our test)." + ], + "pull_requests": [904] +} From 44d23dfa2977e2e89ab055d216b62f67421923d4 Mon Sep 17 00:00:00 2001 From: Zyann <62597503+Zyann7@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:25:56 +0800 Subject: [PATCH 255/365] [Feature] support bf16 MoE router and mxfp4 MoE for MiMo V2 (#57784) Signed-off-by: Zyann7 <62597503+Zyann7@users.noreply.github.com> Co-authored-by: Abatom <182586866+Abatom@users.noreply.github.com> --- vllm/model_executor/models/mimo_v2.py | 13 ++++++++----- vllm/model_executor/models/qwen3_dflash.py | 9 +++++++++ 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 78af26326ca9..40301be49864 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -27,6 +27,7 @@ FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -156,14 +157,16 @@ def __init__( dtype = getattr(config, "moe_router_dtype", "float32") self.gate_dtype = str_dtype_to_torch_dtype(dtype) - self.gate = nn.Linear( + self.gate = GateLinear( config.hidden_size, config.n_routed_experts, bias=False, - dtype=self.gate_dtype, + params_dtype=self.gate_dtype, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", ) self.gate.e_score_correction_bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=self.gate_dtype) + torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype) ) self.experts = FusedMoEFactory( @@ -182,7 +185,7 @@ def __init__( num_expert_group=config.n_group, topk_group=config.topk_group, scoring_func="sigmoid", - router_logits_dtype=self.gate_dtype, + router_logits_dtype=self.gate.out_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -198,7 +201,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: gate_input = hidden_states.to(self.gate_dtype) else: gate_input = hidden_states - router_logits = self.gate(gate_input) + router_logits, _ = self.gate(gate_input) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3b5b01f5afb2..b8ca17bc90c5 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -207,6 +207,7 @@ def __init__( rms_norm_eps: float = 1e-06, attention_bias: bool = False, add_swa_attention_sink_bias: bool = False, + v_scale: float | None = None, sliding_window: int | None = None, causal: bool = False, is_neox_style: bool = True, @@ -232,6 +233,7 @@ def __init__( self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 + self.v_scale = v_scale self.qkv_proj = QKVParallelLinear( hidden_size, @@ -303,6 +305,8 @@ def forward( q, k = self.rotary_emb(positions, q, k) + if self.v_scale is not None: + v = v * self.v_scale attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -351,6 +355,7 @@ def __init__( rms_norm_eps=config.rms_norm_eps, attention_bias=getattr(config, "attention_bias", False), add_swa_attention_sink_bias=add_swa_attention_sink_bias, + v_scale=dflash_config.get("attention_value_scale"), sliding_window=sliding_window, causal=causal, is_neox_style=is_neox_style, @@ -749,6 +754,10 @@ def precompute_and_store_context_kv( if context_slot_mapping is None: return + v_scale = getattr(self.layers[0].self_attn, "v_scale", None) + if v_scale is not None: + all_v.mul_(v_scale) + # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) per_layer = isinstance(context_slot_mapping, (list, tuple)) From b6b22503f84339d44db92d60684a834780739777 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:55:54 +0000 Subject: [PATCH 256/365] test(mimo): verify router dtype and DFlash value scaling --- .../test_mimo_v2_router_dflash.py | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 tests/models/quantization/test_mimo_v2_router_dflash.py diff --git a/tests/models/quantization/test_mimo_v2_router_dflash.py b/tests/models/quantization/test_mimo_v2_router_dflash.py new file mode 100644 index 000000000000..34b460fcc77e --- /dev/null +++ b/tests/models/quantization/test_mimo_v2_router_dflash.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Router dtype and DFlash value scaling, with independent CPU expectations.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.models.mimo_v2 as mimo +import vllm.model_executor.models.qwen3_dflash as dflash + +pytestmark = pytest.mark.cpu_test + + +def test_bf16_router_emits_fp32_logits(monkeypatch): + for module in ( + mimo, + __import__("vllm.model_executor.layers.linear", fromlist=["_"]), + __import__("vllm.model_executor.parameter", fromlist=["_"]), + ): + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(mimo, "get_ep_group", lambda: SimpleNamespace( + device_group=SimpleNamespace(size=lambda: 1))) + config = SimpleNamespace( + hidden_size=32, n_routed_experts=8, hidden_act="silu", + moe_router_dtype="bfloat16", num_experts_per_tok=2, + moe_intermediate_size=64, norm_topk_prob=True, n_group=1, topk_group=1, + ) + vconfig = SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=config), quant_config=None, + parallel_config=SimpleNamespace(use_sequence_parallel_moe=False, + enable_eplb=False, eplb_config=SimpleNamespace(num_redundant_experts=0)), + ) + monkeypatch.setattr(mimo, "get_current_vllm_config", lambda: vconfig) + captured = {} + + class Experts(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + captured.update(kwargs) + + def forward(self, hidden_states, router_logits): + captured["logits"] = router_logits + return hidden_states + + monkeypatch.setattr(mimo, "FusedMoEFactory", Experts) + module = mimo.MiMoV2MoE(vconfig) + with torch.no_grad(): + module.gate.weight.copy_(torch.arange(256).view(8, 32) / 256) + hidden = torch.linspace(-1, 1, 96).view(3, 32).bfloat16() + result = module(hidden) + assert module.gate.weight.dtype == torch.bfloat16 + assert captured["router_logits_dtype"] == torch.float32 + assert module.gate.e_score_correction_bias.dtype == torch.float32 + reference = torch.nn.functional.linear(hidden, module.gate.weight).float() + torch.testing.assert_close(captured["logits"], reference) + torch.testing.assert_close(result, hidden) + + +@pytest.mark.parametrize("value_scale", [None, 0.612]) +def test_dflash_query_values_scaled_once(value_scale): + qkv = torch.arange(24, dtype=torch.float32).view(2, 12) + captured = {} + + def attention(q, k, v): + captured["v"] = v.clone() + return v + + module = SimpleNamespace( + qkv_proj=lambda _: (qkv, None), q_size=4, kv_size=4, head_dim=2, + q_norm=lambda x: x, k_norm=lambda x: x, + rotary_emb=lambda positions, q, k: (q, k), attn=attention, + o_proj=lambda x: (x, None), v_scale=value_scale, + ) + output = dflash.DFlashQwen3Attention.forward( + module, torch.arange(2), torch.empty(2, 4) + ) + expected = qkv[:, 8:] * (1 if value_scale is None else value_scale) + torch.testing.assert_close(captured["v"], expected) + torch.testing.assert_close(output, expected) + + +@pytest.mark.parametrize("value_scale", [None, 0.612]) +@pytest.mark.parametrize("per_layer", [False, True]) +def test_dflash_context_values_scaled_once(monkeypatch, value_scale, per_layer): + # Preserve the complete context method and intercept only native RoPE and + # the cache boundary. Different slot lengths exercise per-layer selection. + monkeypatch.setattr(dflash.ops, "rotary_embedding", lambda *args: None) + original = torch.arange(48, dtype=torch.float32).view(2, 3, 2, 4) + captures = [] + inner = [SimpleNamespace( + kv_cache=torch.empty(0), impl=SimpleNamespace( + do_kv_cache_update=lambda layer, k, v, cache, slots: + captures.append((v.clone(), slots.clone())))) for _ in range(2)] + module = SimpleNamespace( + _num_attn_layers=2, _kv_size=8, _head_dim=4, _num_kv_heads=2, + _project_context_kv=lambda *_: (original.clone(), original.clone()), + _normalize_context_k=lambda x: x, + _rope_cos_sin_cache=torch.zeros(8, 4), _rope_head_size=4, + _rope_is_neox=True, _attn_layers=inner, + layers=[SimpleNamespace(self_attn=SimpleNamespace(v_scale=value_scale)) + for _ in range(2)], + ) + slots = [torch.tensor([0, 1, -1]), torch.tensor([7, -1, -1])] + mapping = slots if per_layer else slots[0] + dflash.DFlashQwen3Model.precompute_and_store_context_kv( + module, torch.empty(3, 8), torch.arange(3), mapping + ) + assert len(captures) == 2 + for index, (values, actual_slots) in enumerate(captures): + expected = original[index] * (1 if value_scale is None else value_scale) + torch.testing.assert_close(values, expected) + torch.testing.assert_close(actual_slots, slots[index] if per_layer else slots[0]) From f4928a8c2432ad448c9f8cbc30894769a2972ca9 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 21:31:28 +0000 Subject: [PATCH 257/365] style(mimo): format router and draft regressions --- .../test_mimo_v2_router_dflash.py | 79 +++++++++++++------ 1 file changed, 57 insertions(+), 22 deletions(-) diff --git a/tests/models/quantization/test_mimo_v2_router_dflash.py b/tests/models/quantization/test_mimo_v2_router_dflash.py index 34b460fcc77e..a16e9ad683eb 100644 --- a/tests/models/quantization/test_mimo_v2_router_dflash.py +++ b/tests/models/quantization/test_mimo_v2_router_dflash.py @@ -21,17 +21,30 @@ def test_bf16_router_emits_fp32_logits(monkeypatch): ): monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) - monkeypatch.setattr(mimo, "get_ep_group", lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1))) + monkeypatch.setattr( + mimo, + "get_ep_group", + lambda: SimpleNamespace(device_group=SimpleNamespace(size=lambda: 1)), + ) config = SimpleNamespace( - hidden_size=32, n_routed_experts=8, hidden_act="silu", - moe_router_dtype="bfloat16", num_experts_per_tok=2, - moe_intermediate_size=64, norm_topk_prob=True, n_group=1, topk_group=1, + hidden_size=32, + n_routed_experts=8, + hidden_act="silu", + moe_router_dtype="bfloat16", + num_experts_per_tok=2, + moe_intermediate_size=64, + norm_topk_prob=True, + n_group=1, + topk_group=1, ) vconfig = SimpleNamespace( - model_config=SimpleNamespace(hf_text_config=config), quant_config=None, - parallel_config=SimpleNamespace(use_sequence_parallel_moe=False, - enable_eplb=False, eplb_config=SimpleNamespace(num_redundant_experts=0)), + model_config=SimpleNamespace(hf_text_config=config), + quant_config=None, + parallel_config=SimpleNamespace( + use_sequence_parallel_moe=False, + enable_eplb=False, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), ) monkeypatch.setattr(mimo, "get_current_vllm_config", lambda: vconfig) captured = {} @@ -69,10 +82,16 @@ def attention(q, k, v): return v module = SimpleNamespace( - qkv_proj=lambda _: (qkv, None), q_size=4, kv_size=4, head_dim=2, - q_norm=lambda x: x, k_norm=lambda x: x, - rotary_emb=lambda positions, q, k: (q, k), attn=attention, - o_proj=lambda x: (x, None), v_scale=value_scale, + qkv_proj=lambda _: (qkv, None), + q_size=4, + kv_size=4, + head_dim=2, + q_norm=lambda x: x, + k_norm=lambda x: x, + rotary_emb=lambda positions, q, k: (q, k), + attn=attention, + o_proj=lambda x: (x, None), + v_scale=value_scale, ) output = dflash.DFlashQwen3Attention.forward( module, torch.arange(2), torch.empty(2, 4) @@ -90,18 +109,32 @@ def test_dflash_context_values_scaled_once(monkeypatch, value_scale, per_layer): monkeypatch.setattr(dflash.ops, "rotary_embedding", lambda *args: None) original = torch.arange(48, dtype=torch.float32).view(2, 3, 2, 4) captures = [] - inner = [SimpleNamespace( - kv_cache=torch.empty(0), impl=SimpleNamespace( - do_kv_cache_update=lambda layer, k, v, cache, slots: - captures.append((v.clone(), slots.clone())))) for _ in range(2)] + inner = [ + SimpleNamespace( + kv_cache=torch.empty(0), + impl=SimpleNamespace( + do_kv_cache_update=lambda layer, k, v, cache, slots: captures.append( + (v.clone(), slots.clone()) + ) + ), + ) + for _ in range(2) + ] module = SimpleNamespace( - _num_attn_layers=2, _kv_size=8, _head_dim=4, _num_kv_heads=2, + _num_attn_layers=2, + _kv_size=8, + _head_dim=4, + _num_kv_heads=2, _project_context_kv=lambda *_: (original.clone(), original.clone()), _normalize_context_k=lambda x: x, - _rope_cos_sin_cache=torch.zeros(8, 4), _rope_head_size=4, - _rope_is_neox=True, _attn_layers=inner, - layers=[SimpleNamespace(self_attn=SimpleNamespace(v_scale=value_scale)) - for _ in range(2)], + _rope_cos_sin_cache=torch.zeros(8, 4), + _rope_head_size=4, + _rope_is_neox=True, + _attn_layers=inner, + layers=[ + SimpleNamespace(self_attn=SimpleNamespace(v_scale=value_scale)) + for _ in range(2) + ], ) slots = [torch.tensor([0, 1, -1]), torch.tensor([7, -1, -1])] mapping = slots if per_layer else slots[0] @@ -112,4 +145,6 @@ def test_dflash_context_values_scaled_once(monkeypatch, value_scale, per_layer): for index, (values, actual_slots) in enumerate(captures): expected = original[index] * (1 if value_scale is None else value_scale) torch.testing.assert_close(values, expected) - torch.testing.assert_close(actual_slots, slots[index] if per_layer else slots[0]) + torch.testing.assert_close( + actual_slots, slots[index] if per_layer else slots[0] + ) From c504861977ecbef4a2828312450be70307c3a919 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:46:17 +0000 Subject: [PATCH 258/365] Compare CPU router logits to an independent FP32 reference --- .../quantization/test_mimo_v2_router_dflash.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/models/quantization/test_mimo_v2_router_dflash.py b/tests/models/quantization/test_mimo_v2_router_dflash.py index a16e9ad683eb..efb1b99f003e 100644 --- a/tests/models/quantization/test_mimo_v2_router_dflash.py +++ b/tests/models/quantization/test_mimo_v2_router_dflash.py @@ -14,6 +14,17 @@ def test_bf16_router_emits_fp32_logits(monkeypatch): + # Exercise the CPU fallback without selecting a CUDA/ROCm kernel merely + # because the test process runs on a GPU host. + from vllm.model_executor.layers.fused_moe.router import gate_linear + + for name in ( + "is_cuda", + "is_rocm", + "is_device_capability", + "is_device_capability_family", + ): + monkeypatch.setattr(gate_linear.current_platform, name, lambda *a, **k: False) for module in ( mimo, __import__("vllm.model_executor.layers.linear", fromlist=["_"]), @@ -67,8 +78,10 @@ def forward(self, hidden_states, router_logits): assert module.gate.weight.dtype == torch.bfloat16 assert captured["router_logits_dtype"] == torch.float32 assert module.gate.e_score_correction_bias.dtype == torch.float32 - reference = torch.nn.functional.linear(hidden, module.gate.weight).float() - torch.testing.assert_close(captured["logits"], reference) + reference = torch.nn.functional.linear(hidden.float(), module.gate.weight.float()) + # The CPU fallback rounds to BF16 before casting the result to FP32. + # SM120 accumulation is checked independently by the graph replay test. + torch.testing.assert_close(captured["logits"], reference, atol=0.01, rtol=0.004) torch.testing.assert_close(result, hidden) From af94e88373402eaa68f035e85217cc3253550f1a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 09:22:47 +0000 Subject: [PATCH 259/365] docs(release): add fragments for vllm-829, vllm-831 and vllm-832 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-829.json | 18 ++++++++++++++++++ .lil/changes/vllm-831.json | 16 ++++++++++++++++ .lil/changes/vllm-832.json | 18 ++++++++++++++++++ 3 files changed, 52 insertions(+) create mode 100644 .lil/changes/vllm-829.json create mode 100644 .lil/changes/vllm-831.json create mode 100644 .lil/changes/vllm-832.json diff --git a/.lil/changes/vllm-829.json b/.lil/changes/vllm-829.json new file mode 100644 index 000000000000..eeb1c172ffc4 --- /dev/null +++ b/.lil/changes/vllm-829.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-829", + "category": "fix", + "summary": "MiMo-V2 router logits stay FP32 and the DFlash drafter applies its value scale", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "Backport of vllm-project/vllm#57784 (Zyann7): MiMo's BF16 gate weights produce FP32 router logits, as in the model's reference code, so near-tie expert choices no longer flip.", + "The DFlash decoder applies its non-unit value scale exactly once to query-path and cached-context values, which raises draft acceptance." + ], + "pull_requests": [ + 829 + ] +} diff --git a/.lil/changes/vllm-831.json b/.lil/changes/vllm-831.json new file mode 100644 index 000000000000..2a085fbc5ee9 --- /dev/null +++ b/.lil/changes/vllm-831.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-831", + "category": "fix", + "summary": "The SM120 BF16-to-FP32 router matmul compiles under fullgraph compilation", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "The CUDA-graph capture guard, its BF16 allocation and the FP32-output matmul run inside one SM120 custom op, so torch.compile no longer traces a Python bool from the capture query. Other platforms and router paths are unchanged." + ], + "pull_requests": [ + 831 + ] +} diff --git a/.lil/changes/vllm-832.json b/.lil/changes/vllm-832.json new file mode 100644 index 000000000000..cab8f05e85d7 --- /dev/null +++ b/.lil/changes/vllm-832.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-832", + "category": "fix", + "summary": "--kv-cache-dtype fp8 reaches the MiMo target and MTP attention layers", + "models": [ + "XiaomiMiMo/MiMo-V2.6-Flash-RL", + "XiaomiMiMo/MiMo-V2.6-Pro-RL" + ], + "compatibility": "No user action required.", + "details": [ + "The MiMo target decoder and native MTP layers now receive the cache configuration, so an FP8 KV request allocates FP8 pages instead of BF16 and keeps scale calibration.", + "Global layers keep full attention: the model-wide sliding-window fallback is cleared on a copy of the configuration for them." + ], + "pull_requests": [ + 832 + ] +} From f85f0957f44a4c429fdbca226db16f1fafd54a53 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 11:02:45 +0000 Subject: [PATCH 260/365] fix(boundary-checkpoint): full-strength identity for image spans content_token_ids (#904) reduced a whole image span to one 30-bit seed (base % 2**30 plus a fixed step), so two different images could map to the same pseudo-token sequence. A reporter found such a PNG pair in 6 s with a birthday search. - v2 pseudo-IDs: counter-mode BLAKE2b keyed by the SHA-256 of the item's identifier and salt; every position carries 30 independent bits. - A checkpoint may not end within the first 5 positions of an item (150 bits); publishing and external imports skip such endpoints. - The version tag changes, so v1 image keys become unreachable. Co-Authored-By: Claude Opus 5.5 --- .../test_boundary_checkpoint_multimodal.py | 53 +++++++++++++++++++ vllm/v1/core/boundary_checkpoint.py | 44 +++++++++++---- vllm/v1/core/kv_cache_manager.py | 5 ++ 3 files changed, 92 insertions(+), 10 deletions(-) diff --git a/tests/v1/core/test_boundary_checkpoint_multimodal.py b/tests/v1/core/test_boundary_checkpoint_multimodal.py index cd8dd60aa936..e1798cc6aabf 100644 --- a/tests/v1/core/test_boundary_checkpoint_multimodal.py +++ b/tests/v1/core/test_boundary_checkpoint_multimodal.py @@ -15,7 +15,9 @@ from vllm.multimodal.inputs import PlaceholderRange from vllm.utils.hashing import sha256 from vllm.v1.core.boundary_checkpoint import ( + MIN_CHECKPOINT_ITEM_TOKENS, BoundaryCheckpointCache, + checkpoint_end_allowed, content_token_ids, ) from vllm.v1.request import RequestStatus @@ -108,3 +110,54 @@ def test_image_inside_the_final_partial_block_is_part_of_the_identity(): publish(manager, image_request("producer", "image-a", 130, span=8), 140) assert hit(manager, image_request("same", "image-a", 130, span=8)) == 140 assert hit(manager, image_request("other", "image-b", 130, span=8)) == 0 + + +def fake_image(identifier, offset=5, length=40, total=60): + return SimpleNamespace( + all_token_ids=list(range(total)), + mm_features=[ + SimpleNamespace( + identifier=identifier, + mm_position=PlaceholderRange(offset=offset, length=length), + ) + ], + ) + + +def test_content_tokens_carry_independent_bits_per_position(): + # Two real PNGs whose processor hashes collided under the v1 scheme, where + # a whole span was a function of one 30-bit seed (Discord report, #904). + red = "2373f7d58a45946947a5179eff2f1a6fc5c6ccac30b9ef4e6999b17e54901057" + blue = "cd07abfe9657e3a97e5a4c82b6baf82ab190e9a150aba160577c7e92b7b82219" + salt = b'{"VLLM_GLM53_VISION_MXFP8": "1"}' + a = content_token_ids(fake_image(red), 5, 45, salt) + b = content_token_ids(fake_image(blue), 5, 45, salt) + assert all(x != y for x, y in zip(a, b)) + # Positions are not an arithmetic progression of one seed. + steps = {(y - x) % (1 << 30) for x, y in zip(a, a[1:])} + assert len(steps) > 1 + # Windows crossing counter blocks match the full view. + full = content_token_ids(fake_image(red), 0, 60) + for start, end in ((3, 22), (20, 21), (21, 38), (44, 60)): + assert content_token_ids(fake_image(red), start, end) == full[start:end] + + +def test_checkpoints_cannot_end_at_the_start_of_an_image(): + request = fake_image("image-a", offset=5, length=40) + allowed = [checkpoint_end_allowed(request, end) for end in range(1, 60)] + shallow = range(6, 5 + MIN_CHECKPOINT_ITEM_TOKENS) + assert [end for end, ok in zip(range(1, 60), allowed) if not ok] == list(shallow) + # A tiny item may end a checkpoint once it is complete. + assert checkpoint_end_allowed(fake_image("icon", offset=5, length=2), 7) + assert not checkpoint_end_allowed(fake_image("icon", offset=5, length=2), 6) + + +def test_shallow_image_endpoint_is_not_published(): + manager = make_manager() + producer = image_request("producer", "image-a", 20) + producer.recurrent_instruction_boundary = 22 + manager.get_computed_blocks(producer) + assert manager.allocate_slots(producer, 22, num_lookahead_tokens=3) + producer.num_computed_tokens = 22 + drain(manager) + assert manager.publish_boundary_checkpoint(producer, 22, kind="instruction") is None diff --git a/vllm/v1/core/boundary_checkpoint.py b/vllm/v1/core/boundary_checkpoint.py index 13e55bf6787f..18a1e852ac0c 100644 --- a/vllm/v1/core/boundary_checkpoint.py +++ b/vllm/v1/core/boundary_checkpoint.py @@ -18,9 +18,12 @@ # Modalities whose placeholder spans are fully identified by their content hash. CHECKPOINT_MODALITIES = frozenset({"image", "video"}) # Placeholder positions map into [2**30, 2**31): above every vocabulary and -# inside LMCache's token wire domain. +# inside LMCache's token wire domain. Each position carries 30 independent +# bits of the item's full content hash. _CONTENT_TOKEN_BASE = 1 << 30 -_CONTENT_TOKEN_STEP = 0x9E3779B1 +# A checkpoint may not end within an item's first positions: with fewer of +# them its key would hold under 150 bits of that item's identity. +MIN_CHECKPOINT_ITEM_TOKENS = 5 def content_token_ids( @@ -50,18 +53,37 @@ def content_token_ids( lo, hi = max(offset, start), min(offset + length, end) if lo >= hi: continue - digest = hashlib.sha256( - b"vllm-mm-content-tokens-v1\0" + salt + b"\0" + feature.identifier.encode() + key = hashlib.sha256( + b"vllm-mm-content-tokens-v2\0" + salt + b"\0" + feature.identifier.encode() ).digest() - base = int.from_bytes(digest[:8], "little") - for position in range(lo, hi): - value = ( - base + (position - offset) * _CONTENT_TOKEN_STEP - ) % _CONTENT_TOKEN_BASE - tokens[position - start] = _CONTENT_TOKEN_BASE + value + # Counter-mode BLAKE2b keyed by the full digest: 16 values per call, + # independent across positions, so a span is not one 30-bit seed. + first, last = lo - offset, hi - offset + for block in range(first // 16, (last - 1) // 16 + 1): + words = hashlib.blake2b( + block.to_bytes(8, "little"), key=key, digest_size=64 + ).digest() + for index in range(max(first, block * 16), min(last, block * 16 + 16)): + word = words[(index % 16) * 4 : (index % 16) * 4 + 4] + value = int.from_bytes(word, "little") & (_CONTENT_TOKEN_BASE - 1) + tokens[offset + index - start] = _CONTENT_TOKEN_BASE + value return tuple(tokens) +def checkpoint_end_allowed(request: Request, num_tokens: int) -> bool: + """Whether a checkpoint may end after ``num_tokens`` of this request. + + It must not end within the first MIN_CHECKPOINT_ITEM_TOKENS positions of + a multimodal item: its key would identify that item by too few bits. + """ + for feature in request.mm_features or (): + offset = feature.mm_position.offset + depth = min(MIN_CHECKPOINT_ITEM_TOKENS, feature.mm_position.length) + if offset < num_tokens < offset + depth: + return False + return True + + # Optional checkpoints have fixed slots; absent slots contain no KV blocks. PROMPT_CHECKPOINT_SLOT = 0 RESPONSE_CHECKPOINT_SLOT = 1 @@ -257,6 +279,8 @@ def stage( raise ValueError("Boundary checkpoints require a text generation request") if not 0 < checkpoint.num_tokens <= request.num_tokens: raise ValueError("Checkpoint must cover an existing, nonempty prefix") + if not checkpoint_end_allowed(request, checkpoint.num_tokens): + raise ValueError("Checkpoint cannot end at the start of a multimodal item") if not 0 <= checkpoint.draft_prefix_len <= checkpoint.num_tokens: raise ValueError("Draft prefix cannot extend beyond the target prefix") if num_ranks < 1: diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 4ba5e0eb46b9..de3b56ff7183 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -18,6 +18,7 @@ BoundaryCheckpointCache, BoundaryCheckpointKind, boundary_checkpoint_slots, + checkpoint_end_allowed, ) from vllm.v1.core.kv_cache_coordinator import ( HybridKVCacheCoordinator, @@ -860,6 +861,8 @@ def publish_boundary_checkpoint( allocation = request.boundary_checkpoint_blocks if cache is None or allocation is None or num_tokens <= 0: return None + if not checkpoint_end_allowed(request, num_tokens): + return None if kind == "response": if ( request.status @@ -984,6 +987,8 @@ def reserve_external_boundary_checkpoint( return None if not 0 < num_tokens <= request.num_tokens: raise ValueError("Imported boundary must cover an existing request prefix") + if not checkpoint_end_allowed(request, num_tokens): + return None if not 0 <= draft_prefix_len <= num_tokens: raise ValueError("Imported draft prefix exceeds the target prefix") if ( From b209dcd90c1053fe2b09e50604e39b3dfcc5903e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 11:02:57 +0000 Subject: [PATCH 261/365] docs(release): add vllm-908 release fragment Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-908.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-908.json diff --git a/.lil/changes/vllm-908.json b/.lil/changes/vllm-908.json new file mode 100644 index 000000000000..eb3bd20bd7d2 --- /dev/null +++ b/.lil/changes/vllm-908.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-908", + "category": "fix", + "summary": "Image checkpoints are keyed by the full image identity, not a 30-bit seed", + "models": ["all"], + "compatibility": "No user action required. Image checkpoints saved by earlier builds are not reused and recompute once; text checkpoints are unaffected.", + "details": [ + "Two different images could get the same boundary-checkpoint key (about 1 in a billion per pair, and constructible on purpose), so one image's cached state could be restored for another.", + "Every image position now carries independent bits of the image's content hash, and no checkpoint ends within an image's first five tokens." + ], + "pull_requests": [908], + "requires": [] +} From 40371c7bb0b7b6547ae75d43ba35db12488fc9ef Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 13:30:31 +0000 Subject: [PATCH 262/365] DS4.1: BF16 sparse attention arithmetic and FP32 logits by default b12x's V4.1 attention tuner may pick FP8 internals (Q and P quantized, the cache re-encoded to FP8), about 3% relative error against DeepSeek's reference sparse_attn, with no end-to-end speed gain on RTX PRO 6000 TP4. Pin BF16 arithmetic by default; VLLM_DS41_ATTENTION_COMPUTE=reference also runs decode rows single-pass (bit-exact with the reference kernel), =auto keeps the tuner choice. The reference head computes FP32 logits; vLLM rounded them to BF16 (steps of up to 1/8). On SM120, DS4.1 target and DSpark draft default to head_dtype=float32 unless --hf-overrides sets it. head_dtype=float32 used to raise for runtime-quantized LM heads, so it could not run with the optional NVFP4 DSpark drafter/Markov heads; such heads now keep their kernel and have their logits widened. Co-Authored-By: Claude Opus 5.5 --- .../vllm-ds41-bf16-attention-fp32-heads.json | 17 ++++++ tests/quantization/test_lm_head.py | 27 +++++++++ vllm/envs.py | 13 +++++ .../model_executor/layers/logits_processor.py | 15 ++++- vllm/model_executor/models/config.py | 4 ++ vllm/models/deepseek_v4_1/attention.py | 57 +++++++++++-------- 6 files changed, 107 insertions(+), 26 deletions(-) create mode 100644 .lil/changes/vllm-ds41-bf16-attention-fp32-heads.json diff --git a/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json new file mode 100644 index 000000000000..9320634c9492 --- /dev/null +++ b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-ds41-bf16-attention-fp32-heads", + "category": "change", + "summary": "DeepSeek-V4.1 computes attention in BF16 and logits in FP32 like the reference", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "VLLM_DS41_ATTENTION_COMPUTE=auto restores the previous tuner choice (FP8 or BF16 per plan); --hf-overrides '{\"head_dtype\": \"model\"}' restores BF16 logits. Output token streams change within numerical noise.", + "details": [ + "b12x V4.1 sparse attention is pinned to BF16 arithmetic (BF16 Q, KV dequantized to BF16, BF16 P, FP32 accumulation) instead of the tuner's FP8 internals, which carried about 3% relative error against DeepSeek's reference sparse_attn. Measured throughput on 4x RTX PRO 6000 is unchanged.", + "VLLM_DS41_ATTENTION_COMPUTE=reference also runs decode rows single-pass, reproducing the reference kernel bit for bit at up to 17% lower low-concurrency decode speed.", + "On RTX PRO / SM120, DeepSeek-V4.1 (target and DSpark draft) defaults to head_dtype=float32: the LM head accumulates into FP32 logits instead of rounding them to BF16 steps of up to 1/8, matching the reference head. Measured decode speed is unchanged.", + "head_dtype=float32 no longer rejects runtime-quantized LM heads (for example the optional NVFP4 DSpark drafter and Markov heads): their logits are widened to FP32." + ], + "authors": ["voipmonitor"], + "pull_requests": [], + "requires": [] +} diff --git a/tests/quantization/test_lm_head.py b/tests/quantization/test_lm_head.py index 63b020a9f9d7..e0d8ee6eb28e 100644 --- a/tests/quantization/test_lm_head.py +++ b/tests/quantization/test_lm_head.py @@ -75,6 +75,33 @@ def mxfp8_head_config(monkeypatch, default_vllm_config): return default_vllm_config +@pytest.mark.cpu_test +def test_fp32_head_dtype_widens_quantized_head_logits(mxfp8_head_config): + """A runtime-quantized head (e.g. the NVFP4 DSpark drafter head) keeps its + own kernel under an FP32 head_dtype; its logits are widened, not rejected.""" + from vllm.model_executor.layers.logits_processor import LogitsProcessor + + mxfp8_head_config.model_config.head_dtype = torch.float32 + processor = LogitsProcessor(8) + calls = [] + + class QuantizedHeadMethod: + def apply(self, layer, x, bias=None): + calls.append(x.dtype) + return x @ layer.weight.t() + + head = SimpleNamespace( + quant_method=QuantizedHeadMethod(), + weight=torch.randn(8, 4).bfloat16(), + tp_size=1, + ) + hidden = torch.randn(3, 4, dtype=torch.bfloat16) + logits = processor(head, hidden, skip_gather=True) + assert calls == [torch.bfloat16] + assert logits.dtype == torch.float32 + torch.testing.assert_close(logits, (hidden @ head.weight.t()).float()) + + @pytest.mark.cpu_test @pytest.mark.parametrize( ("env_value", "enabled"), [(None, False), ("0", False), ("1", True)] diff --git a/vllm/envs.py b/vllm/envs.py index 1eee6d9e2052..8052fb1d502b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -212,6 +212,7 @@ VLLM_MTP_NVFP4_LM_HEAD: bool = True VLLM_DS41_MARKOV_NVFP4: bool = False VLLM_DS41_DRAFT_NVFP4_HEAD: bool = False + VLLM_DS41_ATTENTION_COMPUTE: Literal["bf16", "reference", "auto"] = "bf16" VLLM_DS41_ENGRAM_OVERLAP: bool = True VLLM_QWEN3_8_FLASH_NEXT_OVERLAP: bool = True VLLM_QWEN3_8_FLASH_NEXT_HC_TP: bool = True @@ -1753,6 +1754,18 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_DS41_DRAFT_NVFP4_HEAD": lambda: bool( int(os.getenv("VLLM_DS41_DRAFT_NVFP4_HEAD", "0")) ), + # Arithmetic of the DeepSeek V4.1 b12x sparse attention. "bf16" computes + # like DeepSeek's reference sparse_attn (BF16 Q, KV dequantized to BF16, + # BF16 P, FP32 accumulation) with split-K decode; "reference" also runs + # decode rows single-pass, which reproduces the reference kernel's + # tile-ordered softmax bit for bit at some low-concurrency decode speed; + # "auto" lets the b12x tuner pick FP8 internals (Q and P quantized, KV + # re-encoded; about 3% relative error) or BF16 per plan. + "VLLM_DS41_ATTENTION_COMPUTE": env_with_choices( + "VLLM_DS41_ATTENTION_COMPUTE", + "bf16", + ["bf16", "reference", "auto"], + ), # Read DeepSeek V4.1 disk Engram rows while the target graph starts. "VLLM_DS41_ENGRAM_OVERLAP": lambda: bool( int(os.getenv("VLLM_DS41_ENGRAM_OVERLAP", "1")) diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 05bc57e80007..5a5b8cd3b1d8 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -323,10 +323,19 @@ def _apply_head( if not isinstance( lm_head.quant_method, (UnquantizedEmbeddingMethod, UnquantizedLinearMethod) ): - raise ValueError( - "A head_dtype different from the model dtype is only " - "supported for an unquantized lm_head." + # A runtime-quantized head (for example a drafter's NVFP4 head next + # to an FP32 target head) keeps its own kernel; its logits are only + # widened, since the quantized weights already bound its precision. + logger.warning_once( + "head_dtype=%s: a quantized LM head computes in its own format; " + "its logits are widened to %s.", + self.head_dtype, + self.head_dtype, + ) + logits = lm_head.quant_method.apply( + lm_head, hidden_states, bias=embedding_bias ) + return logits.to(self.head_dtype) if ( self.head_dtype == torch.float32 and (current_platform.is_cuda() or current_platform.is_rocm()) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 8c5ac89c6de8..cc0281322738 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -377,6 +377,10 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: and quant_config.get("quant_method") == "fp8" ): quant_config["quant_method"] = "deepseek_v41_fp8" + # DeepSeek's reference head computes FP32 logits; a BF16 head rounds + # them to steps of up to 1/8. An explicit head_dtype override wins. + if getattr(model_config.hf_config, "head_dtype", None) is None: + model_config.hf_config.head_dtype = "float32" @staticmethod def update_model_config_for_parallelism( diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 62327b41ebfb..2b2ade4bd321 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -14,6 +14,7 @@ from b12x.gemm import wo_projection from torch import nn +from vllm import envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group @@ -676,10 +677,13 @@ def descriptor(tensor, *, page_size, kind): ) self._attention_workspace_specs = {} self._attention_declarations = {} + compute = envs.VLLM_DS41_ATTENTION_COMPUTE for mode, capacity in ( ("decode", min(self.capacity, decode_rows)), ("extend", self.capacity), ): + # "reference" gives decode rows the single-pass kernel as well. + kernel_mode = "extend" if compute == "reference" else mode caps = mla.Caps( device=device, num_q_heads=self.n_local_heads, @@ -690,12 +694,12 @@ def descriptor(tensor, *, page_size, kind): swa_page_size=self.swa_cache_layer.block_size, indexed_page_size=self._main_page, max_page_table_width=self._main_width, - mode=mode, + mode=kernel_mode, # DSpark's two-span reservation can exceed the generic # 256-row cutoff even when replaying a six-row C1 graph. # Keep the entire reserved decode capacity on its decode # split contract; prefill retains its separate plan. - decode_row_capacity=capacity if mode == "decode" else None, + decode_row_capacity=capacity if kernel_mode == "decode" else None, cache_format="deepseek_v41", use_cuda_graph=True, ) @@ -705,29 +709,36 @@ def descriptor(tensor, *, page_size, kind): if self.compress_ratio else None ) - declaration = mla.plan( - caps, - invocation=mla.invocation_from_descriptors( - q={ - "shape": (capacity, self.n_local_heads, 512), - "stride": (self.n_local_heads * 512, 512, 1), - "alignment": 16, - "dtype": "bfloat16", - }, - swa_cache=descriptor( - swa_cache, page_size=self.swa_cache_layer.block_size, kind="swa" - ), - indexed_cache=( - descriptor( - indexed_cache, page_size=self._main_page, kind="indexed" - ) - if self.compress_ratio - else None - ), - attn_sink_present=True, - output_mode="provided", + invocation = mla.invocation_from_descriptors( + q={ + "shape": (capacity, self.n_local_heads, 512), + "stride": (self.n_local_heads * 512, 512, 1), + "alignment": 16, + "dtype": "bfloat16", + }, + swa_cache=descriptor( + swa_cache, page_size=self.swa_cache_layer.block_size, kind="swa" ), + indexed_cache=( + descriptor(indexed_cache, page_size=self._main_page, kind="indexed") + if self.compress_ratio + else None + ), + attn_sink_present=True, + output_mode="provided", ) + declaration = mla.plan(caps, invocation=invocation) + if compute != "auto": + # Pin BF16 arithmetic; the split/single-pass lowering stays the + # plan's own default for this query. + default = declaration.contract.default_config(declaration.query, None) + declaration = mla.plan( + caps, + invocation=invocation, + override=replace( + default, v41_compute_mode="bf16", v41_heads_per_block=16 + ), + ) metadata_specs = ( ((capacity, self.swa_width), torch.int32), ((capacity,), torch.int32), From 534d985ea0907ff5fdfcf4082c5e08db93a107a7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 15:18:46 +0000 Subject: [PATCH 263/365] docs(release): link the DS4.1 numerics fragment to #911 Co-Authored-By: Claude Opus 5.5 --- .../changes/vllm-ds41-bf16-attention-fp32-heads.json | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json index 9320634c9492..79f913f3ed03 100644 --- a/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json +++ b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json @@ -3,7 +3,9 @@ "id": "vllm-ds41-bf16-attention-fp32-heads", "category": "change", "summary": "DeepSeek-V4.1 computes attention in BF16 and logits in FP32 like the reference", - "models": ["DeepSeek-V4.1-Flash"], + "models": [ + "DeepSeek-V4.1-Flash" + ], "compatibility": "VLLM_DS41_ATTENTION_COMPUTE=auto restores the previous tuner choice (FP8 or BF16 per plan); --hf-overrides '{\"head_dtype\": \"model\"}' restores BF16 logits. Output token streams change within numerical noise.", "details": [ "b12x V4.1 sparse attention is pinned to BF16 arithmetic (BF16 Q, KV dequantized to BF16, BF16 P, FP32 accumulation) instead of the tuner's FP8 internals, which carried about 3% relative error against DeepSeek's reference sparse_attn. Measured throughput on 4x RTX PRO 6000 is unchanged.", @@ -11,7 +13,11 @@ "On RTX PRO / SM120, DeepSeek-V4.1 (target and DSpark draft) defaults to head_dtype=float32: the LM head accumulates into FP32 logits instead of rounding them to BF16 steps of up to 1/8, matching the reference head. Measured decode speed is unchanged.", "head_dtype=float32 no longer rejects runtime-quantized LM heads (for example the optional NVFP4 DSpark drafter and Markov heads): their logits are widened to FP32." ], - "authors": ["voipmonitor"], - "pull_requests": [], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 911 + ], "requires": [] } From 5d1af9b2cc56057b502424799945ba2643b1222c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 15:30:30 +0000 Subject: [PATCH 264/365] docs(release): DS4.1 numerics fragment uses the fix category The container publisher accepts breaking, compatibility, feature, fix, internal and performance; "change" stopped the beta release at changelog collection. The fragment has not been published yet. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-ds41-bf16-attention-fp32-heads.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json index 79f913f3ed03..ab295d649538 100644 --- a/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json +++ b/.lil/changes/vllm-ds41-bf16-attention-fp32-heads.json @@ -1,7 +1,7 @@ { "schema": "local-inference-release-change/v1", "id": "vllm-ds41-bf16-attention-fp32-heads", - "category": "change", + "category": "fix", "summary": "DeepSeek-V4.1 computes attention in BF16 and logits in FP32 like the reference", "models": [ "DeepSeek-V4.1-Flash" From ce4be0a112b9b687f93277d117f04e5424779a07 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 16:33:55 +0000 Subject: [PATCH 265/365] DS4.1: share decode page metadata across CUDA graphs Decode attention computes its sparse-MLA page metadata once per KV-cache group per forward (d11d4f35a8). Those per-forward tensors were registered as CUDA graph capture resources, both directly and through the scratch list that also holds them, so every captured graph pinned a private copy. At the full 1M context on 4x RTX PRO 6000 that is a (480, 4096) int32 page table per group: 116.6 MiB per graph, 3.8 GiB per GPU over the 33 DSpark verification graphs. The graph memory profiler extrapolates unsampled FULL graphs by token count and cannot see a constant per-graph cost, so it reserved 1.8 GiB for graphs that took 5.3 GiB. The oversized KV cache left about 1.2 GiB free and DSpark servers at --gpu-memory-utilization 0.95 intermittently ran out of memory when the sampler top-p or adaptive-verification kernels loaded after startup. Every replay rewrites the metadata before reading it, so it is an ordinary graph-pool activation. Retain only the borrowed workspace arena; the graph pool now reuses the metadata storage across graph shapes. Measured on ds41-flash (DSpark, TP4, RTX PRO 6000): CUDA graph memory 5.34 -> 1.90 GiB per GPU (estimate 1.83 GiB), KV cache unchanged at 9.65M tokens, free memory after startup 1.2 -> 4.1 GiB. Co-Authored-By: Claude Opus 5.5 --- ...llm-ds41-decode-metadata-graph-memory.json | 19 +++++++++++++++++++ vllm/models/deepseek_v4_1/attention.py | 12 +++++++----- 2 files changed, 26 insertions(+), 5 deletions(-) create mode 100644 .lil/changes/vllm-ds41-decode-metadata-graph-memory.json diff --git a/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json b/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json new file mode 100644 index 000000000000..a81a6df60869 --- /dev/null +++ b/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json @@ -0,0 +1,19 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-ds41-decode-metadata-graph-memory", + "category": "fix", + "summary": "DeepSeek-V4.1 CUDA graphs no longer pin a private copy of the decode page metadata", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "Every captured decode CUDA graph kept its own copy of the sparse-attention page metadata: about 117 MiB per graph at the full 1M context on 4x RTX PRO 6000, 3.8 GiB per GPU over the 33 DSpark verification graphs. The CUDA graph memory estimate did not see this per-graph cost, so the KV cache was sized for 1.8 GiB of graphs while they used 5.3 GiB, leaving about 1.2 GiB free; DSpark servers at --gpu-memory-utilization 0.95 could then run out of memory when the sampler or adaptive-verification kernels loaded after startup.", + "The metadata is still computed once per KV-cache group per forward, but graphs now share its storage through the CUDA graph pool. Graph memory drops from 5.3 to 1.9 GiB per GPU and matches the estimate; KV cache capacity is unchanged and about 4.1 GiB stays free after startup." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [], + "requires": [] +} diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 2b2ade4bd321..79827e6ff107 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -1495,13 +1495,16 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): # Reuse the shared arena only after indexing completes. Keep full-batch # metadata beside, not overlapping, the native attention scratch. metadata_specs = self._attention_workspace_specs[mode] - buffers = current_workspace_manager().get_simultaneous( + arena = current_workspace_manager().get_simultaneous( *metadata_specs, *((spec.shape, spec.dtype) for spec in state.scratch_plan.scratch_specs()), ) + buffers = arena metadata_count = len(metadata_specs) # Decode layers sharing cache groups map identical pages, so the first # layer of each group keeps its metadata for the rest of the forward. + # Every replay rewrites it before use, so it is an ordinary graph-pool + # activation: retaining it would pin one copy per captured graph. if mode == "decode": key = ( id(swa), @@ -1519,10 +1522,9 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): torch.empty(shape, dtype=dtype, device=q.device) for shape, dtype in metadata_specs ) - retain_cuda_graph_capture_resource(metadata_buffers) self._write_decode_metadata(metadata_buffers, swa, main, rows) entry = shared[key] = (swa, main, metadata_buffers) - buffers = [*entry[2], *buffers[metadata_count:]] + buffers = [*entry[2], *arena[metadata_count:]] else: self._write_decode_metadata(buffers, swa, main, rows) swa_indices, swa_lengths, top_lengths = buffers[:3] @@ -1541,8 +1543,8 @@ def forward_mqa(self, q, kv, positions, output, *, index_query=None): swa_lengths=swa_lengths[:rows], **kwargs, ) - # Keep scratch and metadata, not aliases of the transient query storage. - retain_cuda_graph_capture_resource(buffers) + # Keep the borrowed arena, not aliases of the transient query storage. + retain_cuda_graph_capture_resource(arena) mla.run( binding=binding, swa_k_cache=self.swa_cache_layer.kv_cache, From 5da2ca43459b6686265b9b1ee2a9e06a6964e64f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 17:10:46 +0000 Subject: [PATCH 266/365] test(ds41): capture decode attention with a fresh forward context The graph-replay test reused one forward context for every run, so decode metadata was created by the first eager run and the captured graph never allocated its own. Give every run a fresh context, as set_forward_context does, poison the metadata captured by the graph and require that capture resources release it while replays stay bit-exact. The unfixed ownership fails all nine decode cases; prefill cases are unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/v1/attention/test_b12x_v41_workspace.py | 37 ++++++++++++++----- 1 file changed, 27 insertions(+), 10 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index 66ea9c7bd3e6..b7062e55901b 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -833,15 +833,13 @@ def metadata(page): ), ) - context = SimpleNamespace( - cudagraph_runtime_mode=CUDAGraphMode.NONE, - attn_metadata={ - layer.swa_cache_layer.prefix: metadata(swa_page), - layer.prefix: metadata(main_page), - layer.indexer.k_cache.prefix: metadata(main_page), - }, - ) - monkeypatch.setattr(attention, "get_forward_context", lambda: context) + attn_metadata = { + layer.swa_cache_layer.prefix: metadata(swa_page), + layer.prefix: metadata(main_page), + layer.indexer.k_cache.prefix: metadata(main_page), + } + forward = SimpleNamespace(context=None) + monkeypatch.setattr(attention, "get_forward_context", lambda: forward.context) kv = torch.randn((length, 512), device=device, dtype=torch.bfloat16) for kind, page in (("swa", swa_page), ("indexed", main_page)): cache = torch.empty( @@ -913,6 +911,10 @@ def metadata(page): activation_refs: list[tuple[str, StorageWeakRef]] = [] def run(): + # Like set_forward_context, every forward starts from a fresh context. + forward.context = SimpleNamespace( + cudagraph_runtime_mode=CUDAGraphMode.NONE, attn_metadata=attn_metadata + ) query = q.clone() index_query = packed.clone(), scales.clone(), weights.clone() activation_refs.extend( @@ -956,7 +958,22 @@ def run(): torch.cuda.graph(graph, stream=stream), ): run() - # Graph resources own scratch, not the per-layer query activations. + + # Graph resources own scratch, not the per-layer query activations or the + # per-forward decode metadata: each replay rewrites it, so graphs of other + # shapes share its pool storage instead of pinning one copy per graph. + def poison_decode_metadata(): + shared = getattr(forward.context, "_ds41_decode_metadata", {}) + tensors = [tensor for entry in shared.values() for tensor in entry[2]] + for tensor in tensors: + tensor.fill_(-1) + activation_refs.append( + ("decode_metadata", StorageWeakRef(tensor.untyped_storage())) + ) + return len(tensors) + + assert bool(poison_decode_metadata()) == is_decode + forward.context = None gc.collect() retained = [name for name, reference in activation_refs if not reference.expired()] assert not retained, f"Capture resources retained caller activations: {retained}" From 5e2fdd68f416a904194cb463e415a6c2b383ecf0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 17:11:33 +0000 Subject: [PATCH 267/365] docs(release): link the DS4.1 graph-memory fragment to #918 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-ds41-decode-metadata-graph-memory.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json b/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json index a81a6df60869..83f3dc31fdf6 100644 --- a/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json +++ b/.lil/changes/vllm-ds41-decode-metadata-graph-memory.json @@ -14,6 +14,8 @@ "authors": [ "voipmonitor" ], - "pull_requests": [], + "pull_requests": [ + 918 + ], "requires": [] } From 3b294c22a8c378891fd214c5df3f8deadabde4fb Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 17:14:36 +0000 Subject: [PATCH 268/365] test(ds41): refresh V4.1 workspace tests for BF16 attention and L2 prefetch Two stale expectations failed seven tests in test_b12x_v41_workspace.py: - #911 pins V4.1 sparse attention to BF16 arithmetic, but the CED compact attention oracle still fed the kernel's FP8 view of the cache (E4M3 codes re-canonicalized to per-64 power-of-two scales) through FP8 probability/value products. The BF16 kernel dequantizes the native records directly, so the oracle was off by about 2% RMS. Unless VLLM_DS41_ATTENTION_COMPUTE=auto leaves the choice to the tuner, compare against exactly dequantized BF16 K/V with per-64-key BF16 probabilities; the FP8 model remains for auto. - d11d4f35a8 added L2 prefetch windows to the output projection; the test layers never set them, so four output-projection tests raised AttributeError before reaching the kernels. Co-Authored-By: Claude Opus 5.5 --- tests/v1/attention/test_b12x_v41_workspace.py | 64 ++++++++++++------- 1 file changed, 42 insertions(+), 22 deletions(-) diff --git a/tests/v1/attention/test_b12x_v41_workspace.py b/tests/v1/attention/test_b12x_v41_workspace.py index b7062e55901b..9310969c0321 100644 --- a/tests/v1/attention/test_b12x_v41_workspace.py +++ b/tests/v1/attention/test_b12x_v41_workspace.py @@ -11,6 +11,7 @@ import torch from torch.multiprocessing.reductions import StorageWeakRef +from vllm import envs from vllm.config import CUDAGraphMode pytestmark = pytest.mark.skipif( @@ -124,6 +125,7 @@ def _layer(attention, layer_id=0, swa_page=32): dim=-1, ) ) + layer._l2pf_wo = layer._l2pf_ffn = None layer._ready = False return layer @@ -701,12 +703,13 @@ def test_mhc_fixed_capacity_buckets_preserve_decode_policy( assert bool(torch.isfinite(y).all()) -def _v41_fp8_operands(cache, kind): - """Decode native cache rows into E4M3 operands and per-64 scales. +def _v41_cache_operands(cache, kind, *, fp8): + """Decode native cache rows into attention operands and per-64 scales. - This independent operand reference follows B12X's - tests/_reference/v41_fp8.py. SWA pairs its E8M0 scales; indexed NVFP4 - uses a power-of-two scale that bounds four adjacent scale groups. + BF16 internals dequantize the native records exactly (unit scales). FP8 + internals consume E4M3 operands; this independent reference follows + B12X's tests/_reference/v41_fp8.py. SWA pairs its E8M0 scales; indexed + NVFP4 uses a power-of-two scale that bounds four adjacent scale groups. """ if kind == "swa": rows = cache.reshape(-1, 528) @@ -727,17 +730,21 @@ def _v41_fp8_operands(cache, kind): values = lut[codes] * original.repeat_interleave(16, dim=1) bound = original.reshape(-1, 8, 4).amax(-1) * (6 / 448) exponent = torch.ceil(torch.log2(bound.clamp_min(2.0**-126))) + if not fp8: + return values.float(), torch.ones_like(exponent, dtype=torch.float32) scales = torch.exp2(exponent) codes = (values / scales.repeat_interleave(64, dim=1)).float() return codes.to(torch.float8_e4m3fn).float(), scales.float() -def _v41_fp8_prefill_reference(q, values, scales, valid, sink): - """Model BF16 Q/K and per-64-key FP8 probability/value products. +def _v41_prefill_reference(q, values, scales, valid, sink, *, fp8): + """Model BF16 Q/K and per-64-key probability/value products. - Prefill computes Q/K with BF16 operands. Each output group independently - quantizes scaled probabilities to E4M3 before its product with FP8 V; - partial numerators remain FP32 until the final LSE merge. + Prefill computes Q/K with BF16 operands. With FP8 internals, each output + group independently quantizes scaled probabilities to E4M3 before its + product with FP8 V; BF16 internals round probabilities to BF16 before the + product with dequantized BF16 V. Partial numerators remain FP32 until the + final LSE merge. """ keys = (values * scales.repeat_interleave(64, dim=-1)).bfloat16().float() logits = torch.einsum("rhd,rkd->rhk", q.float(), keys) * 512**-0.5 @@ -749,14 +756,19 @@ def _v41_fp8_prefill_reference(q, values, scales, valid, sink): p = torch.exp(local - maximum[:, :, None]) p = torch.where(valid[:, None, first : first + 64], p, 0) denominator = p.sum(-1) - groups = [] - for group in range(8): - weighted = p * scales[:, None, first : first + 64, group] - pscale = weighted.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 - pq = (weighted / pscale).to(torch.float8_e4m3fn).float() - vq = values[:, first : first + 64, group * 64 : (group + 1) * 64] - groups.append(torch.einsum("rhk,rkd->rhd", pq, vq) * pscale) - partials.append(torch.cat(groups, -1) / denominator.clamp_min(1e-30)[..., None]) + if fp8: + groups = [] + for group in range(8): + weighted = p * scales[:, None, first : first + 64, group] + pscale = weighted.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 + pq = (weighted / pscale).to(torch.float8_e4m3fn).float() + vq = values[:, first : first + 64, group * 64 : (group + 1) * 64] + groups.append(torch.einsum("rhk,rkd->rhd", pq, vq) * pscale) + numerator = torch.cat(groups, -1) + else: + pb = p.bfloat16().float() + numerator = torch.einsum("rhk,rkd->rhd", pb, keys[:, first : first + 64]) + partials.append(numerator / denominator.clamp_min(1e-30)[..., None]) lses.append( torch.where(denominator > 0, maximum + denominator.log(), -torch.inf) ) @@ -1257,6 +1269,7 @@ def run_inv_rope(*args, **kwargs): ) layer._wo_projection_weights = None layer._wo_plans = {3: plan} + layer._l2pf_wo = layer._l2pf_ffn = None layer.setup_wo_projection() output = layer._o_proj( @@ -1659,6 +1672,8 @@ def test_ced_compact_attention_bounded_oracle_and_frozen_replay( layer.config.cache_config.block_size = main_page layer.is_ced_decoder = True layer._prepare(device) + # Only the tuner's own choice ("auto") may select FP8 internals. + fp8 = envs.VLLM_DS41_ATTENTION_COMPUTE == "auto" rows, length, boundary = 128, 256, 128 positions = torch.arange(boundary, length, device=device, dtype=torch.int64) reqs = torch.zeros(rows, device=device, dtype=torch.int32) @@ -1719,7 +1734,7 @@ def metadata(page): layer.swa_cache_layer.kv_cache = cache swa_values, swa_scales = ( part[page : page + length].clone() - for part in _v41_fp8_operands(cache, kind) + for part in _v41_cache_operands(cache, kind, fp8=fp8) ) # NaN FP8 payloads in every old SWA row: reading below the # replay boundary cannot accidentally look like a valid zero page. @@ -1728,7 +1743,7 @@ def metadata(page): layer.kv_cache = cache main_values, main_scales = ( part[page : page + length].clone() - for part in _v41_fp8_operands(cache, kind) + for part in _v41_cache_operands(cache, kind, fp8=fp8) ) layer.indexer.k_cache.kv_cache = torch.zeros( ( @@ -1795,8 +1810,13 @@ def oracle(live): ), dim=1, ).masked_fill(~valid[..., None], 0) - return _v41_fp8_prefill_reference( - q[:live], values, scales, valid, layer.attn_sink + return _v41_prefill_reference( + q[:live], + values, + scales, + valid, + layer.attn_sink, + fp8=fp8, ) workload = B12xWorkload( From 17dae2eb414c6073f8d43f393a484d0ec76f074b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 20:50:13 +0000 Subject: [PATCH 269/365] docs(release): add the missing fragment for test-only #919 The community publisher rejects any vLLM source change without a new .lil/changes fragment, so the beta image after #919 was not published. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-919.json | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 .lil/changes/vllm-919.json diff --git a/.lil/changes/vllm-919.json b/.lil/changes/vllm-919.json new file mode 100644 index 000000000000..7f866642101a --- /dev/null +++ b/.lil/changes/vllm-919.json @@ -0,0 +1,13 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-919", + "category": "internal", + "summary": "DeepSeek-V4.1 workspace tests follow the BF16 attention default and L2 prefetch windows", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "No user action required; test-only change.", + "details": [ + "The CED compact attention oracle compares BF16 attention against exactly dequantized BF16 K/V and keeps its FP8 model for VLLM_DS41_ATTENTION_COMPUTE=auto; hand-built output-projection test layers set the L2 prefetch windows. The V4.1 workspace test file passes in full." + ], + "pull_requests": [919], + "requires": [] +} From 190268a20168964a7875621a079160c1b281ccbf Mon Sep 17 00:00:00 2001 From: lazymio Date: Fri, 18 Sep 2026 13:11:24 +0800 Subject: [PATCH 270/365] [Frontend] Fix the parsing of missing `string=` in DeepSeek V4 (#56271) Co-authored-by: Claude Opus 5 Co-authored-by: Bugen Zhao Signed-off-by: lazymio (cherry picked from commit 39e33db7f3b10ebd3aca6a0008f2990978f118e7) --- .../src/tool/deepseek_dsml/deepseek_v32.rs | 28 ++++++++ .../src/tool/deepseek_dsml/deepseek_v4.rs | 23 ++++++ .../tool/deepseek_dsml/deepseek_v41/tests.rs | 22 ++++++ rust/src/parser/src/tool/deepseek_dsml/mod.rs | 11 +-- tests/parser/engine/test_deepseek_v4.py | 70 +++++++++++++++++++ tests/parser/engine/test_deepseek_v41.py | 2 + vllm/parser/deepseek_v4.py | 32 +++++---- vllm/parser/deepseek_v41.py | 15 +--- 8 files changed, 174 insertions(+), 29 deletions(-) diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs index 632eda600f77..9bbe39367b5e 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs @@ -193,6 +193,34 @@ mod tests { ); } + #[test] + fn deepseek_v32_parse_complete_keeps_parameter_without_string_attr() { + // The model sometimes omits `string="..."`; the parameter must still + // be kept and converted through the schema, like `string="false"`. + let mut parser = DeepSeekV32ToolParser::new(&test_tools()); + let output = parser + .parse_complete( + "<｜DSML｜function_calls>\n\ + <｜DSML｜invoke name=\"convert\">\n\ + <｜DSML｜parameter name=\"whole\">5.0\n\ + <｜DSML｜parameter name=\"flag\" >true\n\ + <｜DSML｜parameter name=\"payload\" string=\"true\">{\"nested\":true}\n\ + \n\ + ", + ) + .unwrap(); + + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ + "whole": 5.0, + "flag": true, + "payload": "{\"nested\":true}", + }) + ); + } + #[test] fn deepseek_v32_parse_complete_preserves_raw_closing_tag_text_in_parameter_value() { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs index b65df235ca20..032cfbe5e284 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs @@ -147,6 +147,29 @@ mod tests { ); } + #[test] + fn deepseek_v4_streaming_keeps_parameter_without_string_attr() { + let mut parser = DeepSeekV4ToolParser::new(&test_tools()); + let output = collect_stream( + &mut parser, + &[ + "<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n", + "<｜DSML｜parameter name=\"loc", + "ation\">Par", + "is\n", + "<｜DSML｜parameter name=\"date\" string=\"true\">tomorrow\n", + "\n", + ], + ); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "location": "Paris", "date": "tomorrow" }) + ); + } + #[test] fn tool_framing_preserves_body_whitespace_across_chunk_boundaries() { assert_tool_framing_preserves_body_whitespace::( diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs index d40f82d533cd..5a6758f5c935 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs @@ -71,6 +71,28 @@ fn spaced_dsml_streaming_preserves_shared_schema_coercion() { } } +#[test] +fn spaced_dsml_keeps_parameter_without_string_attr() { + let wire = concat!( + "<｜DSML｜ calls>\n", + "<｜DSML｜ invoke name=\"lookup\">\n", + "<｜DSML｜ parameter name=\"query\">value\n", + "<｜DSML｜ parameter name=\"limit\">2\n", + "\n", + "", + ); + let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap(); + + let output = collect_stream(parser.as_mut(), &[wire]); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "query": "value", "limit": 2 }) + ); +} + #[test] fn v4_tags_remain_text_in_v41_dialect() { let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap(); diff --git a/rust/src/parser/src/tool/deepseek_dsml/mod.rs b/rust/src/parser/src/tool/deepseek_dsml/mod.rs index 01a3fbe282ea..9c7dd203312d 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/mod.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/mod.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; -use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated}; +use winnow::combinator::{alt, delimited, eof, opt, preceded, repeat, seq, terminated}; use winnow::prelude::*; use winnow::stream::Partial; use winnow::token::{literal, rest, take_until}; @@ -293,13 +293,16 @@ fn parse_invoke_params(invoke_body: &str, tokens: DsmlTokens) -> ModalResult ModalResult { seq! {DsmlParameter { _: literal(tokens.parameter_start), _: ws1, name: name_attr.map(|name: &str| name.to_string()), - _: ws1, - is_string: string_attr.map(|value| value == "true"), + is_string: opt(preceded(ws1, string_attr)).map(|value| value == Some("true")), _: ws0, _: ">", value: take_until(0.., tokens.parameter_end).map(str::to_string), @@ -315,7 +318,7 @@ fn name_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> { /// Parse a string attribute. fn string_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> { - delimited("string=\"", alt(("true", "false")), "\"").parse_next(input) + delimited("string=\"", take_until(0.., "\""), "\"").parse_next(input) } /// Parse a DSML name attribute. diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 5eb369cee969..fcd15f9dce1f 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -212,6 +212,35 @@ def test_empty_args(self): result = json.loads(_dsml_arg_converter("", partial=False)) assert result == {} + def _bare(self, name: str, value: str) -> str: + return f'<｜DSML｜parameter name="{name}">{value}{_PARAM_CLOSE}' + + def test_missing_string_attribute_keeps_literal(self): + """The model sometimes omits ``string``; the parameter must survive.""" + raw = self._bare("name", "alpha") + assert json.loads(_dsml_arg_converter(raw, partial=False)) == {"name": "alpha"} + + def test_missing_string_attribute_parses_json(self): + raw = self._bare("count", "7") + self._bare("opts", '{"k": [1]}') + assert json.loads(_dsml_arg_converter(raw, partial=False)) == { + "count": 7, + "opts": {"k": [1]}, + } + + def test_missing_string_attribute_alongside_annotated(self): + raw = self._bare("a", "x") + _param("b", "true", "42") + assert json.loads(_dsml_arg_converter(raw, partial=False)) == { + "a": "x", + "b": "42", + } + + def test_missing_string_attribute_partial_value(self): + raw = _param("a", "true", "x") + '<｜DSML｜parameter name="n">12' + assert json.loads(_dsml_arg_converter(raw, partial=True)) == { + "a": "x", + "n": 12, + } + def test_invalid_json_fallback(self): raw = self._raw(("data", "false", "[broken")) result = json.loads(_dsml_arg_converter(raw, partial=False)) @@ -327,6 +356,47 @@ def test_streaming_split_next_parameter_tag_is_buffered( } +class TestMissingStringAttribute: + def test_non_streaming_keeps_parameter(self, mock_tokenizer, mock_request): + text = ( + f"{DSML_TOOL_START}" + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" + f'<｜DSML｜parameter name="location">Paris{_PARAM_CLOSE}\n' + f"{_param('date', 'true', 'tomorrow')}\n" + f"{DSML_INVOKE_END}{DSML_TOOL_END}" + ) + + result = DeepSeekV4Parser(mock_tokenizer).extract_tool_calls(text, mock_request) + + assert result.tools_called is True + assert json.loads(result.tool_calls[0].function.arguments) == { + "location": "Paris", + "date": "tomorrow", + } + + def test_streaming_keeps_parameter(self, mock_tokenizer, mock_request): + chunks = [ + DSML_TOOL_START, + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n", + '<｜DSML｜parameter name="loc', + 'ation">Par', + "is", + _PARAM_CLOSE, + f"\n{_param('date', 'true', 'tomorrow')}\n", + DSML_INVOKE_END, + DSML_TOOL_END, + ] + + results = simulate_tool_streaming( + DeepSeekV4Parser(mock_tokenizer), mock_request, chunks + ) + + assert json.loads(collect_tool_arguments(results)) == { + "location": "Paris", + "date": "tomorrow", + } + + # ── Bare absorption and duplicate absorption ───────── diff --git a/tests/parser/engine/test_deepseek_v41.py b/tests/parser/engine/test_deepseek_v41.py index cc7d6e0aedf4..d99ca6b2a17b 100644 --- a/tests/parser/engine/test_deepseek_v41.py +++ b/tests/parser/engine/test_deepseek_v41.py @@ -137,10 +137,12 @@ def test_python_argument_conversion_and_partial_values(): '<｜DSML｜ parameter name="object" string="false">' '{"a": [true, null]}' '<｜DSML｜ parameter name="bad" string="false">[broken' + '<｜DSML｜ parameter name="bare">7' '<｜DSML｜ parameter name="text" string="true"> a", ) -_ESCAPED_DSML = re.escape(_DSML) -_PARAM_RE = re.compile( - rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">' - rf"(.*?)" - rf"(?:|(?=<{_ESCAPED_DSML}parameter\s+name=))", - re.DOTALL, -) -_PARTIAL_PARAM_RE = re.compile( - rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">' - rf"(.*)$", - re.DOTALL, -) + +def _param_patterns( + param_start: str, param_close: str +) -> tuple[re.Pattern, re.Pattern]: + """Complete and trailing-partial parameter regexes for one DSML dialect. + + The ``string`` attribute is optional: the model sometimes omits it, and + dropping such a parameter hands the client a tool call with no arguments. + """ + start, close = re.escape(param_start), re.escape(param_close) + head = rf'{start}\s+name="([^"]+)"(?:\s+string="([^"]*)")?\s*>' + return ( + re.compile(rf"{head}(.*?)(?:{close}|(?={start}\s+name=))", re.DOTALL), + re.compile(rf"{head}(.*)$", re.DOTALL), + ) + + +_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE) def _dsml_arg_converter( @@ -80,6 +86,8 @@ def _dsml_arg_converter( params: dict[str, object] = {} last_end = 0 + # A missing ``string`` attribute is treated like ``string="false"``: try + # JSON, fall back to the literal text. for m in param_re.finditer(raw_args): name, is_str, value = m.group(1), m.group(2), m.group(3) if is_str == "true": diff --git a/vllm/parser/deepseek_v41.py b/vllm/parser/deepseek_v41.py index 668bdd4a6bed..84216005102a 100644 --- a/vllm/parser/deepseek_v41.py +++ b/vllm/parser/deepseek_v41.py @@ -5,11 +5,10 @@ import functools from dataclasses import replace -import regex as re - from vllm.parser.deepseek_v4 import ( DeepSeekV4Parser, _dsml_arg_converter, + _param_patterns, deepseek_v4_config, ) from vllm.parser.engine.parser_engine_config import ParserEngineConfig @@ -21,17 +20,7 @@ DSML_PARAM_START = "<｜DSML｜ parameter" DSML_PARAM_CLOSE = "" -_PARAM_RE = re.compile( - r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">' - r"(.*?)" - r"(?:|(?=<｜DSML｜ parameter\s+name=))", - re.DOTALL, -) -_PARTIAL_PARAM_RE = re.compile( - r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">' - r"(.*)$", - re.DOTALL, -) +_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE) @functools.cache From 574e0feb7f6adc9013b18a673ff9411bee17384d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 21:14:04 +0000 Subject: [PATCH 271/365] docs(release): attribute the DSML optional string-attribute import Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-upstream-56271.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-upstream-56271.json diff --git a/.lil/changes/vllm-upstream-56271.json b/.lil/changes/vllm-upstream-56271.json new file mode 100644 index 000000000000..bd2303dc4726 --- /dev/null +++ b/.lil/changes/vllm-upstream-56271.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-upstream-56271", + "category": "fix", + "summary": "DeepSeek-V4.1 and V4 tool calls keep parameters that omit the string attribute", + "models": ["DeepSeek-V4.1-Flash", "DeepSeek-V4-Flash"], + "compatibility": "No user action required.", + "details": [ + "When the model wrote a DSML tool-call parameter without its string=\"true|false\" attribute, the parser dropped that parameter, so the client received the tool call without it (for example {\"a\": \"x\"} instead of {\"a\": \"x\", \"n\": 12}). A missing attribute is now treated like string=\"false\": the value is parsed as JSON and otherwise kept as literal text, in streaming and non-streaming responses and in the Python and Rust parsers.", + "The change preserves the original author and code from vllm-project/vllm#56271." + ], + "pull_requests": [], + "authors": ["lazymio"], + "evidence": ["https://github.com/vllm-project/vllm/pull/56271"], + "requires": [] +} From 76554c15611f9d90ee6abed392823df5dd779f73 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 27 Sep 2026 21:14:26 +0000 Subject: [PATCH 272/365] docs(release): link the DSML import fragment to #921 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-upstream-56271.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-upstream-56271.json b/.lil/changes/vllm-upstream-56271.json index bd2303dc4726..a47a6091ab83 100644 --- a/.lil/changes/vllm-upstream-56271.json +++ b/.lil/changes/vllm-upstream-56271.json @@ -9,7 +9,7 @@ "When the model wrote a DSML tool-call parameter without its string=\"true|false\" attribute, the parser dropped that parameter, so the client received the tool call without it (for example {\"a\": \"x\"} instead of {\"a\": \"x\", \"n\": 12}). A missing attribute is now treated like string=\"false\": the value is parsed as JSON and otherwise kept as literal text, in streaming and non-streaming responses and in the Python and Rust parsers.", "The change preserves the original author and code from vllm-project/vllm#56271." ], - "pull_requests": [], + "pull_requests": [921], "authors": ["lazymio"], "evidence": ["https://github.com/vllm-project/vllm/pull/56271"], "requires": [] From dfeb39654a580ec3fb15449f3bbbe407534ecd37 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 03:41:46 -0400 Subject: [PATCH 273/365] fix(worker): map circular-buffer groups into their ring block A CircularBufferSpec group owns one block per request and uses it as a ring (DeepSeek-V4.1's compressor partial-state cache). The generic slot mapping indexed block_table[pos // block_size], so every position at or past the ring size mapped to padding and the ring stopped receiving state. Circular groups now map every position into their block at pos % block_size; padding rows and requests without a ring block still map to padding. Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/v1/worker/test_gpu_block_table.py | 38 +++++++++++++++++++++++++ vllm/v1/worker/gpu/block_table.py | 15 +++++++++- vllm/v1/worker/gpu/model_runner.py | 4 +++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 1aa1da0eafef..5d144b5daafc 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -378,3 +378,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() + + +def test_block_tables_circular_group_maps_positions_into_its_ring(): + """A circular-buffer group holds one block per request as a ring, so positions at + or past its block size still map into that block (position modulo the block size). + A request without a ring block keeps padding slots; positional groups are + unchanged.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16, 8], + max_num_reqs=2, + max_num_batched_tokens=64, + max_num_blocks_per_group=[4, 1], + device=device, + kernel_block_sizes=[16, 8], + slot_mapping_circular=[False, True], + ) + block_tables.append_block_ids( + req_index=0, new_block_ids=([3, 4, 5], [9]), overwrite=True + ) + block_tables.append_block_ids(req_index=1, new_block_ids=([6], []), overwrite=True) + block_tables.apply_staged_writes() + + positions = list(range(3, 40)) + [5, 6] + idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 37, 39], dtype=torch.int32, device=device) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + torch.tensor(positions, dtype=torch.int64, device=device), + num_tokens_padded=len(positions), + ) + torch.accelerator.synchronize() + + ring = [9 * 8 + p % 8 for p in positions[:37]] + assert slot_mappings[1].tolist() == ring + [-1, -1] + paged = [[3, 4, 5][p // 16] * 16 + p % 16 for p in positions[:37]] + assert slot_mappings[0].tolist() == paged + [6 * 16 + 5, 6 * 16 + 6] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 2c40aa28a459..5aa9374eae60 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -28,6 +28,7 @@ def __init__( cp_interleave: int = 1, slot_mapping_enabled: list[bool] | None = None, group_cp_sizes: list[int] | None = None, + slot_mapping_circular: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -50,6 +51,12 @@ def __init__( slot_mapping_enabled = [True] * self.num_kv_cache_groups assert len(slot_mapping_enabled) == self.num_kv_cache_groups self._slot_mapping_enabled = slot_mapping_enabled + # A circular group keeps one block per request as a ring: every position maps + # into that block at position modulo the block size. + if slot_mapping_circular is None: + slot_mapping_circular = [False] * self.num_kv_cache_groups + assert len(slot_mapping_circular) == self.num_kv_cache_groups + self._slot_mapping_circular = slot_mapping_circular self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -114,6 +121,9 @@ def init_block_table_layout_tensors(self) -> None: self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) + self.slot_mapping_circular = torch.tensor( + self._slot_mapping_circular, dtype=torch.bool, device=self.device + ) self.group_cp_sizes = torch.tensor( self.group_cp_sizes_list, dtype=torch.int32, device=self.device ) @@ -228,6 +238,7 @@ def compute_slot_mappings( self.block_sizes_tensor, self.kernel_block_sizes_tensor, self.slot_mapping_enabled, + self.slot_mapping_circular, self.group_cp_sizes, slot_mappings, slot_mappings.stride(0), @@ -308,6 +319,7 @@ def _compute_slot_mappings_kernel( block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] + slot_mapping_circular, # [num_kv_cache_groups] group_cp_sizes, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, @@ -339,6 +351,7 @@ def _compute_slot_mappings_kernel( kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) + circular = tl.load(slot_mapping_circular + group_id) group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) @@ -366,7 +379,7 @@ def _compute_slot_mappings_kernel( local_positions = virtual_block_indices * kv_block_size + local_offsets block_indices = tl.where( - mapping_enabled, local_positions // kernel_block_size, 0 + mapping_enabled & ~circular, local_positions // kernel_block_size, 0 ) block_offsets = local_positions % kernel_block_size valid_block = token_mask & (block_indices < num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4a478494054c..55a2dfa07462 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -68,6 +68,7 @@ from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -617,6 +618,7 @@ def initialize_kv_cache( block_sizes = [] max_num_blocks_per_group = [] slot_mapping_enabled = [] + slot_mapping_circular = [] group_cp_sizes = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec @@ -625,6 +627,7 @@ def initialize_kv_cache( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) slot_mapping_enabled.append(layer_spec.uses_slot_mapping) + slot_mapping_circular.append(isinstance(layer_spec, CircularBufferSpec)) group_cp_sizes.append( 1 if getattr(layer_spec, "dcp_replicated", False) else self.dcp_size ) @@ -707,6 +710,7 @@ def initialize_kv_cache( device=self.device, kernel_block_sizes=self.kernel_block_sizes, slot_mapping_enabled=slot_mapping_enabled, + slot_mapping_circular=slot_mapping_circular, cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, From 4d007438e3ee97fb96dd75a4b31304bf864897f5 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 03:42:16 -0400 Subject: [PATCH 274/365] docs(release): record DeepSeek-V4.1 compressor ring fix Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-926.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-926.json diff --git a/.lil/changes/vllm-926.json b/.lil/changes/vllm-926.json new file mode 100644 index 000000000000..5fdc9511ddda --- /dev/null +++ b/.lil/changes/vllm-926.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-926", + "category": "fix", + "summary": "Keep DeepSeek-V4.1 compressor state across decode steps past the ring size", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "Circular-buffer KV groups map every position into their one block, so the compressor state ring is written past its first 8–16 tokens and decode no longer drifts from the model as generation lengthens." + ], + "pull_requests": [ + 926 + ], + "requires": [] +} From 706428761ffb1c20fc8671f15416a57b3a73985b Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 03:41:46 -0400 Subject: [PATCH 275/365] fix(worker): map circular-buffer groups into their ring block A CircularBufferSpec group owns one block per request and uses it as a ring (DeepSeek-V4.1's compressor partial-state cache). The generic slot mapping indexed block_table[pos // block_size], so every position at or past the ring size mapped to padding and the ring stopped receiving state. Circular groups now map every position into their block at pos % block_size; padding rows and requests without a ring block still map to padding. Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit dfeb39654a580ec3fb15449f3bbbe407534ecd37) --- tests/v1/worker/test_gpu_block_table.py | 38 +++++++++++++++++++++++++ vllm/v1/worker/gpu/block_table.py | 15 +++++++++- vllm/v1/worker/gpu/model_runner.py | 4 +++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 1aa1da0eafef..5d144b5daafc 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -378,3 +378,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() + + +def test_block_tables_circular_group_maps_positions_into_its_ring(): + """A circular-buffer group holds one block per request as a ring, so positions at + or past its block size still map into that block (position modulo the block size). + A request without a ring block keeps padding slots; positional groups are + unchanged.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16, 8], + max_num_reqs=2, + max_num_batched_tokens=64, + max_num_blocks_per_group=[4, 1], + device=device, + kernel_block_sizes=[16, 8], + slot_mapping_circular=[False, True], + ) + block_tables.append_block_ids( + req_index=0, new_block_ids=([3, 4, 5], [9]), overwrite=True + ) + block_tables.append_block_ids(req_index=1, new_block_ids=([6], []), overwrite=True) + block_tables.apply_staged_writes() + + positions = list(range(3, 40)) + [5, 6] + idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 37, 39], dtype=torch.int32, device=device) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + torch.tensor(positions, dtype=torch.int64, device=device), + num_tokens_padded=len(positions), + ) + torch.accelerator.synchronize() + + ring = [9 * 8 + p % 8 for p in positions[:37]] + assert slot_mappings[1].tolist() == ring + [-1, -1] + paged = [[3, 4, 5][p // 16] * 16 + p % 16 for p in positions[:37]] + assert slot_mappings[0].tolist() == paged + [6 * 16 + 5, 6 * 16 + 6] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 2c40aa28a459..5aa9374eae60 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -28,6 +28,7 @@ def __init__( cp_interleave: int = 1, slot_mapping_enabled: list[bool] | None = None, group_cp_sizes: list[int] | None = None, + slot_mapping_circular: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -50,6 +51,12 @@ def __init__( slot_mapping_enabled = [True] * self.num_kv_cache_groups assert len(slot_mapping_enabled) == self.num_kv_cache_groups self._slot_mapping_enabled = slot_mapping_enabled + # A circular group keeps one block per request as a ring: every position maps + # into that block at position modulo the block size. + if slot_mapping_circular is None: + slot_mapping_circular = [False] * self.num_kv_cache_groups + assert len(slot_mapping_circular) == self.num_kv_cache_groups + self._slot_mapping_circular = slot_mapping_circular self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -114,6 +121,9 @@ def init_block_table_layout_tensors(self) -> None: self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) + self.slot_mapping_circular = torch.tensor( + self._slot_mapping_circular, dtype=torch.bool, device=self.device + ) self.group_cp_sizes = torch.tensor( self.group_cp_sizes_list, dtype=torch.int32, device=self.device ) @@ -228,6 +238,7 @@ def compute_slot_mappings( self.block_sizes_tensor, self.kernel_block_sizes_tensor, self.slot_mapping_enabled, + self.slot_mapping_circular, self.group_cp_sizes, slot_mappings, slot_mappings.stride(0), @@ -308,6 +319,7 @@ def _compute_slot_mappings_kernel( block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] + slot_mapping_circular, # [num_kv_cache_groups] group_cp_sizes, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, @@ -339,6 +351,7 @@ def _compute_slot_mappings_kernel( kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) + circular = tl.load(slot_mapping_circular + group_id) group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) @@ -366,7 +379,7 @@ def _compute_slot_mappings_kernel( local_positions = virtual_block_indices * kv_block_size + local_offsets block_indices = tl.where( - mapping_enabled, local_positions // kernel_block_size, 0 + mapping_enabled & ~circular, local_positions // kernel_block_size, 0 ) block_offsets = local_positions % kernel_block_size valid_block = token_mask & (block_indices < num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4a478494054c..55a2dfa07462 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -68,6 +68,7 @@ from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -617,6 +618,7 @@ def initialize_kv_cache( block_sizes = [] max_num_blocks_per_group = [] slot_mapping_enabled = [] + slot_mapping_circular = [] group_cp_sizes = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec @@ -625,6 +627,7 @@ def initialize_kv_cache( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) slot_mapping_enabled.append(layer_spec.uses_slot_mapping) + slot_mapping_circular.append(isinstance(layer_spec, CircularBufferSpec)) group_cp_sizes.append( 1 if getattr(layer_spec, "dcp_replicated", False) else self.dcp_size ) @@ -707,6 +710,7 @@ def initialize_kv_cache( device=self.device, kernel_block_sizes=self.kernel_block_sizes, slot_mapping_enabled=slot_mapping_enabled, + slot_mapping_circular=slot_mapping_circular, cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, From 8be242ab7275436d8f54740f8e721493b274c787 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 03:42:16 -0400 Subject: [PATCH 276/365] docs(release): record DeepSeek-V4.1 compressor ring fix Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 4d007438e3ee97fb96dd75a4b31304bf864897f5) --- .lil/changes/vllm-926.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-926.json diff --git a/.lil/changes/vllm-926.json b/.lil/changes/vllm-926.json new file mode 100644 index 000000000000..5fdc9511ddda --- /dev/null +++ b/.lil/changes/vllm-926.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-926", + "category": "fix", + "summary": "Keep DeepSeek-V4.1 compressor state across decode steps past the ring size", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "Circular-buffer KV groups map every position into their one block, so the compressor state ring is written past its first 8–16 tokens and decode no longer drifts from the model as generation lengthens." + ], + "pull_requests": [ + 926 + ], + "requires": [] +} From 5eec8ea30a7d4c4a68f0fbdd7ab1a28f43c9b982 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 05:03:53 -0400 Subject: [PATCH 277/365] fix(kv-cache): give circular-buffer groups a worker slot mapping DeepSeek-V4.1's compressor writes a ring row only when the worker gave that row a valid slot, which is also how it keeps padding out of the ring. CircularBufferSpec.uses_slot_mapping returned False, so the runner disabled slot mapping for the ring group and every row was padding: the ring was never written and decode pooled with stale partial state. Circular groups now take a slot mapping into their ring block (position modulo the block size); padding and requests without a ring block still map to the pad slot. The runner and the tests take the per-group flags from one helper, so the tests cover the settings the runner produces. Qwen3.8-Flash-Next builds its own ring slots and does not read this mapping. Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/models/test_deepseek_v4_1_compressor.py | 76 +++++++++++++++++++ tests/v1/worker/test_gpu_block_table.py | 39 +++++++--- tests/v1/worker/test_gpu_model_runner_v2.py | 4 +- vllm/v1/kv_cache_interface.py | 12 +-- vllm/v1/worker/gpu/block_table.py | 7 ++ vllm/v1/worker/gpu/model_runner.py | 8 +- 6 files changed, 123 insertions(+), 23 deletions(-) diff --git a/tests/models/test_deepseek_v4_1_compressor.py b/tests/models/test_deepseek_v4_1_compressor.py index 9c029b73cdba..ddb9ad232aa1 100644 --- a/tests/models/test_deepseek_v4_1_compressor.py +++ b/tests/models/test_deepseek_v4_1_compressor.py @@ -24,6 +24,82 @@ def test_compressor_cache_rejects_duplicate_prefix_without_replacing_owner(): assert config.compilation_config.static_forward_context[other.prefix] is other +def test_compressor_ring_writes_follow_the_worker_slot_mapping(): + """Ring metadata emits a write only for rows with a valid worker slot. With + the runner's per-spec settings, a verify step past the ring capacity writes + every real row into its request's ring block, and graph padding writes + nothing.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("DeepSeek-V4.1 metadata kernels require SM12x") + + from vllm.models.deepseek_v4_1.compressor import CompressorStateCache + from vllm.models.deepseek_v4_1.sparse_mla import DeepseekV41B12xMetadataBuilder + from vllm.v1.kv_cache_interface import MLAAttentionSpec + from vllm.v1.worker.gpu.block_table import BlockTables, slot_mapping_mode + + device = torch.device("cuda", torch.cuda.current_device()) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=64, max_num_seqs=4), + compilation_config=SimpleNamespace(static_forward_context={}), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + speculative_config=SimpleNamespace( + num_speculative_tokens=7, parallel_drafting=False + ), + ) + ring = CompressorStateCache(config, "layer.state_cache").get_kv_cache_spec(config) + main = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=288, + dtype=torch.uint8, + tokens_per_state=2, + ) + specs = (main, ring) + enabled, circular = zip(*(slot_mapping_mode(spec) for spec in specs)) + tables = BlockTables( + block_sizes=[spec.block_size for spec in specs], + max_num_reqs=4, + max_num_batched_tokens=64, + max_num_blocks_per_group=[2, 1], + device=device, + kernel_block_sizes=[spec.block_size for spec in specs], + slot_mapping_enabled=list(enabled), + slot_mapping_circular=list(circular), + ) + tables.append_block_ids(req_index=0, new_block_ids=([1], [5]), overwrite=True) + tables.append_block_ids(req_index=1, new_block_ids=([2], [6]), overwrite=True) + tables.apply_staged_writes() + + # A verify step (bonus token and seven drafts) at positions 40-47, a single + # token at position 3, then two tokens of graph padding. + positions = list(range(40, 48)) + [3, 0, 0] + idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 8, 9], dtype=torch.int32, device=device) + slot_mappings = tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + torch.tensor(positions, dtype=torch.int64, device=device), + num_tokens_padded=len(positions), + ) + cm = SimpleNamespace( + query_start_loc=query_start_loc, + seq_lens=torch.tensor([48, 4], dtype=torch.int32, device=device), + block_table_tensor=tables.gather_block_tables(idx_mapping, 2)[1], + slot_mapping=slot_mappings[1], + num_reqs=2, + num_actual_tokens=len(positions), + max_query_len=8, + max_seq_len=48, + ) + builder = DeepseekV41B12xMetadataBuilder(ring, ["layer"], config, device) + state = builder.build(0, cm) + torch.accelerator.synchronize() + + size = ring.block_size + expected = [5 * size + p % size for p in range(40, 48)] + [6 * size + 3, -1, -1] + assert state.slot_mapping[: len(positions)].tolist() == expected + + @pytest.fixture def native_compressor_runners(monkeypatch): if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5d144b5daafc..bfa39dbfa8d6 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,7 +5,8 @@ import torch from vllm.platforms import current_platform -from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.kv_cache_interface import CircularBufferSpec, FullAttentionSpec +from vllm.v1.worker.gpu.block_table import BlockTables, slot_mapping_mode pytestmark = pytest.mark.skipif( not current_platform.is_cuda(), @@ -381,19 +382,33 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): def test_block_tables_circular_group_maps_positions_into_its_ring(): - """A circular-buffer group holds one block per request as a ring, so positions at - or past its block size still map into that block (position modulo the block size). - A request without a ring block keeps padding slots; positional groups are - unchanged.""" + """A circular-buffer group holds one block per request as a ring. With the + runner's per-spec settings, every position maps into that block at position + modulo the block size; a request without a ring block and padding tokens map + to padding, and positional groups are unchanged.""" device = torch.device("cuda") + specs = ( + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.bfloat16 + ), + CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=1024, + head_size_v=0, + dtype=torch.float32, + ), + ) + enabled, circular = zip(*(slot_mapping_mode(spec) for spec in specs)) block_tables = BlockTables( - block_sizes=[16, 8], + block_sizes=[spec.block_size for spec in specs], max_num_reqs=2, max_num_batched_tokens=64, max_num_blocks_per_group=[4, 1], device=device, - kernel_block_sizes=[16, 8], - slot_mapping_circular=[False, True], + kernel_block_sizes=[spec.block_size for spec in specs], + slot_mapping_enabled=list(enabled), + slot_mapping_circular=list(circular), ) block_tables.append_block_ids( req_index=0, new_block_ids=([3, 4, 5], [9]), overwrite=True @@ -407,12 +422,12 @@ def test_block_tables_circular_group_maps_positions_into_its_ring(): slot_mappings = block_tables.compute_slot_mappings( idx_mapping, query_start_loc, - torch.tensor(positions, dtype=torch.int64, device=device), - num_tokens_padded=len(positions), + torch.tensor(positions + [0, 0], dtype=torch.int64, device=device), + num_tokens_padded=len(positions) + 2, ) torch.accelerator.synchronize() ring = [9 * 8 + p % 8 for p in positions[:37]] - assert slot_mappings[1].tolist() == ring + [-1, -1] + assert slot_mappings[1].tolist() == ring + [-1] * 4 paged = [[3, 4, 5][p // 16] * 16 + p % 16 for p in positions[:37]] - assert slot_mappings[0].tolist() == paged + [6 * 16 + 5, 6 * 16 + 6] + assert slot_mappings[0].tolist() == paged + [6 * 16 + 5, 6 * 16 + 6, -1, -1] diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 748a2e68bfb4..caa45db41720 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -162,7 +162,9 @@ def capture_block_tables(**kwargs): runner.initialize_kv_cache(kv_cache_config) assert captured["max_num_blocks_per_group"] == [1, 1] - assert captured["slot_mapping_enabled"] == [False, True] + # The worker maps the ring into its block; QSA still builds its own slots. + assert captured["slot_mapping_enabled"] == [True, True] + assert captured["slot_mapping_circular"] == [True, False] @pytest.mark.parametrize( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 89a13469a3eb..6cbd934c1d08 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -277,8 +277,8 @@ def has_layer_views(self) -> bool: def uses_slot_mapping(self) -> bool: """Whether the worker computes a per-token slot mapping for this spec. - Specs that address their pages themselves (raw storage, ring buffers) - take no slot mapping row. + Specs that address their pages themselves (raw storage) take no slot + mapping row. """ return True @@ -925,6 +925,10 @@ class CircularBufferSpec(AttentionSpec): by the speculative lookahead: a speculative step stores all of its rows, drafts included, before acceptance is known, while the next step still reads the open group's committed keys from the ring. + + The worker maps every real row into that block at its position modulo + ``block_size`` and padding rows to the pad slot; owners write only rows + with a valid slot. """ @property @@ -951,10 +955,6 @@ def is_uniform_with_collection( def prefix_cacheable(self) -> bool: return False - @property - def uses_slot_mapping(self) -> bool: - return False - @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 5aa9374eae60..3937f916846f 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -6,6 +6,7 @@ from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import CircularBufferSpec, KVCacheSpec from vllm.v1.worker.gpu.buffer_utils import ( FusedStagedWriter, StagedWriteTensor, @@ -14,6 +15,12 @@ ) +def slot_mapping_mode(layer_spec: KVCacheSpec) -> tuple[bool, bool]: + """Return whether the worker computes slot mappings for a group with this + layer spec, and whether the group is a per-request ring.""" + return layer_spec.uses_slot_mapping, isinstance(layer_spec, CircularBufferSpec) + + class BlockTables: def __init__( self, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 55a2dfa07462..69c82ce7ee97 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -68,7 +68,6 @@ from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( - CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -101,7 +100,7 @@ init_attn_backend, init_kv_cache, ) -from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.block_table import BlockTables, slot_mapping_mode from vllm.v1.worker.gpu.boundary_checkpoint import BoundaryCheckpointState from vllm.v1.worker.gpu.buffer_utils import ( set_default_max_concurrency, @@ -626,8 +625,9 @@ def initialize_kv_cache( layer_spec = ( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) - slot_mapping_enabled.append(layer_spec.uses_slot_mapping) - slot_mapping_circular.append(isinstance(layer_spec, CircularBufferSpec)) + mapping_enabled, mapping_circular = slot_mapping_mode(layer_spec) + slot_mapping_enabled.append(mapping_enabled) + slot_mapping_circular.append(mapping_circular) group_cp_sizes.append( 1 if getattr(layer_spec, "dcp_replicated", False) else self.dcp_size ) From 7b935cb78a3f1f256b04006e82504a651685f72d Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 28 Sep 2026 05:03:53 -0400 Subject: [PATCH 278/365] docs(release): record DeepSeek-V4.1 compressor ring writes Signed-off-by: MadeBy561 Co-Authored-By: Claude Opus 5.5 (1M context) --- .lil/changes/vllm-926-ring-writes.json | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 .lil/changes/vllm-926-ring-writes.json diff --git a/.lil/changes/vllm-926-ring-writes.json b/.lil/changes/vllm-926-ring-writes.json new file mode 100644 index 000000000000..e14e0a8b80c0 --- /dev/null +++ b/.lil/changes/vllm-926-ring-writes.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-926-ring-writes", + "category": "fix", + "summary": "Write the DeepSeek-V4.1 compressor state ring during decode", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "Circular-buffer KV groups take a worker slot mapping into their ring block, so DeepSeek-V4.1 writes its compressor partial-state ring for every real row and decode no longer drifts from the model as generation lengthens.", + "Corrects vllm-926: mapping ring positions alone left circular groups without a slot mapping, so the ring was still never written." + ], + "pull_requests": [ + 926 + ], + "requires": [] +} From b98ece7308a6df04f0a5d8c4500a6cf983ed9adf Mon Sep 17 00:00:00 2001 From: tobymao Date: Sun, 27 Sep 2026 08:46:46 -0700 Subject: [PATCH 279/365] fix(executor): gather every rank's reply without waiting in rank order collective_rpc read the response queues in rank order, each with a blocking dequeue. A failure one rank had already reported sat unread while an earlier rank waited forever, for example in a collective the failed rank would never join. On 4x GB10 at startup, TP1's CUDA graph capture failed and it replied with the error; TP0 hung in the next all-reduce, and the engine waited on TP0 with TP1's failure in its queue until an operator noticed. When every rank replies, the gather now polls each pending queue without blocking and, only when none has replied, blocks on the oldest for 50 ms, so a failure anywhere ends the wait at once and names its rank. It keeps the RPC deadline and logs the same long-wait message on the same interval as a blocking read. A single-rank reply (the per-step output rank) keeps the one blocking read. Signed-off-by: tobymao --- .../test_multiproc_gather_responses.py | 63 +++++++++++ vllm/v1/executor/multiproc_executor.py | 102 +++++++++++++++--- 2 files changed, 149 insertions(+), 16 deletions(-) create mode 100644 tests/v1/executor/test_multiproc_gather_responses.py diff --git a/tests/v1/executor/test_multiproc_gather_responses.py b/tests/v1/executor/test_multiproc_gather_responses.py new file mode 100644 index 000000000000..20459566cf85 --- /dev/null +++ b/tests/v1/executor/test_multiproc_gather_responses.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The executor's gather of every rank's reply does not wait in rank order. + +A rank that never replies (for example, stuck in a collective another rank +failed out of) must not hide a failure another rank has already reported. +""" + +import time + +import pytest + +from vllm.v1.executor.multiproc_executor import WorkerProc, _gather_responses + +SUCCESS = WorkerProc.ResponseStatus.SUCCESS +FAILURE = WorkerProc.ResponseStatus.FAILURE + + +class _Queue: + """Replies once ``ready_after`` dequeue calls have passed; never if None.""" + + def __init__(self, reply=None, ready_after: int = 0) -> None: + self.reply = reply + self.ready_after = ready_after + self.calls = 0 + + def dequeue(self, timeout=None): + self.calls += 1 + if self.reply is not None and self.calls > self.ready_after: + reply, self.reply = self.reply, None + return reply + if timeout: + time.sleep(min(timeout, 0.01)) + raise TimeoutError + + +def test_a_later_rank_failure_surfaces_while_rank_zero_hangs(): + queues = [_Queue(), _Queue((FAILURE, "capture invalidated")), _Queue()] + started = time.monotonic() + with pytest.raises(RuntimeError, match="Worker 1 failed.*capture invalidated"): + _gather_responses(queues, range(3), None, "compile_or_warm_up_model") + assert time.monotonic() - started < 1.0 + + +def test_replies_arriving_out_of_order_come_back_in_rank_order(): + queues = [ + _Queue((SUCCESS, "r0"), ready_after=5), + _Queue((SUCCESS, "r1")), + _Queue((SUCCESS, "r2"), ready_after=2), + ] + assert _gather_responses(queues, range(3), None, "m") == ["r0", "r1", "r2"] + + +def test_the_deadline_still_applies(): + queues = [_Queue(), _Queue((SUCCESS, "r1"))] + with pytest.raises(TimeoutError, match="RPC call to m timed out"): + _gather_responses(queues, range(2), time.monotonic() + 0.2, "m") + + +def test_a_single_reply_names_its_rank(): + with pytest.raises(RuntimeError, match="Worker 3 failed"): + _gather_responses([_Queue((FAILURE, "boom"))], (3,), None, "m") + assert _gather_responses([_Queue((SUCCESS, "ok"))], (3,), None, "m") == ["ok"] diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 598121f9e6af..68562140d013 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -29,7 +29,11 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import destroy_distributed_environment, destroy_model_parallel -from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue +from vllm.distributed.device_communicators.shm_broadcast import ( + LONG_WAIT_TIME_LOG_MSG, + Handle, + MessageQueue, +) from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.parallel_state import ( @@ -76,6 +80,84 @@ logger = init_logger(__name__) +# How long a gather of every rank's reply blocks on one rank before checking +# the others for a reply (a failure among them ends the wait). +RESPONSE_POLL_SECONDS = 0.05 + + +def _check_response(rank: int, status: Any, result: Any) -> Any: + if status != WorkerProc.ResponseStatus.SUCCESS: + raise RuntimeError( + f"Worker {rank} failed with error '{result}', please check the" + " stack trace above for the root cause" + ) + return result + + +def _gather_responses( + response_mqs: Sequence[MessageQueue], + ranks: Sequence[int], + deadline: float | None, + method: str, +) -> list[Any]: + """Collect one reply per queue, failing as soon as any rank reports one. + + Reading the queues in rank order let a failure on one rank sit unread + while an earlier rank waited forever, for example in a collective the + failed rank would never join: TP1's CUDA graph capture fails, TP0 hangs + in the next all-reduce, and the engine waits on TP0 with TP1's error in + its queue. Every pending queue is polled without blocking, and only when + none has replied does the gather block, briefly, on the oldest one. + """ + + def remaining() -> float | None: + if deadline is None: + return None + left = deadline - time.monotonic() + if left <= 0: + raise TimeoutError(f"RPC call to {method} timed out.") + return left + + if len(response_mqs) == 1: + try: + status, result = response_mqs[0].dequeue(timeout=remaining()) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e + return [_check_response(ranks[0], status, result)] + + responses: list[Any] = [None] * len(response_mqs) + pending = list(range(len(response_mqs))) + # A blocking read logs the long-wait message every interval; the short + # polls here would not, and operators watch for it, so log it the same way. + started, warnings = time.monotonic(), 0 + interval = envs.VLLM_RINGBUFFER_WARNING_INTERVAL + while pending: + if time.monotonic() - started >= interval * (warnings + 1): + warnings += 1 + logger.info(LONG_WAIT_TIME_LOG_MSG, interval) + for index in list(pending): + try: + status, result = response_mqs[index].dequeue(timeout=0) + except TimeoutError: + continue + responses[index] = _check_response(ranks[index], status, result) + pending.remove(index) + if not pending: + break + left = remaining() + wait = ( + RESPONSE_POLL_SECONDS if left is None else min(RESPONSE_POLL_SECONDS, left) + ) + index = pending[0] + try: + status, result = response_mqs[index].dequeue(timeout=wait) + except TimeoutError: + continue + responses[index] = _check_response(ranks[index], status, result) + pending.remove(index) + return responses + + class FutureWrapper(Future): def __init__( self, @@ -433,25 +515,13 @@ def _aggregate(outputs: Any) -> Any: self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, output_rank)) response_mqs: Sequence[MessageQueue] = self.response_mqs + ranks: Sequence[int] = range(len(response_mqs)) if output_rank is not None: response_mqs = (response_mqs[output_rank],) + ranks = (output_rank,) def get_response(): - responses = [] - for mq in response_mqs: - dequeue_timeout = ( - None if deadline is None else max(0.0, deadline - time.monotonic()) - ) - try: - status, result = mq.dequeue(timeout=dequeue_timeout) - except TimeoutError as e: - raise TimeoutError(f"RPC call to {method} timed out.") from e - if status != WorkerProc.ResponseStatus.SUCCESS: - raise RuntimeError( - f"Worker failed with error '{result}', please check the" - " stack trace above for the root cause" - ) - responses.append(result) + responses = _gather_responses(response_mqs, ranks, deadline, str(method)) return responses[0] if output_rank is not None else responses future = FutureWrapper( From 34d0206fee6be5d5abb9c987d90f165b5619aa6b Mon Sep 17 00:00:00 2001 From: tobymao Date: Sun, 27 Sep 2026 09:25:18 -0700 Subject: [PATCH 280/365] fix(executor): read a gathered reply only once it has arrived The previous commit probed each pending queue with dequeue(timeout=0). A reply larger than a ring chunk marks its shm slot first and follows over the socket, so a probe could claim the slot, time out before the payload landed, and lose the reply for good. MessageQueue.ready() reports whether a message is waiting for the reader without consuming anything (the slot's written/read flags locally, a zero-timeout socket poll remotely). The gather reads a queue only once it is ready, and then with the RPC's own deadline; between rounds in which nothing arrived it sleeps 1 ms. The fake queues in the tests now fail on any read before ready(), and a test with real queues gathers a 200 KB reply over the socket path beside a small one. Signed-off-by: tobymao --- .../test_multiproc_gather_responses.py | 65 +++++++++++++++---- .../device_communicators/shm_broadcast.py | 18 +++++ vllm/v1/executor/multiproc_executor.py | 48 +++++++------- 3 files changed, 94 insertions(+), 37 deletions(-) diff --git a/tests/v1/executor/test_multiproc_gather_responses.py b/tests/v1/executor/test_multiproc_gather_responses.py index 20459566cf85..a6b929f30090 100644 --- a/tests/v1/executor/test_multiproc_gather_responses.py +++ b/tests/v1/executor/test_multiproc_gather_responses.py @@ -6,10 +6,12 @@ failed out of) must not hide a failure another rank has already reported. """ +import threading import time import pytest +from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.v1.executor.multiproc_executor import WorkerProc, _gather_responses SUCCESS = WorkerProc.ResponseStatus.SUCCESS @@ -17,21 +19,25 @@ class _Queue: - """Replies once ``ready_after`` dequeue calls have passed; never if None.""" + """Has its reply ready after ``ready_after`` readiness checks; never if None. + + A dequeue before the reply is ready fails the test: the gather must only + read a queue that ready() reported, or it can lose a large reply. + """ def __init__(self, reply=None, ready_after: int = 0) -> None: self.reply = reply self.ready_after = ready_after - self.calls = 0 + self.checks = 0 + + def ready(self) -> bool: + self.checks += 1 + return self.reply is not None and self.checks > self.ready_after def dequeue(self, timeout=None): - self.calls += 1 - if self.reply is not None and self.calls > self.ready_after: - reply, self.reply = self.reply, None - return reply - if timeout: - time.sleep(min(timeout, 0.01)) - raise TimeoutError + assert self.reply is not None and self.checks > self.ready_after + reply, self.reply = self.reply, None + return reply def test_a_later_rank_failure_surfaces_while_rank_zero_hangs(): @@ -57,7 +63,44 @@ def test_the_deadline_still_applies(): _gather_responses(queues, range(2), time.monotonic() + 0.2, "m") +def test_a_reply_larger_than_a_ring_chunk_is_not_lost(): + """Real queues: a reply over max_chunk_bytes goes out over the socket after + its slot is marked, so reading it must wait for the payload.""" + pairs = [] + try: + for _ in range(2): + writer = MessageQueue( + n_reader=1, n_local_reader=1, max_chunk_bytes=1024, max_chunks=2 + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + writer.wait_until_ready() + reader.wait_until_ready() + pairs.append((writer, reader)) + big = b"x" * 200_000 + + def send() -> None: + time.sleep(0.05) + pairs[1][0].enqueue((SUCCESS, big)) + time.sleep(0.05) + pairs[0][0].enqueue((SUCCESS, "small")) + + sender = threading.Thread(target=send) + sender.start() + replies = _gather_responses( + [reader for _, reader in pairs], range(2), time.monotonic() + 10, "m" + ) + sender.join() + assert replies == ["small", big] + finally: + for writer, reader in pairs: + writer.shutdown() + reader.shutdown() + + def test_a_single_reply_names_its_rank(): + # One reply is one blocking read (safe for any size), so no readiness check. + boom = _Queue((FAILURE, "boom"), ready_after=-1) with pytest.raises(RuntimeError, match="Worker 3 failed"): - _gather_responses([_Queue((FAILURE, "boom"))], (3,), None, "m") - assert _gather_responses([_Queue((SUCCESS, "ok"))], (3,), None, "m") == ["ok"] + _gather_responses([boom], (3,), None, "m") + ok = _Queue((SUCCESS, "ok"), ready_after=-1) + assert _gather_responses([ok], (3,), None, "m") == ["ok"] diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 2171f79fb7ce..49b3ffe4df1a 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -891,6 +891,24 @@ def oob_callback(buf: PickleBuffer) -> bool: if self.n_remote_reader > 0: self.remote_socket.send_multipart(all_buffers, copy=False) + def ready(self) -> bool: + """Whether a message is waiting for this reader, without consuming it. + + A probe that dequeues with a zero timeout is not safe: a message larger + than a ring chunk marks its slot first and follows over the socket, so + the probe can claim the slot and time out before the payload lands, + losing the message. + """ + if self._is_local_reader: + with self.buffer.get_metadata(self.current_idx) as metadata_buffer: + memory_fence() + return bool(metadata_buffer[0]) and not bool( + metadata_buffer[self.local_reader_rank + 1] + ) + if self._is_remote_reader: + return bool(self.remote_socket.poll(timeout=0)) + raise RuntimeError("Only readers can check for messages") + def dequeue( self, timeout: float | None = None, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 68562140d013..b22acc8f6fe4 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -80,9 +80,8 @@ logger = init_logger(__name__) -# How long a gather of every rank's reply blocks on one rank before checking -# the others for a reply (a failure among them ends the wait). -RESPONSE_POLL_SECONDS = 0.05 +# How often a gather of every rank's reply checks the pending ranks. +RESPONSE_POLL_SECONDS = 0.001 def _check_response(rank: int, status: Any, result: Any) -> Any: @@ -106,8 +105,9 @@ def _gather_responses( while an earlier rank waited forever, for example in a collective the failed rank would never join: TP1's CUDA graph capture fails, TP0 hangs in the next all-reduce, and the engine waits on TP0 with TP1's error in - its queue. Every pending queue is polled without blocking, and only when - none has replied does the gather block, briefly, on the oldest one. + its queue. Every pending queue is checked without consuming anything, each + reply is read as it arrives, and the gather sleeps 1 ms between rounds in + which nothing arrived. """ def remaining() -> float | None: @@ -127,34 +127,30 @@ def remaining() -> float | None: responses: list[Any] = [None] * len(response_mqs) pending = list(range(len(response_mqs))) - # A blocking read logs the long-wait message every interval; the short - # polls here would not, and operators watch for it, so log it the same way. + # A blocking read logs the long-wait message every interval; the polling + # here would not, and operators watch for it, so log it the same way. started, warnings = time.monotonic(), 0 interval = envs.VLLM_RINGBUFFER_WARNING_INTERVAL while pending: - if time.monotonic() - started >= interval * (warnings + 1): - warnings += 1 - logger.info(LONG_WAIT_TIME_LOG_MSG, interval) + progressed = False for index in list(pending): - try: - status, result = response_mqs[index].dequeue(timeout=0) - except TimeoutError: + # ready() never consumes, so a reply is only read once it has + # arrived, and then with the RPC's own deadline. + if not response_mqs[index].ready(): continue + try: + status, result = response_mqs[index].dequeue(timeout=remaining()) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e responses[index] = _check_response(ranks[index], status, result) pending.remove(index) - if not pending: - break - left = remaining() - wait = ( - RESPONSE_POLL_SECONDS if left is None else min(RESPONSE_POLL_SECONDS, left) - ) - index = pending[0] - try: - status, result = response_mqs[index].dequeue(timeout=wait) - except TimeoutError: - continue - responses[index] = _check_response(ranks[index], status, result) - pending.remove(index) + progressed = True + if pending and not progressed: + remaining() # raises once the deadline has passed + if time.monotonic() - started >= interval * (warnings + 1): + warnings += 1 + logger.info(LONG_WAIT_TIME_LOG_MSG, interval) + time.sleep(RESPONSE_POLL_SECONDS) return responses From 4c5e4272e58dfc599e7e8d5fb0cb36406ec7574d Mon Sep 17 00:00:00 2001 From: Costa Tsaousis Date: Sun, 27 Sep 2026 22:34:04 +0000 Subject: [PATCH 281/365] Preserve external checkpoint restores through capacity-constrained admission Signed-off-by: Costa Tsaousis --- .lil/changes/vllm-restore-admission.json | 16 + tests/v1/core/test_boundary_admission.py | 416 ++++++++++++++++++++++- vllm/v1/core/kv_cache_manager.py | 308 +++++++++++++++-- vllm/v1/core/sched/scheduler.py | 60 +++- 4 files changed, 765 insertions(+), 35 deletions(-) create mode 100644 .lil/changes/vllm-restore-admission.json diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json new file mode 100644 index 000000000000..c6f2c86894e2 --- /dev/null +++ b/.lil/changes/vllm-restore-admission.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-restore-admission", + "category": "fix", + "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", + "models": ["all"], + "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", + "details": [ + "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", + "Sliding-window and chunked-local continuation reservations count physical pages rather than skipped positions in the block table.", + "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", + "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." + ], + "authors": ["ktsaou"], + "requires": [] +} diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index 02f02b449331..6f7c0c4f958f 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -11,12 +11,17 @@ from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request from vllm.utils.hashing import sha256 -from vllm.v1.core.boundary_checkpoint import BoundaryCheckpointCache +from vllm.v1.core.boundary_checkpoint import ( + BoundaryCheckpointCache, + boundary_checkpoint_slots, +) from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import RequestStatus @@ -38,7 +43,7 @@ def compute_share_fixture_options(share: float | None) -> dict[str, Any]: return options -def manager(): +def manager(dcp=1, num_blocks=128): attention = MLAAttentionSpec( block_size=64, num_kv_heads=1, @@ -56,13 +61,15 @@ def manager(): groups = [KVCacheGroupSpec(["attention"], attention)] groups += [KVCacheGroupSpec([f"recurrent-{i}"], recurrent) for i in range(3)] return make_kv_cache_manager( - KVCacheConfig(num_blocks=128, kv_cache_tensors=[], kv_cache_groups=groups), + KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ), max_model_len=512, max_in_flight_tokens=32, enable_caching=True, use_eagle=True, num_prefill_lookahead=1, - dcp_world_size=1, + dcp_world_size=dcp, scheduler_block_size=64, hash_block_size=16, enable_boundary_checkpoints=True, @@ -81,6 +88,292 @@ def drain(cache): cache.block_pool.free_blocks(copies) +def reserve_import(cache, consumer, prefix=100): + return cache.reserve_external_boundary_checkpoint( + consumer, + prefix, + cache.boundary_checkpoint_page_positions(prefix), + draft_prefix_len=prefix, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + + +@pytest.mark.parametrize("dcp", [1, 2]) +@pytest.mark.parametrize("prefix", [16, 100, 128]) +@pytest.mark.parametrize("chunk", [1, 16, 33, 127]) +def test_external_import_owns_execution_capacity_through_chunked_prefill( + dcp, prefix, chunk +): + """Other admissions cannot consume a restored request's remaining suffix.""" + cache = manager(dcp) + consumer = request("import", length=400) + free = cache.block_pool.get_num_free_blocks() + checkpoint = reserve_import(cache, consumer, prefix=prefix) + assert checkpoint is not None + credits = cache.external_boundary_reserved_blocks() + assert credits > 0 + assert not cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, 0 + ) + assert cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, 1) + assert not cache.reset_prefix_cache() + pressure = cache.block_pool.get_new_blocks( + cache.block_pool.get_num_free_blocks() - credits + ) + assert cache.allocate_slots(request("competitor", "other"), 1) is None + blocks, hits, _ = cache.get_computed_blocks(consumer) + assert hits == prefix + assert ( + cache.allocate_slots( + consumer, + chunk, + num_new_computed_tokens=prefix, + new_computed_blocks=blocks, + num_lookahead_tokens=3, + full_sequence_must_fit=True, + ) + is not None + ) + consumer.num_computed_tokens = prefix + chunk + drain(cache) + assert cache.has_external_boundary_admission(consumer.request_id) + assert cache.external_boundary_reserved_blocks() > 0 + while consumer.num_computed_tokens < consumer.num_tokens: + count = min(chunk, consumer.num_tokens - consumer.num_computed_tokens) + assert cache.allocate_slots(consumer, count, num_lookahead_tokens=3) is not None + consumer.num_computed_tokens += count + drain(cache) + assert not cache.has_external_boundary_admission(consumer.request_id) + cache.free(consumer) + cache.block_pool.free_blocks(pressure) + assert cache.block_pool.get_num_free_blocks() == free + + +@pytest.mark.parametrize( + "attention_kind,retained", [("sliding", 0), ("sliding", 48), ("chunked", 0)] +) +@pytest.mark.parametrize( + "prefix,length", [(128, 400), (128, 128), (128, 129), (100, 100), (100, 101)] +) +def test_scheduler_preserves_recycling_import_capacity( + attention_kind, retained, prefix, length +): + """Null positions must not release credits while the live window moves.""" + from vllm.v1.core.sched.scheduler import Scheduler + from vllm.v1.outputs import ModelRunnerOutput + + template = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + max_num_batched_tokens=32, + long_prefill_token_threshold=16, + block_size=64, + max_model_len=512, + ) + config = template.vllm_config + config.cache_config.mamba_cache_mode = "align" + config.cache_config.num_gpu_blocks = 64 + attention_args = dict(block_size=64, num_kv_heads=1, head_size=1, dtype=torch.uint8) + attention = ( + SlidingWindowSpec( + **attention_args, sliding_window=128, extra_retained_tokens=retained + ) + if attention_kind == "sliding" + else ChunkedLocalAttentionSpec(**attention_args, attention_chunk_size=128) + ) + recurrent = MambaSpec( + block_size=16, shapes=((1,),), dtypes=(torch.uint8,), mamba_cache_mode="align" + ) + scheduler = Scheduler( + vllm_config=config, + kv_cache_config=KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], attention), + KVCacheGroupSpec(["recurrent"], recurrent), + ], + ), + structured_output_manager=template.structured_output_manager, + block_size=64, + hash_block_size=16, + ) + scheduler.use_v2_model_runner = True + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + consumer = request("recycling-import", length=length) + scheduler.add_request(consumer) + checkpoint = reserve_import(cache, consumer, prefix=prefix) + assert checkpoint is not None + if length <= prefix + 1: + # One attention page and two private recurrent states suffice for + # this suffix; imported read-only history need not be replaced. + assert cache.external_boundary_reserved_blocks() <= ( + 3 + 3 * len(boundary_checkpoint_slots(consumer)) + ) + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + pressure = [] + while consumer.num_computed_tokens < consumer.num_tokens: + spare = ( + cache.block_pool.get_num_free_blocks() + - cache.external_boundary_reserved_blocks() + ) + pressure.extend(cache.block_pool.get_new_blocks(max(0, spare))) + expected = min(16, max(1, length - max(prefix, consumer.num_computed_tokens))) + output = scheduler.schedule() + assert output.num_scheduled_tokens == {consumer.request_id: expected} + assert consumer.num_preemptions == 0 + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=[consumer.request_id], + req_id_to_index={consumer.request_id: 0}, + sampled_token_ids=[[]], + ), + ) + cache.free(consumer) + cache.block_pool.free_blocks(pressure) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.block_pool.get_num_free_blocks() == 63 + + +@pytest.mark.parametrize("cancel_at", ["copying", "ready"]) +def test_external_import_cancellation_releases_only_after_copy_drain(cancel_at): + cache = manager() + consumer = request("cancel") + free = cache.block_pool.get_num_free_blocks() + checkpoint = reserve_import(cache, consumer) + assert checkpoint is not None + if cancel_at == "ready": + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, rank + ) + cache.free(consumer) + if cancel_at == "copying": + assert cache.block_pool.get_num_free_blocks() < free + assert cache.external_boundary_reserved_blocks() > 0 + cache.discard_external_boundary_checkpoint(checkpoint.checkpoint_id) + cache.release_external_boundary_admission(consumer.request_id) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.block_pool.get_num_free_blocks() == free + + +def test_external_imports_reserve_slots_without_serializing_available_imports(): + cache = manager() + cache.set_external_boundary_admission_context(2, 0) + first, second, third = [request(name, name) for name in ("one", "two", "three")] + assert reserve_import(cache, first) is not None + assert reserve_import(cache, second) is not None + free = cache.block_pool.get_num_free_blocks() + assert reserve_import(cache, third) is None + assert not cache.can_admit_external_boundary_request(third.request_id) + assert cache.block_pool.get_num_free_blocks() == free + assert cache.can_admit_external_boundary_request(first.request_id) + + +def test_external_import_preflight_does_not_evict_cache_on_execution_pressure(): + cache = manager() + consumer = request("consumer") + cache.set_external_boundary_admission_context( + 1, cache.block_pool.get_num_free_blocks() + ) + free = cache.block_pool.get_num_free_blocks() + assert reserve_import(cache, consumer) is None + assert cache.block_pool.get_num_free_blocks() == free + assert cache.external_boundary_reserved_blocks() == 0 + + +def test_capacity_waiter_prevents_newer_admissions_from_stealing_released_space(): + cache = manager() + older, newer = request("older"), request("newer", "newer") + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert reserve_import(cache, older) is None + cache.block_pool.free_blocks(pressure) + assert not cache.can_admit_external_boundary_request(newer.request_id) + assert reserve_import(cache, newer) is None + assert reserve_import(cache, older) is not None + assert cache.can_admit_external_boundary_request(newer.request_id) + + +def test_cancelled_capacity_waiter_releases_admission_barrier(): + cache = manager() + consumer = request("cancel-waiter") + cache.set_external_boundary_admission_context( + 1, cache.block_pool.get_num_free_blocks() + ) + assert reserve_import(cache, consumer) is None + assert not cache.can_admit_external_boundary_request("later") + cache.release_external_boundary_admission(consumer.request_id) + assert cache.can_admit_external_boundary_request("later") + + +def test_import_preflight_preserves_watermark_and_classifies_impossible_restore(): + cache = manager() + consumer = request("watermark") + free = cache.block_pool.get_num_free_blocks() + cache.watermark_blocks = free + cache.set_external_boundary_admission_context(1, 0, has_scheduled_reqs=True) + assert reserve_import(cache, consumer) is None + assert cache.block_pool.get_num_free_blocks() == free + # The last runnable request finishing removes the normal admission watermark. + cache.set_external_boundary_admission_context(1, 0, has_scheduled_reqs=False) + assert reserve_import(cache, consumer) is not None + cache = manager(num_blocks=16) + free = cache.block_pool.get_num_free_blocks() + with pytest.raises(ValueError, match="exceed the GPU pool"): + reserve_import(cache, consumer) + assert cache.block_pool.get_num_free_blocks() == free + assert not cache.has_pending_external_boundary_admissions() + + +def test_eight_external_sessions_rotate_without_cold_prefill_or_leaked_blocks(): + """Aggregate checkpoints exceed the pool; each admitted turn still restores.""" + cache = manager(dcp=2, num_blocks=64) + free = cache.block_pool.get_num_free_blocks() + consumers = [request(f"session-{i}", str(i), length=400) for i in range(8)] + pages = sum(map(len, cache.boundary_checkpoint_page_positions(400))) + 1 + assert pages * len(consumers) > free + remaining = list(consumers) + completed = [] + while remaining: + admitted = [] + cache.set_external_boundary_admission_context(8, 0) + for consumer in remaining: + checkpoint = reserve_import(cache, consumer, prefix=400) + if checkpoint is not None: + admitted.append((consumer, checkpoint)) + assert admitted, "imports must make progress when previous owners drain" + for consumer, checkpoint in admitted: + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, rank + ) + blocks, hits, _ = cache.get_computed_blocks(consumer) + assert hits == consumer.num_tokens + assert ( + cache.allocate_slots( + consumer, + 1, + num_new_computed_tokens=hits, + new_computed_blocks=blocks, + num_lookahead_tokens=3, + full_sequence_must_fit=True, + ) + is not None + ) + drain(cache) + cache.free(consumer) + remaining.remove(consumer) + completed.append(consumer.request_id) + assert completed == [consumer.request_id for consumer in consumers] + assert cache.block_pool.get_num_free_blocks() == free + assert cache.external_boundary_reserved_blocks() == 0 + + def seed(cache, salt): producer = request("producer-" + salt, salt) cache.get_computed_blocks(producer) @@ -270,6 +563,121 @@ def make_scheduler(**kwargs): return create_scheduler(**kwargs) +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +@pytest.mark.parametrize("lanes", [1, 2]) +@pytest.mark.parametrize("limit", ["slots", "blocks"]) +def test_ready_import_behind_blocked_head_gets_saved_logits_step(policy, lanes, limit): + """Reserved imports remain reachable while unrelated requests wait.""" + from unittest.mock import Mock + + from tests.v1.core.utils import create_requests + + scheduler = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + async_scheduling=True, + max_num_seqs=2 if limit == "slots" else 3, + max_parallel_prefills=lanes, + scheduling_policy=policy, + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + running, consumer = create_requests( + num_requests=2, + num_tokens=32, + req_ids=["running", "consumer"], + ) + (cold,) = create_requests(num_requests=1, num_tokens=160, req_ids=["cold"]) + scheduler.add_request(running) + assert scheduler.schedule().num_scheduled_tokens == {"running": 32} + checkpoint = cache.reserve_external_boundary_checkpoint( + consumer, + 32, + cache.boundary_checkpoint_page_positions(32), + draft_prefix_len=32, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + assert checkpoint is not None + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + pressure = [] + if limit == "blocks": + pressure = cache.block_pool.get_new_blocks( + cache.block_pool.get_num_free_blocks() + - cache.external_boundary_reserved_blocks() + ) + scheduler.connector = Mock() + scheduler.connector.poll_boundary_checkpoint.return_value = True + scheduler.connector.boundary_checkpoint_external_tokens.side_effect = lambda req: ( + 32 if req is consumer else 0 + ) + scheduler.connector.get_num_new_matched_tokens.return_value = (0, False) + scheduler.add_request(cold) + scheduler.add_request(consumer) + output = scheduler.schedule() + assert output.boundary_logits_only + assert output.num_scheduled_tokens == {"consumer": 1} + assert consumer.prefill_stats.num_computed_tokens == 0 + assert consumer.prefill_stats.num_external_cached_tokens == 32 + assert cold.status == RequestStatus.WAITING + cache.block_pool.free_blocks(pressure) + + +def test_streaming_owner_can_resume_ahead_of_capacity_waiter(): + from unittest.mock import Mock + + from tests.v1.core.utils import create_requests + + scheduler = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + max_num_seqs=8, + num_blocks=12, + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + (session,) = create_requests(num_requests=1, num_tokens=100, req_ids=["stream"]) + session.resumable = True + scheduler.add_request(session) + assert scheduler.schedule().num_scheduled_tokens == {"stream": 100} + session.num_in_flight_tokens = 0 + scheduler.running.remove(session) + assert not scheduler._handle_stopped_request(session) + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + (consumer,) = create_requests(num_requests=1, num_tokens=64, req_ids=["restore"]) + connector = Mock() + + def poll(req): + if req is not consumer: + return True + restored = cache.reserve_external_boundary_checkpoint( + req, + 64, + cache.boundary_checkpoint_page_positions(64), + draft_prefix_len=64, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + assert restored is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + connector.boundary_checkpoint_external_tokens.return_value = 0 + connector.get_num_new_matched_tokens.return_value = (0, False) + scheduler.connector = connector + scheduler.add_request(consumer) + assert scheduler.schedule().num_scheduled_tokens == {} + assert not cache.can_admit_external_boundary_request(session.request_id) + (update,) = create_requests(num_requests=1, num_tokens=1, req_ids=["stream"]) + update.resumable = True + scheduler.add_request(update) + assert session.status == RequestStatus.WAITING + assert scheduler.schedule().num_scheduled_tokens == {"stream": 1} + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index de3b56ff7183..a737fa9e6aa1 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -26,13 +26,18 @@ ) from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy -from vllm.v1.core.single_type_kv_cache_manager import MambaManager +from vllm.v1.core.single_type_kv_cache_manager import ( + MambaManager, + SingleTypeKVCacheManager, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, + ChunkedLocalAttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, KVCacheConfig, MambaSpec, + SlidingWindowSpec, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, iter_layer_specs, @@ -43,6 +48,15 @@ logger = init_logger(__name__) +@dataclass +class _BoundaryImportAdmission: + request: Request + checkpoint: BoundaryCheckpoint + remaining_blocks: int + ready: bool = False + admitted: bool = False + + @dataclass class KVCacheBlocks: """ @@ -171,6 +185,7 @@ def __init__( # always supplies the real value at runtime. if max_in_flight_tokens is None: max_in_flight_tokens = max_model_len + self._boundary_max_in_flight_tokens = max_in_flight_tokens self.enable_caching = enable_caching self.enable_kv_cache_events = enable_kv_cache_events @@ -247,6 +262,11 @@ def __init__( else None ) self._boundary_imports: set[int] = set() + self._boundary_import_admissions: dict[str, _BoundaryImportAdmission] = {} + self._boundary_import_waiters: dict[str, None] = {} + self._external_admission_slots: int | None = None + self._external_admission_reserved_blocks = 0 + self._external_admission_has_scheduled_reqs = False self._boundary_allocations: dict[str, list[KVCacheBlock]] = {} self._boundary_readers: dict[str, BoundaryCheckpoint] = {} self._boundary_reader_horizons: dict[str, int] = {} @@ -339,7 +359,12 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int return self.empty_kv_cache_blocks, 0, 0 if boundary_cache is not None and boundary_cache.supports_request(request): - checkpoint = boundary_cache.find(request, request.num_tokens) + admission = self._boundary_import_admissions.get(request.request_id) + checkpoint = ( + admission.checkpoint + if admission is not None and admission.ready and not admission.admitted + else boundary_cache.find(request, request.num_tokens) + ) request.boundary_checkpoint = checkpoint if checkpoint is not None: checkpoint_blocks = tuple( @@ -597,9 +622,7 @@ def allocate_slots( boundary_blocks = 0 boundary_replay_managers = [] if request.use_boundary_checkpoints: - if request.request_id not in self._boundary_allocations: - checkpoint_slots = len(boundary_checkpoint_slots(request)) - boundary_blocks = checkpoint_slots * (self.num_kv_cache_groups + 1) + boundary_blocks = self._boundary_capture_blocks(request) if request.boundary_checkpoint is not None: boundary_blocks += sum( self.block_pool.blocks[i].ref_cnt == 0 @@ -610,15 +633,9 @@ def allocate_slots( and new_computed_blocks is not None and num_local_computed_tokens > 0 ): - boundary_replay_managers = [ - manager - for manager in self.coordinator.single_type_managers - if num_local_computed_tokens % manager.block_size == 0 - and not any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(manager.kv_cache_spec) - ) - ] + boundary_replay_managers = self._boundary_replay_managers( + num_local_computed_tokens + ) boundary_blocks += len(boundary_replay_managers) # The watermark is applied to waiting/preempted requests only, and only # when there's at least one request already scheduled. @@ -645,7 +662,10 @@ def allocate_slots( required_blocks = ( num_blocks_to_allocate + watermark_blocks + boundary_blocks ) - if required_blocks > self.block_pool.get_num_free_blocks(): + if required_blocks > ( + self.block_pool.get_num_free_blocks() + - self.external_boundary_reserved_blocks(request.request_id) + ): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -681,7 +701,11 @@ def allocate_slots( # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. - available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks + available_blocks = ( + self.block_pool.get_num_free_blocks() + - reserved_blocks + - self.external_boundary_reserved_blocks(request.request_id) + ) required_blocks = num_blocks_to_allocate + watermark_blocks + boundary_blocks if required_blocks > available_blocks: # Cannot allocate new blocks @@ -763,9 +787,14 @@ def allocate_slots( and request.request_id not in self._boundary_readers ): assert self.boundary_checkpoints is not None - checkpoint = self.boundary_checkpoints.acquire( - request.boundary_checkpoint.checkpoint_id - ) + admission = self._boundary_import_admissions.get(request.request_id) + if admission is not None and admission.ready and not admission.admitted: + checkpoint = admission.checkpoint + admission.admitted = True + else: + checkpoint = self.boundary_checkpoints.acquire( + request.boundary_checkpoint.checkpoint_id + ) if checkpoint is None: request.boundary_checkpoint = None return None @@ -813,6 +842,19 @@ def allocate_slots( by_slot.get(slot, (0,) * width) for slot in range(max(slots) + 1) ) + admission = self._boundary_import_admissions.get(request.request_id) + if admission is not None and admission.admitted: + if num_tokens_main_model >= request.num_tokens: + del self._boundary_import_admissions[request.request_id] + else: + admission.remaining_blocks = self._external_boundary_execution_blocks( + request, + num_tokens_main_model, + self.empty_kv_cache_blocks.blocks, + restoring=False, + ) + self._boundary_import_waiters.pop(request.request_id, None) + # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. if delay_cache_blocks: @@ -844,6 +886,7 @@ def free(self, request: Request) -> None: self.coordinator.free(request.request_id) def _pop_boundary_blocks(self, request: Request) -> list[KVCacheBlock]: + self.release_external_boundary_admission(request.request_id) blocks = self._boundary_allocations.pop(request.request_id, []) reader = self._boundary_readers.pop(request.request_id, None) self._boundary_reader_horizons.pop(request.request_id, None) @@ -967,6 +1010,7 @@ def reserve_external_boundary_checkpoint( draft_prefix_len: int, kind: BoundaryCheckpointKind, num_ranks: int, + reserve_admission: bool = False, ) -> BoundaryCheckpoint | None: """Reserve private destinations for an externally stored checkpoint. @@ -974,9 +1018,15 @@ def reserve_external_boundary_checkpoint( acknowledges its completed H2D copies. Destinations are pinned in the ordinary KV pool, including one page for target/draft auxiliary state. - Returns None if request-boundary caching is unavailable or the pool - cannot admit every live page. Raises ValueError for incompatible cache - geometry; callers must treat that as an external cache miss. + With reserve_admission, also reserve a running slot and continuation + credits. Successful publication retains the consumer pin until ordinary + allocation takes ownership. The scheduler must refresh admission context + before each candidate; cancellation must release the admission after all + copies drain. + + Returns None for temporary resource pressure or unavailable caching. + Raises ValueError for incompatible geometry or an import whose checkpoint + and execution state cannot fit even an otherwise empty pool. """ cache = self.boundary_checkpoints if ( @@ -988,6 +1038,10 @@ def reserve_external_boundary_checkpoint( if not 0 < num_tokens <= request.num_tokens: raise ValueError("Imported boundary must cover an existing request prefix") if not checkpoint_end_allowed(request, num_tokens): + if reserve_admission: + raise ValueError( + "Imported checkpoint ends inside an unauthenticated span" + ) return None if not 0 <= draft_prefix_len <= num_tokens: raise ValueError("Imported draft prefix exceeds the target prefix") @@ -999,7 +1053,37 @@ def reserve_external_boundary_checkpoint( if page_positions != self.boundary_checkpoint_page_positions(num_tokens): raise ValueError("External checkpoint does not cover every live cache page") count = sum(len(group) for group in page_positions) + 1 - if count > self.block_pool.get_num_free_blocks(): + execution_blocks = 0 + if reserve_admission: + if request.request_id in self._boundary_import_admissions: + return None + if not self.can_admit_external_boundary_request(request.request_id): + return None + # Estimation must not evict cached pages or mutate real block metadata. + pinned = KVCacheBlock(block_id=-1, ref_cnt=1) + prospective = tuple( + tuple( + pinned if i in positions else self.block_pool.null_block + for i in range(positions[-1] + 1) + ) + for positions in page_positions + ) + execution_blocks = self._external_boundary_execution_blocks( + request, + num_tokens, + prospective, + restoring=True, + ) + if count + execution_blocks > self.block_pool.num_gpu_blocks - 1: + raise ValueError("Checkpoint and execution state exceed the GPU pool") + reserved = self.external_boundary_reserved_blocks() + if reserve_admission: + reserved += self._external_admission_reserved_blocks + if self._external_admission_has_scheduled_reqs: + reserved += self.watermark_blocks + if count + execution_blocks + reserved > self.block_pool.get_num_free_blocks(): + if reserve_admission: + self._boundary_import_waiters.setdefault(request.request_id, None) return None allocation = self.block_pool.get_new_blocks(count) try: @@ -1020,6 +1104,11 @@ def reserve_external_boundary_checkpoint( ) cache.stage(request, checkpoint, num_ranks=num_ranks) self._boundary_imports.add(checkpoint.checkpoint_id) + if reserve_admission: + self._boundary_import_waiters.pop(request.request_id, None) + self._boundary_import_admissions[request.request_id] = ( + _BoundaryImportAdmission(request, checkpoint, execution_blocks) + ) return checkpoint finally: # stage() owns a separate pin on every dependency until publication @@ -1036,6 +1125,20 @@ def acknowledge_external_boundary_checkpoint( published = cache.acknowledge(checkpoint_id, rank) if not cache.is_pending(checkpoint_id): self._boundary_imports.remove(checkpoint_id) + for request_id, admission in tuple( + self._boundary_import_admissions.items() + ): + if admission.checkpoint.checkpoint_id != checkpoint_id: + continue + if published: + # No scheduler allocation may intervene between publication + # and acquiring the consumer's ownership. + acquired = cache.acquire(checkpoint_id) + assert acquired is not None + admission.ready = True + admission.request.boundary_checkpoint = admission.checkpoint + else: + del self._boundary_import_admissions[request_id] return published def discard_external_boundary_checkpoint(self, checkpoint_id: int) -> None: @@ -1048,6 +1151,165 @@ def discard_external_boundary_checkpoint(self, checkpoint_id: int) -> None: assert self.boundary_checkpoints is not None self.boundary_checkpoints.discard(checkpoint_id) self._boundary_imports.remove(checkpoint_id) + for request_id, admission in tuple( + self._boundary_import_admissions.items() + ): + if admission.checkpoint.checkpoint_id == checkpoint_id: + del self._boundary_import_admissions[request_id] + + def set_external_boundary_admission_context( + self, + available_slots: int, + reserved_blocks: int, + *, + has_scheduled_reqs: bool = False, + ) -> None: + """Refresh scheduler capacity immediately before polling a candidate. + + Args: + available_slots: Slots remaining after running and paused sessions. + reserved_blocks: Remaining needs of prefills without import tickets. + has_scheduled_reqs: Apply the normal admission watermark when true. + """ + self._external_admission_slots = available_slots + self._external_admission_reserved_blocks = reserved_blocks + self._external_admission_has_scheduled_reqs = has_scheduled_reqs + + def has_external_boundary_admission(self, request_id: str) -> bool: + """Whether this request already owns import continuation credits.""" + return request_id in self._boundary_import_admissions + + def external_boundary_reserved_blocks(self, owner: str | None = None) -> int: + """Return outstanding import credits, excluding the allocating owner.""" + return sum( + admission.remaining_blocks + for request_id, admission in self._boundary_import_admissions.items() + if request_id != owner + ) + + def can_admit_external_boundary_request(self, request_id: str) -> bool: + """Whether a candidate can use a running slot without stealing a ticket.""" + if request_id in self._boundary_import_admissions: + return True + if self._boundary_import_waiters and request_id != next( + iter(self._boundary_import_waiters) + ): + return False + return ( + self._external_admission_slots is None + or sum( + not admission.admitted + for admission in self._boundary_import_admissions.values() + ) + < self._external_admission_slots + ) + + def has_pending_external_boundary_admissions(self) -> bool: + """Whether an import still needs ordinary scheduler admission.""" + return any(not a.admitted for a in self._boundary_import_admissions.values()) + + def external_boundary_admission_ready(self, request_id: str) -> bool: + """Whether a waiting import has finished copying and owns its checkpoint.""" + admission = self._boundary_import_admissions.get(request_id) + return admission is not None and admission.ready and not admission.admitted + + def release_external_boundary_admission(self, request_id: str) -> None: + """Release a finished consumer's credits; in-flight copies must drain first.""" + self._boundary_import_waiters.pop(request_id, None) + admission = self._boundary_import_admissions.get(request_id) + if ( + admission is None + or admission.checkpoint.checkpoint_id in self._boundary_imports + ): + return + del self._boundary_import_admissions[request_id] + if admission.ready and not admission.admitted: + assert self.boundary_checkpoints is not None + self.boundary_checkpoints.release(admission.checkpoint) + + def _external_boundary_execution_blocks( + self, + request: Request, + prefix: int, + blocks: tuple[Sequence[KVCacheBlock], ...], + *, + restoring: bool, + ) -> int: + full = min(request.num_tokens + 1, self.max_model_len) + needed = self.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=min(full + self._boundary_restore_lookahead, self.max_model_len), + new_computed_blocks=blocks, + num_encoder_tokens=0, + total_computed_tokens=prefix, + num_local_computed_tokens=prefix, + num_tokens_main_model=full, + apply_admission_cap=True, + ) + reader = self._boundary_readers.get(request.request_id) + dependencies = set(reader.dependencies) if reader is not None else set() + for group, manager in zip(blocks, self.coordinator.single_type_managers): + spec = manager.kv_cache_spec + if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align": + # Continuation overlaps running/previous states while the + # imported source stays pinned. + peak = ( + 2 + spec.num_speculative_blocks + spec.num_prefill_checkpoint_blocks + ) + elif isinstance(spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): + # Logical tables retain null positions after a window moves; + # only private physical pages can satisfy the continuation peak. + peak = spec.max_admission_blocks_per_request( + self._boundary_max_in_flight_tokens, self.max_model_len + ) + else: + continue + owned = sum( + not block.is_null and block.block_id not in dependencies + for block in manager.req_to_blocks.get(request.request_id, ()) + ) + endpoint_need = manager.get_num_blocks_to_allocate( + request.request_id, + min(full + self._boundary_restore_lookahead, self.max_model_len), + group, + prefix, + prefix, + full, + apply_admission_cap=True, + ) + remaining = max(0, peak - owned) + if isinstance(spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): + # A short suffix may never fill a complete private window. + suffix_need = manager.get_num_blocks_to_allocate( + request.request_id, + min(full + self._boundary_restore_lookahead, self.max_model_len), + group, + prefix, + prefix, + full, + ) + remaining = min(remaining, suffix_need) + needed += max(0, remaining - endpoint_need) + needed += self._boundary_capture_blocks(request) + if restoring and self.use_eagle: + needed += len(self._boundary_replay_managers(prefix)) + return needed + + def _boundary_capture_blocks(self, request: Request) -> int: + if request.request_id in self._boundary_allocations: + return 0 + return len(boundary_checkpoint_slots(request)) * (self.num_kv_cache_groups + 1) + + def _boundary_replay_managers(self, prefix: int) -> list[SingleTypeKVCacheManager]: + return [ + manager + for manager in self.coordinator.single_type_managers + if prefix % manager.block_size == 0 + and not any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(manager.kv_cache_spec) + ) + ] def remove_skipped_blocks( self, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f2150fd9b51d..a61b5fac1c9d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -782,14 +782,30 @@ def _has_waiting_boundary_logits(self) -> bool: queue = self._select_waiting_queue_for_scheduling() if queue is None: return False - request = queue.peek_request() - checkpoint = request.boundary_checkpoint - return ( + pending_imports = ( + self.kv_cache_manager.has_pending_external_boundary_admissions() + ) + candidates = itertools.chain( + (queue.peek_request(),), + ( + request + for request in ( + itertools.chain(self.skipped_waiting, self.waiting) + if pending_imports + else () + ) + if self.kv_cache_manager.external_boundary_admission_ready( + request.request_id + ) + ), + ) + return any( request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) and request.num_computed_tokens == 0 - and checkpoint is not None - and checkpoint.num_tokens == request.num_tokens + and request.boundary_checkpoint is not None + and request.boundary_checkpoint.num_tokens == request.num_tokens and (request.num_stale_output_tokens == 0 or request.drop_stale_output) + for request in candidates ) def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: @@ -1373,6 +1389,32 @@ def schedule_running_requests( external_boundary_hit_tokens = 0 # Get already-cached tokens. + if ( + request.num_computed_tokens == 0 + and self.connector is not None + and self.kv_cache_manager.boundary_checkpoints is not None + ): + self.kv_cache_manager.set_external_boundary_admission_context( + self.max_num_running_reqs - num_running, + sum( + self._request_remaining_blocks(req) + for req in self._inflight_prefills + if not ( + self.kv_cache_manager.has_external_boundary_admission( + req.request_id + ) + ) + ), + has_scheduled_reqs=bool(self.running), + ) + if not self.kv_cache_manager.can_admit_external_boundary_request( + request.request_id + ): + request_queue.remove_request(request) + if prefill_interleave_step is not None: + prefill_interleave_step.mark_unavailable(request_id) + step_skipped_waiting.prepend_request(request) + continue if request.num_computed_tokens == 0: if ( self.connector is not None @@ -1753,11 +1795,13 @@ def schedule_running_requests( # manager if request.has_encoder_inputs: self.encoder_cache_manager.free(request) - if not interleaved_local_prefill: + if not interleaved_local_prefill and not ( + self.kv_cache_manager.has_pending_external_boundary_admissions() + ): break request_queue.remove_request(request) - assert prefill_interleave_step is not None - prefill_interleave_step.mark_unavailable(request_id) + if prefill_interleave_step is not None: + prefill_interleave_step.mark_unavailable(request_id) step_skipped_waiting.prepend_request(request) continue From ae02fd76f68a4e1127210133316dd1abb4d767b0 Mon Sep 17 00:00:00 2001 From: Costa Tsaousis Date: Sun, 27 Sep 2026 22:35:20 +0000 Subject: [PATCH 282/365] Link checkpoint admission release note to PR 922 Signed-off-by: Costa Tsaousis --- .lil/changes/vllm-restore-admission.json | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index c6f2c86894e2..1d7fdf96e909 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -3,7 +3,9 @@ "id": "vllm-restore-admission", "category": "fix", "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", - "models": ["all"], + "models": [ + "all" + ], "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", "details": [ "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", @@ -11,6 +13,8 @@ "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." ], - "authors": ["ktsaou"], - "requires": [] + "requires": [], + "pull_requests": [ + 922 + ] } From 64f0d170c73c30a7749c75b0688a8f78c644c3fd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 14:11:47 +0000 Subject: [PATCH 283/365] fix(executor): surface a failed rank during single-rank RPCs A rank that does not own the reply of execute_model, sample_tokens, execute_dummy_batch or take_draft_token_ids swallowed its exception, so when TP1 hit a CUDA OOM mid-step, TP0 waited in the next all-reduce and the engine waited on TP0 until VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (300 s), with the server up but answering nothing. Every rank now reports a failure, and the gather watches the other ranks' queues while it waits on the output rank, so the first failure ends the wait with its rank and error. A reply from a watched rank that is not a failure (stale) is dropped with a warning. The any-order gather no longer sleeps 1 ms between polls: it waits on the oldest pending queue the way a blocking read does (spinning while replies are frequent, parked on the notification otherwise) in 50 ms slices, via MessageQueue.wait_for_message, which never consumes a message. Co-Authored-By: Claude Opus 5.5 --- .../vllm-worker-failure-fast-exit.json | 16 +++ .../test_multiproc_gather_responses.py | 131 +++++++++++++++++- .../device_communicators/shm_broadcast.py | 16 +++ vllm/v1/executor/multiproc_executor.py | 66 +++++++-- 4 files changed, 214 insertions(+), 15 deletions(-) create mode 100644 .lil/changes/vllm-worker-failure-fast-exit.json diff --git a/.lil/changes/vllm-worker-failure-fast-exit.json b/.lil/changes/vllm-worker-failure-fast-exit.json new file mode 100644 index 000000000000..910baccc54dc --- /dev/null +++ b/.lil/changes/vllm-worker-failure-fast-exit.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-worker-failure-fast-exit", + "category": "fix", + "summary": "A GPU worker that fails mid-step stops the engine at once instead of leaving the server hung", + "models": ["all"], + "compatibility": "No user action required. With a container restart policy, the server comes back on its own after such a failure.", + "details": [ + "When one tensor-parallel rank failed inside a step, for example with a CUDA out-of-memory error on TP1, the other ranks waited in the next all-reduce and the engine waited on rank 0. The server stayed up but answered nothing until VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (300 s) ran out.", + "Every rank now reports a failure, including ranks that do not return the step's output, and the engine reads replies from all ranks in any order. The first failure stops the engine within a fraction of a second, and the log names the failing rank and its error.", + "While replies are frequent the wait still spins, as before; when idle it parks on the queue notification, so steps do not get slower." + ], + "pull_requests": [], + "authors": ["tobymao", "Local Inference Lab"], + "requires": [] +} diff --git a/tests/v1/executor/test_multiproc_gather_responses.py b/tests/v1/executor/test_multiproc_gather_responses.py index a6b929f30090..5653520714a8 100644 --- a/tests/v1/executor/test_multiproc_gather_responses.py +++ b/tests/v1/executor/test_multiproc_gather_responses.py @@ -3,7 +3,8 @@ """The executor's gather of every rank's reply does not wait in rank order. A rank that never replies (for example, stuck in a collective another rank -failed out of) must not hide a failure another rank has already reported. +failed out of) must not hide a failure another rank has already reported, +including a rank that does not own the reply of a single-rank RPC. """ import threading @@ -39,6 +40,9 @@ def dequeue(self, timeout=None): reply, self.reply = self.reply, None return reply + def wait_for_message(self, timeout_ms: int) -> None: + self.waits = getattr(self, "waits", 0) + 1 + def test_a_later_rank_failure_surfaces_while_rank_zero_hangs(): queues = [_Queue(), _Queue((FAILURE, "capture invalidated")), _Queue()] @@ -97,6 +101,131 @@ def send() -> None: reader.shutdown() +def _queue_pair(max_chunk_bytes: int = 1024) -> tuple[MessageQueue, MessageQueue]: + writer = MessageQueue( + n_reader=1, n_local_reader=1, max_chunk_bytes=max_chunk_bytes, max_chunks=2 + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + writer.wait_until_ready() + reader.wait_until_ready() + return writer, reader + + +def test_a_watched_rank_failure_ends_a_single_rank_wait(): + """TP1 fails out of execute_model while TP0, the output rank, waits in a + collective TP1 left: the engine must not wait out the RPC timeout.""" + output = _Queue() + watched = [(1, _Queue((FAILURE, "CUDA out of memory"), ready_after=3))] + started = time.monotonic() + with pytest.raises(RuntimeError, match="Worker 1 failed.*out of memory"): + _gather_responses([output], (0,), time.monotonic() + 300, "m", watched) + assert time.monotonic() - started < 1.0 + assert output.waits >= 3 + + +def test_a_single_rank_reply_is_read_while_watching(): + output = _Queue((SUCCESS, "r0"), ready_after=2) + quiet = _Queue() + assert _gather_responses([output], (0,), None, "m", [(1, quiet)]) == ["r0"] + assert quiet.checks >= 2 and quiet.reply is None + + +def test_an_unexpected_success_from_a_watched_rank_is_dropped(): + stale = _Queue((SUCCESS, "late")) + output = _Queue((SUCCESS, "r0"), ready_after=1) + assert _gather_responses([output], (0,), None, "m", [(1, stale)]) == ["r0"] + assert stale.reply is None + + +def test_the_deadline_applies_while_watching(): + with pytest.raises(TimeoutError, match="RPC call to m timed out"): + _gather_responses( + [_Queue()], (0,), time.monotonic() + 0.2, "m", [(1, _Queue())] + ) + + +def test_a_watched_failure_on_real_queues_wakes_the_parked_gather(): + """Real queues with the output rank idle: the gather parks on the output + rank's notification in slices and still sees the other rank's failure.""" + pairs = [_queue_pair(), _queue_pair()] + try: + + def fail() -> None: + time.sleep(0.2) + pairs[1][0].enqueue((FAILURE, "CUDA out of memory")) + + failer = threading.Thread(target=fail) + started = time.monotonic() + failer.start() + with pytest.raises(RuntimeError, match="Worker 1 failed"): + _gather_responses( + [pairs[0][1]], (0,), time.monotonic() + 30, "m", [(1, pairs[1][1])] + ) + failer.join() + assert time.monotonic() - started < 2.0 + finally: + for writer, reader in pairs: + writer.shutdown() + reader.shutdown() + + +def test_wait_for_message_does_not_consume(): + writer, reader = _queue_pair() + try: + started = time.monotonic() + reader.wait_for_message(100) + assert not reader.ready() + assert time.monotonic() - started < 5.0 + writer.enqueue((SUCCESS, "r")) + reader.wait_for_message(1000) + assert reader.ready() and reader.ready() + assert reader.dequeue(timeout=1) == (SUCCESS, "r") + assert not reader.ready() + finally: + writer.shutdown() + reader.shutdown() + + +class _Worker: + def fail(self): + raise RuntimeError("CUDA out of memory") + + def ok(self): + return "done" + + +class _ResponseQueue: + def __init__(self) -> None: + self.sent: list = [] + + def enqueue(self, obj) -> None: + self.sent.append(obj) + + +def _worker_proc(rank: int) -> WorkerProc: + proc = object.__new__(WorkerProc) + proc.worker = _Worker() + proc.rank = rank + proc.use_async_scheduling = False + proc.worker_response_mq = _ResponseQueue() + return proc + + +def test_a_rank_that_does_not_own_the_reply_reports_its_failure(): + proc = _worker_proc(rank=1) + proc._execute_worker_rpc(("fail", (), {}, 0)) + assert proc.worker_response_mq.sent == [(FAILURE, "CUDA out of memory")] + + +def test_a_rank_that_does_not_own_the_reply_stays_quiet_on_success(): + proc = _worker_proc(rank=1) + proc._execute_worker_rpc(("ok", (), {}, 0)) + assert proc.worker_response_mq.sent == [] + owner = _worker_proc(rank=0) + owner._execute_worker_rpc(("ok", (), {}, 0)) + assert owner.worker_response_mq.sent == [(SUCCESS, "done")] + + def test_a_single_reply_names_its_rank(): # One reply is one blocking read (safe for any size), so no readiness check. boom = _Queue((FAILURE, "boom"), ready_after=-1) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 49b3ffe4df1a..c930438da2ab 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -909,6 +909,22 @@ def ready(self) -> bool: return bool(self.remote_socket.poll(timeout=0)) raise RuntimeError("Only readers can check for messages") + def wait_for_message(self, timeout_ms: int) -> None: + """Wait at most ``timeout_ms`` for a message, without consuming one. + + Waits the way a blocking read does, spinning while reads are frequent + and parking on the writer's notification otherwise, and may return + early, so callers check ``ready()`` again afterwards. + """ + if self._is_local_reader: + self._spin_condition.wait(timeout_ms=timeout_ms) + elif self._is_remote_reader: + self.remote_socket.poll(timeout=timeout_ms) + else: + raise RuntimeError("Only readers can wait for messages") + if self.shutting_down: + raise RuntimeError("cancelled") + def dequeue( self, timeout: float | None = None, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index b22acc8f6fe4..5f1f1483fb51 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -80,8 +80,9 @@ logger = init_logger(__name__) -# How often a gather of every rank's reply checks the pending ranks. -RESPONSE_POLL_SECONDS = 0.001 +# How long a gather parks on one queue's notification before checking every +# other queue it waits on again. +RESPONSE_WAIT_SLICE_MS = 50 def _check_response(rank: int, status: Any, result: Any) -> Any: @@ -98,6 +99,7 @@ def _gather_responses( ranks: Sequence[int], deadline: float | None, method: str, + watched: Sequence[tuple[int, MessageQueue]] = (), ) -> list[Any]: """Collect one reply per queue, failing as soon as any rank reports one. @@ -105,9 +107,16 @@ def _gather_responses( while an earlier rank waited forever, for example in a collective the failed rank would never join: TP1's CUDA graph capture fails, TP0 hangs in the next all-reduce, and the engine waits on TP0 with TP1's error in - its queue. Every pending queue is checked without consuming anything, each - reply is read as it arrives, and the gather sleeps 1 ms between rounds in - which nothing arrived. + its queue. ``watched`` are the other ranks of a single-rank RPC such as + ``execute_model``: they reply only when they fail, and the output rank is + then blocked in a collective the failed rank has left, so a message from + one of them ends the wait with its error. + + Every queue is checked without consuming anything, and a reply is read + only once it has arrived, then with the RPC's own deadline. Between checks + the gather waits on the oldest pending queue the way a blocking read does: + spinning while replies are frequent, otherwise parked on its notification + for at most RESPONSE_WAIT_SLICE_MS. """ def remaining() -> float | None: @@ -118,7 +127,7 @@ def remaining() -> float | None: raise TimeoutError(f"RPC call to {method} timed out.") return left - if len(response_mqs) == 1: + if len(response_mqs) == 1 and not watched: try: status, result = response_mqs[0].dequeue(timeout=remaining()) except TimeoutError as e: @@ -127,8 +136,8 @@ def remaining() -> float | None: responses: list[Any] = [None] * len(response_mqs) pending = list(range(len(response_mqs))) - # A blocking read logs the long-wait message every interval; the polling - # here would not, and operators watch for it, so log it the same way. + # A blocking read logs the long-wait message every interval; the waits + # here are sliced, so log it the same way, since operators watch for it. started, warnings = time.monotonic(), 0 interval = envs.VLLM_RINGBUFFER_WARNING_INTERVAL while pending: @@ -145,12 +154,29 @@ def remaining() -> float | None: responses[index] = _check_response(ranks[index], status, result) pending.remove(index) progressed = True + for rank, mq in watched: + if not mq.ready(): + continue + try: + status, result = mq.dequeue(timeout=remaining()) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e + _check_response(rank, status, result) + logger.warning( + "Dropped a reply from rank %d to %s, which only rank %d answers.", + rank, + method, + ranks[0], + ) if pending and not progressed: - remaining() # raises once the deadline has passed + left = remaining() # raises once the deadline has passed if time.monotonic() - started >= interval * (warnings + 1): warnings += 1 logger.info(LONG_WAIT_TIME_LOG_MSG, interval) - time.sleep(RESPONSE_POLL_SECONDS) + wait_ms = RESPONSE_WAIT_SLICE_MS + if left is not None: + wait_ms = max(1, min(wait_ms, int(left * 1000))) + response_mqs[pending[0]].wait_for_message(wait_ms) return responses @@ -512,12 +538,22 @@ def _aggregate(outputs: Any) -> Any: response_mqs: Sequence[MessageQueue] = self.response_mqs ranks: Sequence[int] = range(len(response_mqs)) + watched: Sequence[tuple[int, MessageQueue]] = () if output_rank is not None: + # The other ranks reply only if they fail; watch them so a failure + # does not wait out the RPC timeout behind the output rank. + watched = [ + (rank, mq) + for rank, mq in enumerate(response_mqs) + if rank != output_rank + ] response_mqs = (response_mqs[output_rank],) ranks = (output_rank,) def get_response(): - responses = _gather_responses(response_mqs, ranks, deadline, str(method)) + responses = _gather_responses( + response_mqs, ranks, deadline, str(method), watched + ) return responses[0] if output_rank is not None else responses future = FutureWrapper( @@ -1135,9 +1171,11 @@ def _execute_worker_rpc( e.add_note(traceback.format_exc()) logger.exception("WorkerProc hit an exception.") # enqueue_output converts the exception to a FAILURE response - # containing its string representation before transport. - if output_rank is None or self.rank == output_rank: - self.handle_output(e) + # containing its string representation before transport. A rank + # that does not own the reply sends it too: its peers may be + # blocked in a collective it has left, so the output rank would + # never answer, and the executor watches every rank for this. + self.handle_output(e) @staticmethod def setup_proc_title_and_log_prefix(enable_ep: bool) -> None: From a8ae18928e07c97b535e70b01a2d486e292e881d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 14:34:48 +0000 Subject: [PATCH 284/365] fix(serve): exit with 1 when the server stops because its engine failed A crashed engine ended vllm serve with status 0, since the API server shuts down gracefully on EngineDeadError. Supervisors and restart policies such as Docker's on-failure then treated the crash as a clean stop and left the server down. BackgroundResources now records engine_failed only when an engine reported its death or its process exited unexpectedly (engine_dead is also set by an orderly shutdown), and the API server raises SystemExit(1) after its shutdown completes if it is set. Co-Authored-By: Claude Opus 5.5 --- .../launchers/test_engine_failure_exit.py | 40 +++++++++++++++++++ .../entrypoints/launchers/api_server/entry.py | 10 +++++ vllm/v1/engine/core_client.py | 6 +++ 3 files changed, 56 insertions(+) create mode 100644 tests/entrypoints/launchers/test_engine_failure_exit.py diff --git a/tests/entrypoints/launchers/test_engine_failure_exit.py b/tests/entrypoints/launchers/test_engine_failure_exit.py new file mode 100644 index 000000000000..cc36bdf48f62 --- /dev/null +++ b/tests/entrypoints/launchers/test_engine_failure_exit.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A server that stops because its engine died exits with a failure status. + +Restart policies and supervisors (Docker ``on-failure``, systemd) only act on +a non-zero exit; a crashed engine used to end ``vllm serve`` with status 0. +""" + +from types import SimpleNamespace + +import pytest +import zmq + +from vllm.entrypoints.launchers.api_server.entry import exit_if_engine_failed +from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.engine.core_client import BackgroundResources +from vllm.v1.engine.exceptions import EngineDeadError + + +def _client(resources: BackgroundResources) -> SimpleNamespace: + return SimpleNamespace(engine_core=SimpleNamespace(resources=resources)) + + +def test_a_reported_engine_death_is_a_failure(): + resources = BackgroundResources(ctx=zmq.Context()) + with pytest.raises(EngineDeadError): + resources.validate_alive([zmq.Frame(EngineCoreProc.ENGINE_CORE_DEAD)]) + assert resources.engine_dead and resources.engine_failed + with pytest.raises(SystemExit) as exited: + exit_if_engine_failed(_client(resources)) + assert exited.value.code == 1 + + +def test_an_orderly_shutdown_is_not_a_failure(): + resources = BackgroundResources(ctx=zmq.Context()) + resources() # the cleanup an orderly shutdown runs + assert resources.engine_dead and not resources.engine_failed + exit_if_engine_failed(_client(resources)) + # Clients without these resources (for example a remote engine) never fail it. + exit_if_engine_failed(SimpleNamespace()) diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py index 37d7ca381455..7f557b2bd763 100644 --- a/vllm/entrypoints/launchers/api_server/entry.py +++ b/vllm/entrypoints/launchers/api_server/entry.py @@ -199,6 +199,16 @@ async def run_server_worker( await shutdown_task finally: sock.close() + exit_if_engine_failed(engine_client) + + +def exit_if_engine_failed(engine_client) -> None: + """Exit with a failure status when the server stopped because its engine + died, so a supervisor or restart policy sees a crash, not a clean stop.""" + resources = getattr(getattr(engine_client, "engine_core", None), "resources", None) + if getattr(resources, "engine_failed", False): + logger.error("API server stopped because its engine failed; exiting with 1") + raise SystemExit(1) def main(): diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index a976551f8569..34910a543c1b 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -486,6 +486,10 @@ class BackgroundResources: # Set if any of the engines are dead. Here so that the output # processing threads can access it without holding a ref to the client. engine_dead: bool = False + # Set only when an engine died on its own (it reported its death or its + # process exited unexpectedly), not by an orderly shutdown, so the server + # can exit with a failure status that restart policies act on. + engine_failed: bool = False def __call__(self): """Clean up background resources.""" @@ -552,6 +556,7 @@ def close_sockets_and_tasks(): def validate_alive(self, frames: Sequence[zmq.Frame]): if len(frames) == 1 and (frames[0].buffer == EngineCoreProc.ENGINE_CORE_DEAD): self.engine_dead = True + self.engine_failed = True raise EngineDeadError() @@ -808,6 +813,7 @@ def monitor_engine_cores(): if not _self or not _self._finalizer.alive or _self.resources.engine_dead: return _self.resources.engine_dead = True + _self.resources.engine_failed = True logger.warning_once( "[shutdown] MPClient: engine core exited unexpectedly; starting cleanup" ) From eb9fec38611f6fd8d8230bd1872de0655b009086 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 14:49:38 +0000 Subject: [PATCH 285/365] docs(release): record the fast exit on a failed rank Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-928.json | 24 +++++++++++++++++++ .../vllm-worker-failure-fast-exit.json | 16 ------------- 2 files changed, 24 insertions(+), 16 deletions(-) create mode 100644 .lil/changes/vllm-928.json delete mode 100644 .lil/changes/vllm-worker-failure-fast-exit.json diff --git a/.lil/changes/vllm-928.json b/.lil/changes/vllm-928.json new file mode 100644 index 000000000000..a0d2ca50ff57 --- /dev/null +++ b/.lil/changes/vllm-928.json @@ -0,0 +1,24 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-928", + "category": "fix", + "summary": "A GPU worker that fails mid-step stops the engine at once, and the server exits with status 1 so a restart policy can bring it back", + "models": [ + "all" + ], + "compatibility": "No user action required. To have the server restart after such a failure, run the container with a restart policy (for example --restart on-failure).", + "details": [ + "When one tensor-parallel rank failed inside a step, for example with a CUDA out-of-memory error on TP1, its peers could stay blocked in a collective the failed rank never joined. The engine read replies in rank order and only from the output rank, so the server stayed up but answered nothing until VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (300 s) ran out.", + "Every rank now reports its failure, including ranks that do not return the step's output, and the engine reads replies from all ranks in any order (includes #917). The first failure stops the engine within a second, and the log names the failing rank and its error.", + "When the server stops because its engine failed, it now exits with status 1 instead of 0, so Docker's on-failure policy and other supervisors see a crash.", + "Waiting for a step still spins while replies are frequent, so decode speed is unchanged (GLM-5.3-Flash Spark TP2: steps/s within 0.4%)." + ], + "pull_requests": [ + 928 + ], + "authors": [ + "tobymao", + "Local Inference Lab" + ], + "requires": [] +} diff --git a/.lil/changes/vllm-worker-failure-fast-exit.json b/.lil/changes/vllm-worker-failure-fast-exit.json deleted file mode 100644 index 910baccc54dc..000000000000 --- a/.lil/changes/vllm-worker-failure-fast-exit.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "schema": "local-inference-release-change/v1", - "id": "vllm-worker-failure-fast-exit", - "category": "fix", - "summary": "A GPU worker that fails mid-step stops the engine at once instead of leaving the server hung", - "models": ["all"], - "compatibility": "No user action required. With a container restart policy, the server comes back on its own after such a failure.", - "details": [ - "When one tensor-parallel rank failed inside a step, for example with a CUDA out-of-memory error on TP1, the other ranks waited in the next all-reduce and the engine waited on rank 0. The server stayed up but answered nothing until VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (300 s) ran out.", - "Every rank now reports a failure, including ranks that do not return the step's output, and the engine reads replies from all ranks in any order. The first failure stops the engine within a fraction of a second, and the log names the failing rank and its error.", - "While replies are frequent the wait still spins, as before; when idle it parks on the queue notification, so steps do not get slower." - ], - "pull_requests": [], - "authors": ["tobymao", "Local Inference Lab"], - "requires": [] -} From 1d5a74c62a039f860501fb26d54da73c58b38b9f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:08:59 +0000 Subject: [PATCH 286/365] fix(serve): count only a failed engine exit as an engine failure A signal to the whole process group (Ctrl-C on vllm serve, a supervisor's killpg) ends the engine cleanly with status 0 before the API server's orderly shutdown marks it dead; that must not make the server exit with 1. Co-Authored-By: Claude Opus 5.5 --- vllm/v1/engine/core_client.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 34910a543c1b..865da2763515 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -487,8 +487,8 @@ class BackgroundResources: # processing threads can access it without holding a ref to the client. engine_dead: bool = False # Set only when an engine died on its own (it reported its death or its - # process exited unexpectedly), not by an orderly shutdown, so the server - # can exit with a failure status that restart policies act on. + # process exited with a failure status), not by an orderly shutdown, so + # the server can exit with a failure status that restart policies act on. engine_failed: bool = False def __call__(self): @@ -813,7 +813,11 @@ def monitor_engine_cores(): if not _self or not _self._finalizer.alive or _self.resources.engine_dead: return _self.resources.engine_dead = True - _self.resources.engine_failed = True + # A signal to the whole process group (Ctrl-C, a supervisor's + # killpg) also ends the engine, cleanly and with status 0. + _self.resources.engine_failed = ( + getattr(engine_manager, "failed_proc_name", None) is not None + ) logger.warning_once( "[shutdown] MPClient: engine core exited unexpectedly; starting cleanup" ) From 3f94f17e19fcdf5998d2b4afa5bed4db7d3603ae Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:13:56 +0000 Subject: [PATCH 287/365] fix(kv-cache): hold import reservations only against new admissions Follow-up to #922 review findings on the admission reservations: - Import credits gated every allocation, so a running decode at a block boundary could be preempted (often the admitted consumer itself) to protect credits of a restore that had not started. Credits now apply to waiting and preempted admissions only; running requests keep growing and resolve pressure through ordinary preemption. - A capacity waiter blocked every other new admission, including requests the connector does not handle and higher-priority ones. A refused restore now only holds back admissions considered after it in the same scheduler step, and those still run when they fit beside the blocks it waits for. Requests ordered ahead of it are unaffected. - A capacity wait is bounded by VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S (default 60, 0 recomputes immediately). Once it expires, one WARNING records the wait, the reservation keeps returning None, and the scheduler admits the request to recompute its prompt even though the connector still defers it for capacity. - Cancelling an import mid-copy kept its credits and running slot until the copy drained. Only the staged pages must wait for the copy now; credits and the slot are released at once. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 219 ++++++++++++++++++++--- vllm/envs.py | 7 + vllm/v1/core/kv_cache_manager.py | 145 +++++++++++---- vllm/v1/core/sched/scheduler.py | 5 + 4 files changed, 320 insertions(+), 56 deletions(-) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index 6f7c0c4f958f..d8f0df707d69 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -240,23 +240,30 @@ def test_scheduler_preserves_recycling_import_capacity( assert cache.block_pool.get_num_free_blocks() == 63 +def publish(cache, checkpoint): + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + + @pytest.mark.parametrize("cancel_at", ["copying", "ready"]) -def test_external_import_cancellation_releases_only_after_copy_drain(cancel_at): +def test_external_import_cancellation_releases_only_pages_after_copy_drain(cancel_at): + """Cancellation frees the slot and credits at once; pages wait for copies.""" cache = manager() + cache.set_external_boundary_admission_context(1, 0) consumer = request("cancel") free = cache.block_pool.get_num_free_blocks() checkpoint = reserve_import(cache, consumer) assert checkpoint is not None + assert not cache.can_admit_external_boundary_request("other") if cancel_at == "ready": - for rank in range(2): - cache.acknowledge_external_boundary_checkpoint( - checkpoint.checkpoint_id, rank - ) + publish(cache, checkpoint) cache.free(consumer) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.can_admit_external_boundary_request("other") if cancel_at == "copying": assert cache.block_pool.get_num_free_blocks() < free - assert cache.external_boundary_reserved_blocks() > 0 cache.discard_external_boundary_checkpoint(checkpoint.checkpoint_id) + # The connector releases again once the cancelled copy has drained. cache.release_external_boundary_admission(consumer.request_id) assert cache.external_boundary_reserved_blocks() == 0 assert cache.block_pool.get_num_free_blocks() == free @@ -287,28 +294,86 @@ def test_external_import_preflight_does_not_evict_cache_on_execution_pressure(): assert cache.external_boundary_reserved_blocks() == 0 -def test_capacity_waiter_prevents_newer_admissions_from_stealing_released_space(): +def refused_waiter_need(cache, waiter): + """Refuse a restore on a full pool and return the space it waits for.""" + probe = manager() + free = probe.block_pool.get_num_free_blocks() + assert reserve_import(probe, waiter) is not None + need = free - probe.block_pool.get_num_free_blocks() + need += probe.external_boundary_reserved_blocks() + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert reserve_import(cache, waiter) is None + return pressure, need + + +@pytest.mark.parametrize("released", ["waiter_only", "both"]) +def test_capacity_waiter_holds_back_only_admissions_that_would_take_its_space( + released, +): cache = manager() older, newer = request("older"), request("newer", "newer") - pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) - assert reserve_import(cache, older) is None - cache.block_pool.free_blocks(pressure) - assert not cache.can_admit_external_boundary_request(newer.request_id) - assert reserve_import(cache, newer) is None - assert reserve_import(cache, older) is not None + pressure, need = refused_waiter_need(cache, older) + cache.block_pool.free_blocks(pressure[: need * (2 if released == "both" else 1)]) assert cache.can_admit_external_boundary_request(newer.request_id) + # A later restore may use only the space beyond the waiter's need. + assert (reserve_import(cache, newer) is not None) == (released == "both") + assert reserve_import(cache, older) is not None + + +def test_capacity_waiter_does_not_hold_back_requests_considered_before_it(): + cache = manager() + older, newer = request("older"), request("newer", "newer") + pressure, need = refused_waiter_need(cache, older) + cache.new_step_starts() + cache.block_pool.free_blocks(pressure[:need]) + assert reserve_import(cache, newer) is not None + assert reserve_import(cache, older) is None def test_cancelled_capacity_waiter_releases_admission_barrier(): cache = manager() - consumer = request("cancel-waiter") - cache.set_external_boundary_admission_context( - 1, cache.block_pool.get_num_free_blocks() - ) - assert reserve_import(cache, consumer) is None - assert not cache.can_admit_external_boundary_request("later") + consumer, later = request("cancel-waiter"), request("later", "later") + pressure, need = refused_waiter_need(cache, consumer) + cache.block_pool.free_blocks(pressure[:need]) cache.release_external_boundary_admission(consumer.request_id) - assert cache.can_admit_external_boundary_request("later") + assert reserve_import(cache, later) is not None + + +def bounded_wait_clock(monkeypatch, max_wait): + """Drive restore waits from a fake clock and capture the manager's warnings.""" + import time + + from vllm.v1.core import kv_cache_manager + + clock = [100.0] + warnings = [] + monkeypatch.setattr(time, "monotonic", lambda: clock[0]) + monkeypatch.setattr( + kv_cache_manager.logger, "warning", lambda *args: warnings.append(args) + ) + monkeypatch.setenv("VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", str(max_wait)) + return clock, warnings + + +def test_capacity_waiter_stops_waiting_after_bounded_wait(monkeypatch): + clock, warnings = bounded_wait_clock(monkeypatch, 30) + cache = manager() + waiter, later = request("waiter"), request("later", "later") + pressure, need = refused_waiter_need(cache, waiter) + clock[0] += 29.0 + assert reserve_import(cache, waiter) is None + assert not warnings and not cache.external_boundary_wait_expired("waiter") + clock[0] += 1.5 + assert reserve_import(cache, waiter) is None + assert cache.external_boundary_wait_expired(waiter.request_id) + assert len(warnings) == 1 and warnings[0][2] == pytest.approx(30.5) + # The request recomputes: it neither restores nor holds later requests back. + cache.block_pool.free_blocks(pressure[:need]) + assert reserve_import(cache, waiter) is None + assert reserve_import(cache, later) is not None + assert len(warnings) == 1 + cache.release_external_boundary_admission(waiter.request_id) + assert not cache.external_boundary_wait_expired(waiter.request_id) def test_import_preflight_preserves_watermark_and_classifies_impossible_restore(): @@ -670,7 +735,7 @@ def poll(req): scheduler.connector = connector scheduler.add_request(consumer) assert scheduler.schedule().num_scheduled_tokens == {} - assert not cache.can_admit_external_boundary_request(session.request_id) + assert consumer.request_id in cache._boundary_import_waiters (update,) = create_requests(num_requests=1, num_tokens=1, req_ids=["stream"]) update.resumable = True scheduler.add_request(update) @@ -678,6 +743,118 @@ def poll(req): assert scheduler.schedule().num_scheduled_tokens == {"stream": 1} +def import_scheduler(**options): + """A real scheduler with checkpoint caching and a restore connector mock. + + Tests install the connector after any setup step that processes output. + """ + from unittest.mock import Mock + + scheduler = make_scheduler( + **{ + "enable_prefix_caching": True, + "use_v2_model_runner": True, + "async_scheduling": True, + **options, + } + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + connector = Mock() + connector.poll_boundary_checkpoint.return_value = True + connector.boundary_checkpoint_external_tokens.return_value = 0 + connector.get_num_new_matched_tokens.return_value = (0, False) + connector.has_pending_block_frees.return_value = False + return scheduler, cache, connector + + +def start_decode(scheduler, name="running"): + from tests.v1.core.utils import create_requests + + (running,) = create_requests(num_requests=1, num_tokens=16, req_ids=[name]) + scheduler.add_request(running) + assert scheduler.schedule().num_scheduled_tokens == {name: 16} + return running + + +def test_unadmitted_import_credits_do_not_preempt_running_decode(): + """Reservations gate new admissions, never the growth of running requests.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=64) + checkpoint = reserve_import(cache, importer, prefix=16) + assert checkpoint is not None + publish(cache, checkpoint) + credits = cache.external_boundary_reserved_blocks() + assert credits > 0 + # The decode crosses a block boundary with only the import's credits free. + cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks() - credits) + scheduler.connector = connector + scheduler.add_request(importer) + output = scheduler.schedule() + assert output.num_scheduled_tokens["running"] == 1 + assert running.status == RequestStatus.RUNNING + assert running.num_preemptions == 0 + + +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +def test_capacity_waiter_does_not_block_requests_ahead_of_it(policy): + """A restore waiting for capacity cannot overtake requests ordered before it.""" + scheduler, cache, connector = import_scheduler( + num_blocks=64, scheduling_policy=policy + ) + earlier = request("earlier", "earlier", length=16) + waiter = request("waiter", "waiter", length=400) + lookup_done = [False] + + def poll(req): + if req is earlier: + return lookup_done[0] + assert reserve_import(cache, req, prefix=384) is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + scheduler.connector = connector + scheduler.add_request(earlier) + scheduler.add_request(waiter) + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert scheduler.schedule().num_scheduled_tokens == {} + assert waiter.request_id in cache._boundary_import_waiters + # Space returns for the earlier request, but not for the restore. + cache.block_pool.free_blocks(pressure[:16]) + lookup_done[0] = True + assert scheduler.schedule().num_scheduled_tokens == {"earlier": 16} + assert waiter.status == RequestStatus.WAITING + + +def test_restore_past_its_bounded_wait_is_admitted_to_recompute(monkeypatch): + """The connector keeps deferring the refused restore; the scheduler admits it.""" + clock, warnings = bounded_wait_clock(monkeypatch, 30) + scheduler, cache, connector = import_scheduler( + num_blocks=64, long_prefill_token_threshold=64 + ) + # A first chunk fits where the whole restore does not. + scheduler.scheduler_reserve_full_isl = False + waiter = request("waiter", "waiter", length=400) + + def poll(req): + assert reserve_import(cache, req, prefix=384) is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + scheduler.connector = connector + scheduler.add_request(waiter) + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert scheduler.schedule().num_scheduled_tokens == {} + cache.block_pool.free_blocks(pressure[:20]) + assert scheduler.schedule().num_scheduled_tokens == {} + clock[0] += 31.0 + output = scheduler.schedule() + assert output.num_scheduled_tokens == {"waiter": 64} + assert len(warnings) == 1 and warnings[0][2] == pytest.approx(31.0) + assert not cache.external_boundary_wait_expired(waiter.request_id) + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/envs.py b/vllm/envs.py index 8052fb1d502b..88f38fc34505 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -174,6 +174,7 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS: bool = False + VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S: float = 60.0 VLLM_MLA_DISABLE: bool = False VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" @@ -1508,6 +1509,11 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS": lambda: ( os.getenv("VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS", "0") == "1" ), + # Seconds an external checkpoint restore may wait for GPU capacity before + # its request recomputes the prompt instead. 0 recomputes immediately. + "VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S": lambda: max( + 0.0, float(os.getenv("VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", "60")) + ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), # If set, vLLM will pick up the provided Flash Attention MLA @@ -2531,6 +2537,7 @@ def compile_factors() -> dict[str, object]: "VLLM_LOGGING_COLOR", "VLLM_LOG_STATS_INTERVAL", "VLLM_REQUEST_STALL_WARNING_S", + "VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", "VLLM_LOG_REQUEST_TIMELINE", "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "VLLM_TUNED_CONFIG_FOLDER", diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index a737fa9e6aa1..8ae6a51d5916 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools +import time from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Literal, overload +import vllm.envs as envs from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger from vllm.utils.math_utils import cdiv @@ -57,6 +59,12 @@ class _BoundaryImportAdmission: admitted: bool = False +@dataclass +class _BoundaryImportWaiter: + need: int + since: float + + @dataclass class KVCacheBlocks: """ @@ -263,7 +271,14 @@ def __init__( ) self._boundary_imports: set[int] = set() self._boundary_import_admissions: dict[str, _BoundaryImportAdmission] = {} - self._boundary_import_waiters: dict[str, None] = {} + self._boundary_import_waiters: dict[str, _BoundaryImportWaiter] = {} + # Requests whose restore outwaited its capacity bound. They recompute + # the prompt instead of restoring until they are admitted. + self._boundary_import_expired: set[str] = set() + # The first restore refused for capacity in this scheduler step and + # the blocks it needs. Admissions considered after it must leave them. + self._boundary_import_barrier: tuple[str, int] | None = None + self._boundary_import_max_wait = envs.VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S self._external_admission_slots: int | None = None self._external_admission_reserved_blocks = 0 self._external_admission_has_scheduled_reqs = False @@ -639,11 +654,17 @@ def allocate_slots( boundary_blocks += len(boundary_replay_managers) # The watermark is applied to waiting/preempted requests only, and only # when there's at least one request already scheduled. - if has_scheduled_reqs and request.status in ( + new_admission = request.status in ( RequestStatus.WAITING, RequestStatus.PREEMPTED, - ): + ) + if has_scheduled_reqs and new_admission: watermark_blocks = self.watermark_blocks + # Import reservations hold back new admissions only. Running requests + # keep growing and resolve pressure through ordinary preemption. + import_reserve = ( + self._blocks_reserved_for_imports(request) if new_admission else 0 + ) if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. @@ -662,10 +683,7 @@ def allocate_slots( required_blocks = ( num_blocks_to_allocate + watermark_blocks + boundary_blocks ) - if required_blocks > ( - self.block_pool.get_num_free_blocks() - - self.external_boundary_reserved_blocks(request.request_id) - ): + if required_blocks > self.block_pool.get_num_free_blocks() - import_reserve: return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -702,9 +720,7 @@ def allocate_slots( # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. available_blocks = ( - self.block_pool.get_num_free_blocks() - - reserved_blocks - - self.external_boundary_reserved_blocks(request.request_id) + self.block_pool.get_num_free_blocks() - reserved_blocks - import_reserve ) required_blocks = num_blocks_to_allocate + watermark_blocks + boundary_blocks if required_blocks > available_blocks: @@ -854,6 +870,7 @@ def allocate_slots( restoring=False, ) self._boundary_import_waiters.pop(request.request_id, None) + self._boundary_import_expired.discard(request.request_id) # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. @@ -1019,12 +1036,15 @@ def reserve_external_boundary_checkpoint( ordinary KV pool, including one page for target/draft auxiliary state. With reserve_admission, also reserve a running slot and continuation - credits. Successful publication retains the consumer pin until ordinary - allocation takes ownership. The scheduler must refresh admission context - before each candidate; cancellation must release the admission after all - copies drain. + credits that later new admissions must leave free. Successful + publication retains the consumer pin until ordinary allocation takes + ownership. The scheduler must refresh admission context whenever it + admits a request; cancellation releases credits and the slot at once, + while staged pages stay pinned until their copies drain. Returns None for temporary resource pressure or unavailable caching. + A restore refused for VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S keeps getting + None, and the scheduler admits its request to recompute the prompt. Raises ValueError for incompatible geometry or an import whose checkpoint and execution state cannot fit even an otherwise empty pool. """ @@ -1055,7 +1075,10 @@ def reserve_external_boundary_checkpoint( count = sum(len(group) for group in page_positions) + 1 execution_blocks = 0 if reserve_admission: - if request.request_id in self._boundary_import_admissions: + if ( + request.request_id in self._boundary_import_admissions + or request.request_id in self._boundary_import_expired + ): return None if not self.can_admit_external_boundary_request(request.request_id): return None @@ -1076,14 +1099,17 @@ def reserve_external_boundary_checkpoint( ) if count + execution_blocks > self.block_pool.num_gpu_blocks - 1: raise ValueError("Checkpoint and execution state exceed the GPU pool") - reserved = self.external_boundary_reserved_blocks() - if reserve_admission: - reserved += self._external_admission_reserved_blocks + reserved = ( + self._blocks_reserved_for_imports(request) + + self._external_admission_reserved_blocks + ) if self._external_admission_has_scheduled_reqs: reserved += self.watermark_blocks + else: + reserved = self.external_boundary_reserved_blocks() if count + execution_blocks + reserved > self.block_pool.get_num_free_blocks(): if reserve_admission: - self._boundary_import_waiters.setdefault(request.request_id, None) + self._wait_for_import_capacity(request, count + execution_blocks) return None allocation = self.block_pool.get_new_blocks(count) try: @@ -1164,7 +1190,9 @@ def set_external_boundary_admission_context( *, has_scheduled_reqs: bool = False, ) -> None: - """Refresh scheduler capacity immediately before polling a candidate. + """Refresh scheduler capacity before polling candidates. + + The scheduler refreshes it again after every admission in its pass. Args: available_slots: Slots remaining after running and paused sessions. @@ -1187,14 +1215,54 @@ def external_boundary_reserved_blocks(self, owner: str | None = None) -> int: if request_id != owner ) + def _blocks_reserved_for_imports(self, request: Request) -> int: + """Blocks a new admission must leave free for imports ahead of it.""" + reserved = self.external_boundary_reserved_blocks(request.request_id) + barrier = self._boundary_import_barrier + if ( + barrier is not None + and barrier[0] != request.request_id + and request.num_computed_tokens == 0 + and request.request_id not in self._boundary_import_admissions + ): + reserved += barrier[1] + return reserved + + def _wait_for_import_capacity(self, request: Request, need: int) -> None: + """Record a capacity refusal and bound how long the restore may wait.""" + now = time.monotonic() + waiter = self._boundary_import_waiters.get(request.request_id) + if waiter is None: + waiter = _BoundaryImportWaiter(need, now) + self._boundary_import_waiters[request.request_id] = waiter + waited = now - waiter.since + if waited >= self._boundary_import_max_wait: + self._forget_import_waiter(request.request_id) + self._boundary_import_expired.add(request.request_id) + logger.warning( + "Checkpoint restore for request %s waited %.1f s for %d GPU " + "blocks; recomputing its prompt instead", + request.request_id, + waited, + need, + ) + return + waiter.need = need + # Later candidates in this step are behind the waiter; earlier ones, + # and requests considered first in the next step, are not. + if self._boundary_import_barrier is None: + self._boundary_import_barrier = (request.request_id, need) + + def _forget_import_waiter(self, request_id: str) -> None: + self._boundary_import_waiters.pop(request_id, None) + barrier = self._boundary_import_barrier + if barrier is not None and barrier[0] == request_id: + self._boundary_import_barrier = None + def can_admit_external_boundary_request(self, request_id: str) -> bool: """Whether a candidate can use a running slot without stealing a ticket.""" if request_id in self._boundary_import_admissions: return True - if self._boundary_import_waiters and request_id != next( - iter(self._boundary_import_waiters) - ): - return False return ( self._external_admission_slots is None or sum( @@ -1213,19 +1281,25 @@ def external_boundary_admission_ready(self, request_id: str) -> bool: admission = self._boundary_import_admissions.get(request_id) return admission is not None and admission.ready and not admission.admitted + def external_boundary_wait_expired(self, request_id: str) -> bool: + """Whether a restore outwaited its capacity bound, so the prompt recomputes.""" + return request_id in self._boundary_import_expired + def release_external_boundary_admission(self, request_id: str) -> None: - """Release a finished consumer's credits; in-flight copies must drain first.""" - self._boundary_import_waiters.pop(request_id, None) - admission = self._boundary_import_admissions.get(request_id) - if ( - admission is None - or admission.checkpoint.checkpoint_id in self._boundary_imports - ): + """Release an import's credits and slot, or end its capacity wait. + + A copy that is still in flight keeps only its staged pages pinned, until + the connector acknowledges or discards it. + """ + self._forget_import_waiter(request_id) + self._boundary_import_expired.discard(request_id) + admission = self._boundary_import_admissions.pop(request_id, None) + if admission is None or not admission.ready or admission.admitted: return - del self._boundary_import_admissions[request_id] - if admission.ready and not admission.admitted: - assert self.boundary_checkpoints is not None - self.boundary_checkpoints.release(admission.checkpoint) + assert self.boundary_checkpoints is not None + self.boundary_checkpoints.release(admission.checkpoint) + if admission.request.boundary_checkpoint is admission.checkpoint: + admission.request.boundary_checkpoint = None def _external_boundary_execution_blocks( self, @@ -1655,4 +1729,5 @@ def finalize_partial_tail_offloads( def new_step_starts(self) -> None: """Notify the coordinator that a new step is starting.""" + self._boundary_import_barrier = None self.coordinator.new_step_starts() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a61b5fac1c9d..35100581b80d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1420,6 +1420,11 @@ def schedule_running_requests( self.connector is not None and self.kv_cache_manager.boundary_checkpoints is not None and not self.connector.poll_boundary_checkpoint(request) + # The connector still defers a restore that outwaited + # its capacity bound; the request recomputes instead. + and not self.kv_cache_manager.external_boundary_wait_expired( + request_id + ) ): request_queue.remove_request(request) if prefill_interleave_step is not None: From 0dd394dfb15dfd0f2b4ce78d9558dc13acf287a7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:05 +0000 Subject: [PATCH 288/365] fix(scheduler): keep queue order and step cost while imports are pending Follow-up to #922 review findings on the waiting pass: - A failed allocation no longer stopped the waiting pass while any import was pending, so every later request could overtake a blocked queue head (FCFS and priority alike). The first failure stops ordinary admission again; only restored imports, which already own their capacity and slot, are still admitted past it. - With an import pending, every step walked the whole waiting queue (connector poll, checkpoint lookup and allocation per request), the saved-logits check scanned both queues (copying the priority heap), and the admission context summed in-flight prefills per candidate. Ready imports are now found through the allocator's admissions, the context is refreshed only after admissions, and priority queues test membership without the ordered heap copy. With 500 blocked 2K-token requests a step costs the same as without imports (0.06 ms instead of 15.7 ms FCFS / 28.3 ms priority). - A ready full-prefix import anywhere in the queue suppressed the running pass; if an older request was admitted instead, decodes lost the step. The import now gets its isolated saved-logits step first, and the running pass runs, without preempting new admissions, whenever that step is not taken. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 97 +++++++++ .../test_boundary_admission_continuation.py | 2 +- vllm/v1/core/kv_cache_manager.py | 8 + vllm/v1/core/sched/request_queue.py | 4 + vllm/v1/core/sched/scheduler.py | 206 ++++++++++++------ 5 files changed, 247 insertions(+), 70 deletions(-) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index d8f0df707d69..cbf1adf902ab 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -777,6 +777,103 @@ def start_decode(scheduler, name="running"): return running +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +@pytest.mark.parametrize("state", ["copying", "ready"]) +def test_blocked_queue_head_admits_only_restored_imports_past_it(policy, state): + """Ordinary requests keep their queue order while imports are pending.""" + from tests.v1.core.utils import create_requests + + scheduler, cache, connector = import_scheduler( + num_blocks=64, scheduling_policy=policy + ) + importer = request("import", "import", length=48) + checkpoint = reserve_import(cache, importer, prefix=32) + assert checkpoint is not None + if state == "ready": + publish(cache, checkpoint) + (big,) = create_requests(num_requests=1, num_tokens=2000, req_ids=["big"]) + small = request("small", "small", length=16) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer or state == "ready" + ) + scheduler.connector = connector + for req in (big, small, importer): + scheduler.add_request(req) + expected = {"import": 16} if state == "ready" else {} + assert scheduler.schedule().num_scheduled_tokens == expected + assert small.status == big.status == RequestStatus.WAITING + + +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +def test_pending_import_step_cost_does_not_scale_with_blocked_queue(policy): + """A copying import leaves a blocked 500-request queue at one lookup a step.""" + import time + + from tests.v1.core.utils import create_requests + + def blocked(with_import): + scheduler, cache, connector = import_scheduler( + num_blocks=256, scheduling_policy=policy + ) + queued = create_requests(num_requests=500, num_tokens=2048) + importer = request("import", "import", length=64) + if with_import: + assert reserve_import(cache, importer, prefix=32) is not None + cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer + ) + scheduler.connector = connector + for req in (*queued, importer): + scheduler.add_request(req) + assert scheduler.schedule().num_scheduled_tokens == {} + connector.poll_boundary_checkpoint.reset_mock() + steps = 20 + start = time.perf_counter() + for _ in range(steps): + assert scheduler.schedule().num_scheduled_tokens == {} + elapsed = (time.perf_counter() - start) / steps + return elapsed, connector.poll_boundary_checkpoint.call_count / steps + + baseline, baseline_polls = blocked(with_import=False) + elapsed, polls = blocked(with_import=True) + assert polls == baseline_polls == 1 + assert elapsed < 10 * baseline + 1e-3 + + +@pytest.mark.parametrize("state", ["copying", "ready", "ready_but_skipped"]) +def test_ready_import_behind_older_request_keeps_running_decode(state): + """Running decodes skip a step only for a saved-logits step actually taken.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=32) + checkpoint = reserve_import(cache, importer, prefix=32) + assert checkpoint is not None + if state != "copying": + publish(cache, checkpoint) + older = request("older", "older", length=16) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer or state == "ready" + ) + connector.boundary_checkpoint_external_tokens.side_effect = lambda req: ( + 32 if req is importer else 0 + ) + scheduler.connector = connector + scheduler.add_request(older) + scheduler.add_request(importer) + output = scheduler.schedule() + if state == "ready": + assert output.boundary_logits_only + assert output.num_scheduled_tokens == {"import": 1} + output = scheduler.schedule() + assert output.num_scheduled_tokens["running"] == 1 + assert output.num_scheduled_tokens["older"] == 16 + else: + assert not output.boundary_logits_only + assert output.num_scheduled_tokens == {"running": 1, "older": 16} + assert running.num_preemptions == 0 + + def test_unadmitted_import_credits_do_not_preempt_running_decode(): """Reservations gate new admissions, never the growth of running requests.""" scheduler, cache, connector = import_scheduler() diff --git a/tests/v1/core/test_boundary_admission_continuation.py b/tests/v1/core/test_boundary_admission_continuation.py index 6e929a73f571..46b9d4c82b60 100644 --- a/tests/v1/core/test_boundary_admission_continuation.py +++ b/tests/v1/core/test_boundary_admission_continuation.py @@ -111,7 +111,7 @@ def test_nonisolated_continuation_deferral_serves_existing_decoder( ) base.queue_matching(scheduler, unrelated) base.victim_at_head(cache, checkpoint) - assert not scheduler._has_waiting_boundary_logits() + assert scheduler._waiting_boundary_logits_request() is None output = scheduler.schedule() assert not output.boundary_logits_only assert output.num_scheduled_tokens == {"first": 4} diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 8ae6a51d5916..4e99e1b4a82f 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -1285,6 +1285,14 @@ def external_boundary_wait_expired(self, request_id: str) -> bool: """Whether a restore outwaited its capacity bound, so the prompt recomputes.""" return request_id in self._boundary_import_expired + def ready_external_boundary_requests(self) -> list[Request]: + """Waiting consumers whose imports finished copying, in reservation order.""" + return [ + admission.request + for admission in self._boundary_import_admissions.values() + if admission.ready and not admission.admitted + ] + def release_external_boundary_admission(self, request_id: str) -> None: """Release an import's credits and slot, or end its capacity wait. diff --git a/vllm/v1/core/sched/request_queue.py b/vllm/v1/core/sched/request_queue.py index 38c7db94a45e..6bf56500d694 100644 --- a/vllm/v1/core/sched/request_queue.py +++ b/vllm/v1/core/sched/request_queue.py @@ -191,6 +191,10 @@ def __len__(self) -> int: """Get number of requests in queue.""" return len(self._heap) + def __contains__(self, request: object) -> bool: + """Check membership without the ordered iteration's heap copy.""" + return request in self._heap + def __iter__(self) -> Iterator[Request]: """Iterate over the queue according to priority policy.""" heap_copy = self._heap[:] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 35100581b80d..216e968629a3 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -56,6 +56,7 @@ ) from vllm.v1.core.sched.prefill_interleave import ( PrefillInterleaveController, + PrefillInterleaveStep, resolve_decode_refill_target, resolve_max_parallel_prefills, ) @@ -770,44 +771,68 @@ def _local_prefill_pressure(self, now: float) -> tuple[float, int]: pressure = max(pressure, 1.0 + age_seconds / expected_seconds) return pressure, backlog_tokens - def _has_waiting_boundary_logits(self) -> bool: - """Whether the queue head needs an isolated saved-logits step.""" + def _waiting_boundary_logits_request(self) -> Request | None: + """Return the waiting request that needs an isolated saved-logits step. + + This is the queue head or, failing that, a restored import that already + owns its reserved capacity. + """ if ( not (self.waiting or self.skipped_waiting) or self._pause_state != PauseState.UNPAUSED or len(self.running) + self.num_waiting_for_streaming_input >= self.max_num_running_reqs ): - return False + return None queue = self._select_waiting_queue_for_scheduling() if queue is None: - return False - pending_imports = ( - self.kv_cache_manager.has_pending_external_boundary_admissions() - ) - candidates = itertools.chain( - (queue.peek_request(),), + return None + return next( ( request for request in ( - itertools.chain(self.skipped_waiting, self.waiting) - if pending_imports - else () - ) - if self.kv_cache_manager.external_boundary_admission_ready( - request.request_id + queue.peek_request(), + *self._ready_import_requests(), ) + if request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) + and request.num_computed_tokens == 0 + and request.boundary_checkpoint is not None + and request.boundary_checkpoint.num_tokens == request.num_tokens + and (request.num_stale_output_tokens == 0 or request.drop_stale_output) ), - ) - return any( - request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) - and request.num_computed_tokens == 0 - and request.boundary_checkpoint is not None - and request.boundary_checkpoint.num_tokens == request.num_tokens - and (request.num_stale_output_tokens == 0 or request.drop_stale_output) - for request in candidates + None, ) + def _ready_import_requests(self) -> list[Request]: + """Restored imports that wait for admission, in scheduling order.""" + requests = self.kv_cache_manager.ready_external_boundary_requests() + if self.policy == SchedulingPolicy.PRIORITY: + requests.sort() + return requests + + def _ready_imports_to_consider( + self, + considered_ids: set[str], + prefill_interleave_step: PrefillInterleaveStep | None, + ) -> list[Request]: + """Restored imports not yet considered in this step's waiting pass.""" + return [ + request + for request in self._ready_import_requests() + if request.request_id not in considered_ids + and ( + prefill_interleave_step is None + or prefill_interleave_step.is_selected(request.request_id) + ) + ] + + def _waiting_queue_of(self, request: Request) -> RequestQueue | None: + """Return the waiting queue that holds the request, if any.""" + for queue in (self.skipped_waiting, self.waiting): + if request in queue: + return queue + return None + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: @@ -1293,7 +1318,9 @@ def schedule_running_requests( # A full recurrent-cache hit cannot share a model step with forwards. # Its cache lookup in the waiting pass records that requirement. Give # it an isolated admission step instead of waiting for decodes to end. - reserve_boundary_logits_step = self._has_waiting_boundary_logits() + boundary_logits_request = self._waiting_boundary_logits_request() + reserve_boundary_logits_step = boundary_logits_request is not None + boundary_logits_taken = False if not reserve_boundary_logits_step: schedule_running_requests(initial_service_class) @@ -1310,6 +1337,15 @@ def schedule_running_requests( # Next, schedule the WAITING requests. if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: step_skipped_waiting = create_request_queue(self.policy) + external_admission = ( + self.connector is not None + and self.kv_cache_manager.boundary_checkpoints is not None + ) + admission_context_stale = True + # Set once ordinary admission stops: restored imports already own + # their capacity and slot, so only they may still be admitted. + ready_imports: list[Request] | None = None + considered_ids: set[str] = set() while (self.waiting or self.skipped_waiting) and token_budget > 0: if ( @@ -1323,20 +1359,36 @@ def schedule_running_requests( if num_running >= self.max_num_running_reqs: break - request_queue = self._select_waiting_queue_for_scheduling() - assert request_queue is not None - request = request_queue.peek_request() - if ( - prefill_interleave_step is not None - and self._request_has_local_prefill(request) - ): - waiting_selection = prefill_interleave_step.select_waiting_request( - (self.skipped_waiting, self.waiting) - ) - if waiting_selection is None: + selected: Request | None = None + if boundary_logits_request is not None: + selected, boundary_logits_request = boundary_logits_request, None + elif ready_imports is not None: + if not ready_imports: break - request_queue, request = waiting_selection + selected = ready_imports.pop(0) + if selected is None: + request_queue = self._select_waiting_queue_for_scheduling() + assert request_queue is not None + request = request_queue.peek_request() + if ( + prefill_interleave_step is not None + and self._request_has_local_prefill(request) + ): + waiting_selection = ( + prefill_interleave_step.select_waiting_request( + (self.skipped_waiting, self.waiting) + ) + ) + if waiting_selection is None: + break + request_queue, request = waiting_selection + else: + request_queue = self._waiting_queue_of(selected) + if request_queue is None: + continue + request = selected request_id = request.request_id + considered_ids.add(request_id) # try to promote blocked statuses while traversing skipped queue. if self._is_blocked_waiting_status( @@ -1389,31 +1441,23 @@ def schedule_running_requests( external_boundary_hit_tokens = 0 # Get already-cached tokens. - if ( - request.num_computed_tokens == 0 - and self.connector is not None - and self.kv_cache_manager.boundary_checkpoints is not None - ): - self.kv_cache_manager.set_external_boundary_admission_context( - self.max_num_running_reqs - num_running, - sum( - self._request_remaining_blocks(req) - for req in self._inflight_prefills - if not ( - self.kv_cache_manager.has_external_boundary_admission( - req.request_id - ) - ) - ), - has_scheduled_reqs=bool(self.running), - ) + if request.num_computed_tokens == 0 and external_admission: + if admission_context_stale: + # Only admissions change slots and in-flight needs. + self.kv_cache_manager.set_external_boundary_admission_context( + self.max_num_running_reqs - num_running, + self._unticketed_inflight_prefill_blocks(), + has_scheduled_reqs=bool(self.running), + ) + admission_context_stale = False if not self.kv_cache_manager.can_admit_external_boundary_request( - request.request_id + request_id ): - request_queue.remove_request(request) - if prefill_interleave_step is not None: - prefill_interleave_step.mark_unavailable(request_id) - step_skipped_waiting.prepend_request(request) + # Every remaining slot is reserved for a restored import. + if ready_imports is None: + ready_imports = self._ready_imports_to_consider( + considered_ids, prefill_interleave_step + ) continue if request.num_computed_tokens == 0: if ( @@ -1455,7 +1499,9 @@ def schedule_running_requests( and num_new_local_computed_tokens == request.num_tokens ) if boundary_logits_only and num_scheduled_tokens: - break + if ready_imports is None: + break + continue # Get externally-cached tokens if using a KVConnector. if self.connector is not None: @@ -1800,13 +1846,19 @@ def schedule_running_requests( # manager if request.has_encoder_inputs: self.encoder_cache_manager.free(request) - if not interleaved_local_prefill and not ( - self.kv_cache_manager.has_pending_external_boundary_admissions() - ): - break + if not interleaved_local_prefill: + # FCFS stops here; imports that already own their + # capacity cannot take it from this request. + if ready_imports is None: + ready_imports = self._ready_imports_to_consider( + considered_ids, prefill_interleave_step + ) + if not ready_imports: + break + continue request_queue.remove_request(request) - if prefill_interleave_step is not None: - prefill_interleave_step.mark_unavailable(request_id) + assert prefill_interleave_step is not None + prefill_interleave_step.mark_unavailable(request_id) step_skipped_waiting.prepend_request(request) continue @@ -1838,6 +1890,7 @@ def schedule_running_requests( ) request_queue.remove_request(request) + admission_context_stale = True if load_kv_async: # If loading async, allocate memory and put request # into the WAITING_FOR_REMOTE_KV state. @@ -1928,6 +1981,7 @@ def schedule_running_requests( # interleaving policy, including policies that do not # maintain a separate micro-prefill budget. token_budget = 0 + boundary_logits_taken = True break # re-queue requests skipped in this pass ahead of older skipped items. @@ -1939,10 +1993,15 @@ def schedule_running_requests( if not defer_prefills: self.prefill_capacity_bound = bool(self.waiting) - if reserve_boundary_logits_step and not num_scheduled_tokens: + if reserve_boundary_logits_step and not boundary_logits_taken: # Cache eviction, admission limits, or deferred input can prevent - # the isolated step. Running requests must still make progress. - schedule_running_requests(initial_service_class) + # the isolated step. Running requests must still make progress, + # without evicting requests admitted above. + schedule_running_requests( + initial_service_class, + allow_preemption=not num_scheduled_tokens, + enforce_lora_limit=bool(num_scheduled_tokens), + ) if ( prefill_interleave_step is not None @@ -3922,6 +3981,15 @@ def _inflight_prefill_reserved_blocks(self) -> int: self._request_remaining_blocks(req) for req in self._inflight_prefills ) + def _unticketed_inflight_prefill_blocks(self) -> int: + """Blocks in-flight prefills without import credits still need.""" + manager = self.kv_cache_manager + return sum( + self._request_remaining_blocks(req) + for req in self._inflight_prefills + if not manager.has_external_boundary_admission(req.request_id) + ) + def _update_waiting_for_remote_kv(self, request: Request) -> None: """ KV Connector: update request state after async recv is finished. From 7023056e0d7ce9aa978794e7e9affe28ab7f6ddd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:11 +0000 Subject: [PATCH 289/365] fix(scheduler): release restored imports on a forced prefix-cache reset reset_prefix_cache(reset_running_requests=True) preempts every running request so the pool can be cleared, but a restore that finished copying and was not yet admitted still pinned its pages, and the reset raised RuntimeError. pause_scheduler(mode="keep", clear_cache=True) and sleep(level>=1, mode="keep") reach this path. Such restores are now released with the running requests; their requests fall back to an ordinary lookup. A reset without preemption still reports failure while a restored import holds its pages. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 19 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 6 ++++++ 2 files changed, 25 insertions(+) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index cbf1adf902ab..7d73321e1bad 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -952,6 +952,25 @@ def poll(req): assert not cache.external_boundary_wait_expired(waiter.request_id) +def test_reset_with_running_requests_releases_ready_unadmitted_import(): + """A cache reset drops a finished restore; its request then recomputes.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=32) + checkpoint = reserve_import(cache, importer, prefix=16) + assert checkpoint is not None + publish(cache, checkpoint) + scheduler.connector = connector + scheduler.add_request(importer) + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert running.status == RequestStatus.PREEMPTED + assert importer.boundary_checkpoint is None + assert not cache.has_external_boundary_admission(importer.request_id) + assert cache.block_pool.get_num_free_blocks() == cache.block_pool.num_gpu_blocks - 1 + output = scheduler.schedule() + assert output.num_scheduled_tokens["import"] == 32 + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 216e968629a3..9247348da340 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3750,6 +3750,12 @@ def reset_prefix_cache( while self.running: request = self.running.pop() self._preempt_request(request, timestamp, drop_stale_output=True) + # Restores that finished copying but were never admitted hold their + # pages the same way; their requests fall back to ordinary lookup. + for request in self.kv_cache_manager.ready_external_boundary_requests(): + self.kv_cache_manager.release_external_boundary_admission( + request.request_id + ) # Clear scheduled request ids cache. Since we are forcing preemption # + resumption in the same step, we must act as if these requests were From f9082f2223d4a92ecb99f5df8a201c4ca70352cd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:16 +0000 Subject: [PATCH 290/365] docs(release): describe fair, bounded checkpoint-restore admission Update the #922 fragment to the reviewed behavior: reservations hold back only new admissions, queue order and step cost are kept while imports are pending, saved-logits steps no longer cost decodes, cancellation frees credits at once, forced resets release restored imports, and VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S bounds capacity waits. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-restore-admission.json | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index 1d7fdf96e909..2d9b75986d23 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -2,15 +2,18 @@ "schema": "local-inference-release-change/v1", "id": "vllm-restore-admission", "category": "fix", - "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", + "summary": "External recurrent checkpoint restores wait for GPU capacity instead of recomputing cached prefixes, with fair scheduling and a bounded wait.", "models": [ "all" ], - "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", + "compatibility": "Use with the paired LMCache restore-admission change. VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S (default 60) bounds how long a restore may wait for capacity before its request recomputes the prompt; 0 recomputes immediately.", "details": [ - "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", + "Restores reserve GPU blocks and a running slot before copies start, and keep restored pages owned until execution takes over.", + "Reservations hold back only new admissions. Running requests keep growing, requests ordered ahead of a waiting restore are never blocked by it, and later requests are admitted when they fit beside it.", + "Queue order is kept: a request that does not fit still stops ordinary admission, and only restores that already own their capacity are admitted past it, without rescanning the waiting queue each step.", + "A finished full-prefix restore gets its isolated saved-logits step directly, and running decodes skip a step only when that step is taken.", + "Cancelling a restore frees its slot and reserved capacity at once; only the copied pages wait for in-flight transfers. A forced prefix-cache reset releases finished restores that were not yet admitted.", "Sliding-window and chunked-local continuation reservations count physical pages rather than skipped positions in the block table.", - "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." ], "requires": [], From 1ac6c2deb5b5cb1543b75a5e93176ffb7d048d00 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:47:26 +0000 Subject: [PATCH 291/365] docs(release): link the restore-admission fragment to PR 929 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-restore-admission.json | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index 2d9b75986d23..ce684ef3819b 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -18,6 +18,10 @@ ], "requires": [], "pull_requests": [ - 922 + 929 + ], + "authors": [ + "ktsaou", + "Local Inference Lab" ] } From 0928737d9f2effa93cfa2070404910f6dfc30aa7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:26:03 +0000 Subject: [PATCH 292/365] fix(structured-output): keep MTP draft rows constrained after a skipped step With async scheduling the scheduler holds only -1 placeholders for a request's drafts and back-fills them from take_draft_token_ids() before it builds the grammar bitmask, while the V2 model runner verifies the drafts it holds on the GPU. The hand-off returned the drafts proposed by the last sampled step, so a structured-output request that skipped that step (a compute-share prefill turn during concurrent prefills, or the isolated saved-logits step of a checkpoint restore) kept its placeholders. Its step was not even deferred, because it had no output in flight. grammar_bitmask then left every row after the first draft unconstrained and the runner accepted drafts and a bonus token the grammar never allowed: "Failed to advance FSM ... grammar rejected tokens [62, 100928, 23, 62]" (#726). - AsyncScheduler defers sampling for every structured-output request with scheduled drafts (V2 runner), not only for requests with output in flight. - With async scheduling the V2 DraftTokensHandler records the drafts each step's inputs consumed, so the deferred back-fill always sees the drafts the runner verifies, including those of requests that skipped a step. - The scheduler marks draft positions it could not constrain as invalid (num_invalid_spec_tokens), so the runner rejects them instead of sampling from rows that allow every token, whatever the reason the back-fill missed. - A deferred step with no earlier batch in flight is sampled immediately instead of popping an empty batch queue. Co-Authored-By: Claude Opus 5.5 --- .../test_structured_output_draft_handoff.py | 357 ++++++++++++++++++ tests/v1/worker/test_draft_tokens_handler.py | 78 ++++ vllm/v1/core/sched/async_scheduler.py | 27 +- vllm/v1/core/sched/scheduler.py | 39 ++ vllm/v1/engine/core.py | 60 +-- vllm/v1/worker/gpu/model_runner.py | 12 +- vllm/v1/worker/gpu/spec_decode/utils.py | 55 ++- 7 files changed, 596 insertions(+), 32 deletions(-) create mode 100644 tests/v1/engine/test_structured_output_draft_handoff.py create mode 100644 tests/v1/worker/test_draft_tokens_handler.py diff --git a/tests/v1/engine/test_structured_output_draft_handoff.py b/tests/v1/engine/test_structured_output_draft_handoff.py new file mode 100644 index 000000000000..20fd9e8a7d54 --- /dev/null +++ b/tests/v1/engine/test_structured_output_draft_handoff.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured output with async speculative decoding must stay constrained +when a request skips a step. + +The V2 model runner verifies the drafts it holds itself, while the scheduler +builds the grammar bitmask from -1 placeholders that deferred sampling +back-fills. A request that skipped the previous step (here a compute-share +prefill turn) used to be scheduled with placeholders nobody back-filled, so its +bitmask rows past the first draft allowed every token and the runner accepted +tokens the grammar rejected (issue local-inference-lab/vllm#726). +""" + +import itertools +from collections import deque +from concurrent.futures import Future +from contextlib import nullcontext +from dataclasses import dataclass, field +from unittest.mock import Mock, patch + +import numpy as np +import pytest +import torch + +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.engine.core import EngineCore +from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm.v1.structured_output.backend_types import StructuredOutputGrammar + +from ..core.utils import EOS_TOKEN_ID, create_requests, create_scheduler + +pytestmark = pytest.mark.cpu_test + +NUM_SPEC_TOKENS = 3 +VOCAB_SIZE = 128 +# The only output the grammar accepts, one token per position. +GRAMMAR_TOKENS = list(range(10, 34)) +# Token the model prefers at odd positions; the grammar never allows it. +OFF_GRAMMAR_TOKEN = 99 + + +def _preferred_token(position: int) -> int: + if position % 2 or position >= len(GRAMMAR_TOKENS): + return OFF_GRAMMAR_TOKEN + return GRAMMAR_TOKENS[position] + + +@dataclass +class SequenceGrammar(StructuredOutputGrammar): + """Accepts exactly GRAMMAR_TOKENS.""" + + position: int = 0 + + def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: + for token in tokens: + if self.is_terminated() or token != GRAMMAR_TOKENS[self.position]: + return False + self.position += 1 + return True + + def validate_tokens(self, tokens: list[int]) -> list[int]: + valid = list( + itertools.takewhile( + lambda item: item[1] == GRAMMAR_TOKENS[self.position + item[0]], + enumerate(tokens[: len(GRAMMAR_TOKENS) - self.position]), + ) + ) + return [token for _, token in valid] + + def rollback(self, num_tokens: int) -> None: + self.position -= num_tokens + + def fill_bitmask(self, bitmask: torch.Tensor, batch_index: int) -> None: + bitmask[batch_index].zero_() + if not self.is_terminated(): + token = GRAMMAR_TOKENS[self.position] + # Bit 31 is the int32 sign bit. + bitmask[batch_index, token // 32] = int( + np.array(1 << (token % 32), dtype=np.uint32).view(np.int32) + ) + + def is_terminated(self) -> bool: + return self.position >= len(GRAMMAR_TOKENS) + + def reset(self): + self.position = 0 + + +class BitmaskBackend: + def allocate_token_bitmask(self, max_num_seqs: int) -> torch.Tensor: + return torch.full((max_num_seqs, VOCAB_SIZE // 32), -1, dtype=torch.int32) + + def destroy(self): + pass + + +def _allowed(row: np.ndarray, token: int) -> bool: + return bool((int(row[token // 32]) >> (token % 32)) & 1) + + +@dataclass +class _RunnerRequest: + prefill_len: int + num_computed: int + num_output: int = 0 + # Persistent draft slot, zeroed on admission like RequestState.draft_tokens. + drafts: list[int] = field(default_factory=lambda: [0] * NUM_SPEC_TOKENS) + + +class FakeV2Executor: + """Emulates the V2 runner's draft ownership with a greedy model. + + The model prefers ``_preferred_token``; grammar rows constrain it. Drafts + are verified from the runner's own slot, never from the scheduler's + placeholders, and invalidated per ``num_invalid_spec_tokens`` exactly as + ``_build_grammar_mapping`` does. ``consumed_handoff`` selects what + ``take_draft_token_ids`` returns: the drafts the latest step consumed, or + (the previous contract) the drafts the latest sampled step proposed. + """ + + def __init__(self, consumed_handoff: bool): + self.consumed_handoff = consumed_handoff + self.requests: dict[str, _RunnerRequest] = {} + self.pending: tuple[SchedulerOutput, dict[str, int]] | None = None + self.consumed = DraftTokenIds([], []) + self.proposed = DraftTokenIds([], []) + # Per sampled step: request id -> number of accepted drafts. + self.accepted_per_step: list[dict[str, int]] = [] + + @staticmethod + def _done(value) -> Future: + future: Future = Future() + future.set_result(value) + return future + + def execute_model(self, scheduler_output: SchedulerOutput, non_block=False): + for req_id in itertools.chain( + scheduler_output.finished_req_ids, scheduler_output.preempted_req_ids or () + ): + self.requests.pop(req_id, None) + for new_req in scheduler_output.scheduled_new_reqs: + assert new_req.prefill_token_ids is not None + self.requests[new_req.req_id] = _RunnerRequest( + len(new_req.prefill_token_ids), new_req.num_computed_tokens + ) + num_drafts = { + req_id: len(scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())) + for req_id in scheduler_output.num_scheduled_tokens + } + drafting = [req_id for req_id, n in num_drafts.items() if n] + self.consumed = DraftTokenIds( + drafting, + [self.requests[req_id].drafts[: num_drafts[req_id]] for req_id in drafting], + ) + if scheduler_output.total_num_scheduled_tokens == 0: + return self._done(ModelRunnerOutput(req_ids=[], req_id_to_index={})) + self.pending = (scheduler_output, num_drafts) + return self._done(None) + + def sample_tokens(self, grammar_output: GrammarOutput | None, non_block=False): + assert self.pending is not None + scheduler_output, num_drafts = self.pending + self.pending = None + rows: dict[tuple[str, int], np.ndarray] = {} + num_invalid: dict[str, int] = {} + if grammar_output is not None: + row_iter = iter(grammar_output.grammar_bitmask) + for req_id in grammar_output.structured_output_request_ids: + for position in range(num_drafts[req_id] + 1): + rows[req_id, position] = next(row_iter) + num_invalid = grammar_output.num_invalid_spec_tokens or {} + + req_ids = list(scheduler_output.num_scheduled_tokens) + sampled: list[list[int]] = [] + accepted: dict[str, int] = {} + self.accepted_per_step.append(accepted) + for req_id in req_ids: + state = self.requests[req_id] + end = state.num_computed + scheduler_output.num_scheduled_tokens[req_id] + if end < state.prefill_len: + state.num_computed = end + sampled.append([]) + continue + n = num_drafts[req_id] + n_valid = n - num_invalid.get(req_id, 0) + drafts = state.drafts[:n_valid] + [-1] * (n - n_valid) + tokens: list[int] = [] + for position in range(n + 1): + target = _preferred_token(state.num_output + position) + row = rows.get((req_id, position)) + if row is not None and not _allowed(row, target): + target = next(t for t in range(VOCAB_SIZE) if _allowed(row, t)) + tokens.append(target) + if position == n or drafts[position] != target: + break + accepted[req_id] = len(tokens) - 1 + state.num_computed = end - (n + 1 - len(tokens)) + state.num_output += len(tokens) + state.drafts = [ + _preferred_token(state.num_output + i) for i in range(NUM_SPEC_TOKENS) + ] + sampled.append(tokens) + self.proposed = DraftTokenIds( + req_ids, + [ + self.requests[req_id].drafts + if scheduler_output.has_structured_output_requests + else [-1] * NUM_SPEC_TOKENS + for req_id in req_ids + ], + ) + return self._done( + ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=sampled, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + ) + + def take_draft_token_ids(self) -> DraftTokenIds: + return self.consumed if self.consumed_handoff else self.proposed + + +def _engine(scheduler, executor) -> EngineCore: + """The real step_with_batch_queue loop around a fake executor.""" + engine = object.__new__(EngineCore) + engine.scheduler = scheduler + engine.model_executor = executor + engine.batch_queue_size = 2 + engine.batch_queue = deque(maxlen=2) + engine.is_ec_consumer = True + engine.is_pooling_model = False + engine.check_for_draft_tokens = True + engine._last_model_completion_time = None + engine._should_throttle_prefills = Mock(return_value=False) + engine.log_error_detail = lambda _: nullcontext() + engine.capture_iteration_details = lambda _: nullcontext() + engine._wait_for_boundary_checkpoint_copies = Mock() + engine._process_aborts_queue = Mock() + engine._attach_iteration_details = Mock() + return engine + + +@pytest.fixture +def opt_model_path(tmp_path): + (tmp_path / "config.json").write_text( + '{"architectures":["OPTForCausalLM"],"model_type":"opt",' + '"max_position_embeddings":32768}' + ) + return str(tmp_path) + + +def _structured_request() -> Request: + params = SamplingParams( + max_tokens=len(GRAMMAR_TOKENS), + structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), + ) + params.update_from_generation_config({}, EOS_TOKEN_ID) + request = Request( + request_id="structured", + prompt_token_ids=list(range(8)), + sampling_params=params, + pooling_params=None, + ) + assert request.structured_output_request is not None + request.structured_output_request.grammar = SequenceGrammar() + return request + + +@pytest.mark.parametrize("consumed_handoff", [True, False]) +def test_skipped_step_keeps_draft_rows_constrained(opt_model_path, consumed_handoff): + """consumed_handoff=False is a runner that still hands back the drafts the + last sampled step proposed: the scheduler must then reject the unresolved + drafts rather than let the runner accept them against open rows.""" + scheduler = create_scheduler( + model=opt_model_path, + skip_tokenizer_init=True, + device="cpu", + async_scheduling=True, + use_v2_model_runner=True, + num_speculative_tokens=NUM_SPEC_TOKENS, + speculative_method="ngram_gpu", + prefill_compute_share=0.5, + max_num_batched_tokens=32, + max_model_len=512, + ) + scheduler.structured_output_manager.backend = BitmaskBackend() + executor = FakeV2Executor(consumed_handoff) + engine = _engine(scheduler, executor) + + request = _structured_request() + scheduler.add_request(request) + # A long plain prompt contends with the structured decode, so prefill + # turns leave the decode out of some steps. + (prefill,) = create_requests( + num_requests=1, num_tokens=320, max_tokens=2, req_ids=["prefill"] + ) + + clock = itertools.count(step=0.001) + with patch("vllm.v1.engine.core.time.perf_counter", lambda: next(clock)): + for step in range(200): + if step == 2: + scheduler.add_request(prefill) + if not (scheduler.has_requests() or engine.batch_queue): + break + engine.step_with_batch_queue() + + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED + assert list(request.output_token_ids) == GRAMMAR_TOKENS + + steps = [i for i, a in enumerate(executor.accepted_per_step) if "structured" in a] + after_skip = [b for a, b in itertools.pairwise(steps) if b - a > 1] + assert after_skip, "the prefill turns never skipped the structured decode" + accepted_after_skip = sum( + executor.accepted_per_step[i]["structured"] for i in after_skip + ) + if consumed_handoff: + # The skipped request's drafts reach the grammar, so it keeps + # speculating instead of verifying one token per step. + assert accepted_after_skip > 0 + else: + assert accepted_after_skip == 0 + + +def test_deferred_step_without_inflight_batch_samples_immediately(): + """A deferred step may have no earlier batch to wait for: a request that + skipped a step has its output processed already, and only its drafts still + need the back-fill.""" + done: Future = Future() + done.set_result(None) + scheduler_output = SchedulerOutput.make_empty() + scheduler_output.total_num_scheduled_tokens = 4 + scheduler_output.pending_structured_output_tokens = True + scheduler = Mock() + scheduler.has_requests.return_value = True + scheduler.schedule.return_value = scheduler_output + executor = Mock() + executor.execute_model.return_value = done + executor.sample_tokens.return_value = done + engine = _engine(scheduler, executor) + + assert engine.step_with_batch_queue() == (None, True) + + draft_token_ids = executor.take_draft_token_ids.return_value + scheduler.update_draft_token_ids_in_output.assert_called_once_with( + draft_token_ids, scheduler_output + ) + scheduler.get_grammar_bitmask.assert_called_once_with(scheduler_output) + executor.sample_tokens.assert_called_once_with( + scheduler.get_grammar_bitmask.return_value, non_block=True + ) + assert [entry[1] for entry in engine.batch_queue] == [scheduler_output] diff --git a/tests/v1/worker/test_draft_tokens_handler.py b/tests/v1/worker/test_draft_tokens_handler.py new file mode 100644 index 000000000000..ffa7bdc51451 --- /dev/null +++ b/tests/v1/worker/test_draft_tokens_handler.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The draft hand-off must return the drafts a step's inputs consumed. + +Under async scheduling the scheduler back-fills a step's draft placeholders +after that step's inputs were prepared. A request that skipped the step whose +drafts were proposed last must still get its own drafts back. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="DraftTokensHandler copies on CUDA streams" +) + + +def _batch(req_ids, slots, num_drafts, structured=True): + return SimpleNamespace( + req_ids=req_ids, + idx_mapping=torch.tensor(slots, dtype=torch.int32, device="cuda"), + num_draft_tokens_per_req=( + None if num_drafts is None else np.array(num_drafts, dtype=np.int32) + ), + has_structured_output_reqs=structured, + ) + + +def _draft_slots(): + # One row of drafts per persistent request slot. + return torch.tensor( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.int64, device="cuda" + ) + + +def test_consumed_drafts_include_requests_that_skipped_the_proposing_step(): + handler = DraftTokensHandler(torch.device("cuda"), track_consumed_drafts=True) + drafts = _draft_slots() + + # The latest sampled step proposed drafts for slot 0 only. + handler.set_draft_tokens(_batch(["a"], [0], [3]), drafts[:1]) + # The next step also verifies slot 2, which skipped the proposing step. + handler.set_consumed_draft_tokens(_batch(["a", "b"], [0, 2], [3, 2]), drafts) + # Proposals of the step being sampled must not replace the hand-off. + handler.set_draft_tokens(_batch(["a", "b"], [0, 2], [3, 2]), drafts[:2]) + + draft_token_ids = handler.get_draft_tokens() + assert draft_token_ids.req_ids == ["a", "b"] + assert draft_token_ids.draft_token_ids == [[1, 2, 3], [7, 8, 9]] + + +def test_step_without_drafts_clears_the_hand_off(): + handler = DraftTokensHandler(torch.device("cuda"), track_consumed_drafts=True) + drafts = _draft_slots() + handler.set_consumed_draft_tokens(_batch(["a"], [0], [3]), drafts) + + handler.set_consumed_draft_tokens(_batch(["b"], [1], None), drafts) + + draft_token_ids = handler.get_draft_tokens() + assert draft_token_ids.req_ids == ["b"] + assert draft_token_ids.draft_token_ids == [[]] + + +def test_synchronous_scheduling_hands_back_proposals(): + handler = DraftTokensHandler(torch.device("cuda")) + drafts = _draft_slots() + + handler.set_consumed_draft_tokens(_batch(["a", "b"], [0, 2], [3, 3]), drafts) + handler.set_draft_tokens(_batch(["c"], [1], [3]), drafts[1:2]) + + draft_token_ids = handler.get_draft_tokens() + assert draft_token_ids.req_ids == ["c"] + assert draft_token_ids.draft_token_ids == [[4, 5, 6]] diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index a1715c7b7831..990d3810f0fb 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -28,14 +28,18 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: if request.is_prefill_chunk: continue - scheduler_output.pending_structured_output_tokens |= ( - request.use_structured_output and request.num_output_placeholders > 0 - ) # The request will generate num_sampled_tokens_per_step new tokens # plus num_spec_tokens in this scheduling step. Diffusion has no AR # bonus token (num_sampled_tokens_per_step == 0) — only the canvas # (spec) tokens. cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) + scheduler_output.pending_structured_output_tokens |= ( + request.use_structured_output + and ( + request.num_output_placeholders > 0 + or self._needs_draft_backfill(cur_num_spec_tokens) + ) + ) request.num_output_placeholders += ( self.num_sampled_tokens_per_step + cur_num_spec_tokens ) @@ -48,6 +52,23 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: # scheduled for decode (for PP microbatching). request.next_decode_eligible_step = self.current_step + self.pp_size + def _needs_draft_backfill(self, num_scheduled_spec_tokens: int) -> bool: + """Whether a structured-output request's drafts must be handed back + before its grammar bitmask is built. + + This scheduler only holds -1 placeholders for the drafts; the V2 model + runner verifies its own copy. A request that skipped the previous step + (a prefill turn, an isolated saved-logits step) has no output in + flight, but its placeholders still need the back-fill that deferred + sampling performs, or its bitmask rows past the first draft would + allow every token. + """ + return ( + num_scheduled_spec_tokens > 0 + and self.use_v2_model_runner + and self.num_sampled_tokens_per_step > 0 + ) + def _update_request_with_output( self, request: Request, new_token_ids: list[int], is_stale: bool = False ) -> tuple[list[int], bool]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9247348da340..66bd89742b79 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -87,6 +87,7 @@ from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm.v1.structured_output import StructuredOutputManager +from vllm.v1.structured_output.utils import strip_speculative_padding from vllm.v1.utils import record_function_or_nullcontext logger = init_logger(__name__) @@ -2819,6 +2820,9 @@ def get_grammar_bitmask( if not structured_output_request_ids: return None + self._reject_unconstrained_drafts( + scheduler_output, structured_output_request_ids + ) bitmask = self.structured_output_manager.grammar_bitmask( self.requests, structured_output_request_ids, @@ -2830,6 +2834,41 @@ def get_grammar_bitmask( scheduler_output.num_invalid_spec_tokens, ) + def _reject_unconstrained_drafts( + self, + scheduler_output: SchedulerOutput, + structured_output_request_ids: list[str], + ) -> None: + """Invalidate the drafts whose grammar bitmask rows are unconstrained. + + The grammar bitmask follows this scheduler's copy of the drafts, but the + V2 model runner verifies the drafts it holds itself. Rows after a -1 + placeholder that was never back-filled allow every token, so the + runner must reject those drafts instead of accepting tokens the + grammar never allowed. Back-filled drafts that the grammar rejected + already carry the same count. + """ + if not self.use_v2_model_runner or self.num_sampled_tokens_per_step == 0: + return + spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens + if not spec_decode_tokens: + return + num_invalid_spec_tokens = scheduler_output.num_invalid_spec_tokens or {} + for req_id in structured_output_request_ids: + drafts = spec_decode_tokens.get(req_id) + if not drafts: + continue + num_unconstrained = len(drafts) - len(strip_speculative_padding(drafts)) + if num_unconstrained > num_invalid_spec_tokens.get(req_id, 0): + num_invalid_spec_tokens[req_id] = num_unconstrained + logger.warning_once( + "Structured-output drafts were not handed back to the " + "scheduler; rejecting them so every sampled token stays " + "grammar-constrained." + ) + if num_invalid_spec_tokens: + scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens + def update_from_output( self, scheduler_output: SchedulerOutput, diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 316e5fd22961..f5511d4fda1f 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -778,6 +778,13 @@ def step_with_batch_queue( # Don't block on next worker response unless the queue is full # or there are no more requests to schedule. return None, model_executed + elif not batch_queue: + # Only this step's drafts are outstanding (no prior output is in + # flight), so hand them back and sample now. + self._sample_deferred_batch( + deferred_scheduler_output, exec_future, deferred_execution_timing + ) + return None, model_executed elif not batch_queue: # Queue is empty. We should not reach here since this method should @@ -821,34 +828,39 @@ def step_with_batch_queue( # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. if deferred_scheduler_output: - # When draft tokens are used with structured output, validate them - # before computing the grammar bitmask for the deferred request. - if self.check_for_draft_tokens: - draft_token_ids = self.model_executor.take_draft_token_ids() - if draft_token_ids is not None: - # Update the draft token ids in the scheduler output to - # filter out the invalid spec tokens, which will be padded - # with -1 and skipped by the grammar bitmask computation. - self.scheduler.update_draft_token_ids_in_output( - draft_token_ids, deferred_scheduler_output - ) - # We now have the tokens needed to compute the bitmask for the - # deferred request. Get the bitmask and call sample tokens. - grammar_output = self.scheduler.get_grammar_bitmask( - deferred_scheduler_output - ) - future = self.model_executor.sample_tokens(grammar_output, non_block=True) - batch_queue.appendleft( - ( - future, - deferred_scheduler_output, - exec_future, - deferred_execution_timing, - ) + self._sample_deferred_batch( + deferred_scheduler_output, exec_future, deferred_execution_timing ) return engine_core_outputs, model_executed + def _sample_deferred_batch( + self, + scheduler_output: SchedulerOutput, + exec_future: Future[Any], + execution_timing: float | None, + ) -> None: + """Back-fill the step's drafts, build its bitmask and queue sampling.""" + assert self.batch_queue is not None + # When draft tokens are used with structured output, validate them + # before computing the grammar bitmask for the deferred request. + if self.check_for_draft_tokens: + draft_token_ids = self.model_executor.take_draft_token_ids() + if draft_token_ids is not None: + # Update the draft token ids in the scheduler output to + # filter out the invalid spec tokens, which will be padded + # with -1 and skipped by the grammar bitmask computation. + self.scheduler.update_draft_token_ids_in_output( + draft_token_ids, scheduler_output + ) + # We now have the tokens needed to compute the bitmask for the + # deferred request. Get the bitmask and call sample tokens. + grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) + future = self.model_executor.sample_tokens(grammar_output, non_block=True) + self.batch_queue.appendleft( + (future, scheduler_output, exec_future, execution_timing) + ) + def _process_aborts_queue(self): if not self.aborts_queue.empty(): request_ids = [] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 55a2dfa07462..299089453968 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -306,7 +306,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.use_aux_hidden_state_outputs = True # Draft tokens propagation - for spec-dec + struct outputs. - self.draft_tokens_handler = DraftTokensHandler(self.device) + self.draft_tokens_handler = DraftTokensHandler( + self.device, + track_consumed_drafts=bool(self.scheduler_config.async_scheduling) + and self.speculative_config is not None + and self.is_last_pp_rank, + ) self.pcp_manager: pcp.PCPManager | None = None @@ -1962,6 +1967,11 @@ def execute_model( input_batch = self.prepare_inputs( scheduler_output, batch_req_state, batch_desc ) + # Deferred structured-output sampling back-fills the scheduler's + # draft placeholders with exactly the drafts these inputs hold. + self.draft_tokens_handler.set_consumed_draft_tokens( + input_batch, self.req_states.draft_tokens + ) if scheduler_output.boundary_logits_only: assert self.boundary_checkpoint_state is not None assert input_batch.num_reqs == 1 diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 6657ff195c0e..98964a72b1f4 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -55,8 +55,24 @@ def limit_draft_tokens( class DraftTokensHandler: - def __init__(self, device: torch.device | None = None): + """Hands draft tokens back to the scheduler for grammar validation. + + With synchronous scheduling the scheduler takes the drafts proposed for the + next step. With async scheduling (``track_consumed_drafts``) it takes them + only to back-fill the -1 placeholders of a step whose inputs were already + prepared, so the handler records the drafts that step's inputs consumed. + The drafts of the last proposing step would miss every request that + skipped it, and the scheduler would then build unconstrained grammar rows + for drafts that the model still verifies. + """ + + def __init__( + self, + device: torch.device | None = None, + track_consumed_drafts: bool = False, + ): self.device = device + self.track_consumed_drafts = track_consumed_drafts self.copy_stream = torch.cuda.Stream(device) # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. self.copy_event = torch.cuda.Event(blocking=True) @@ -68,9 +84,39 @@ def __init__(self, device: torch.device | None = None): def set_draft_tokens( self, input_batch: InputBatch, draft_tokens: torch.Tensor ) -> None: - self.req_ids = input_batch.req_ids + """Record the drafts proposed for the next step.""" + if self.track_consumed_drafts: + return + self._set( + input_batch.req_ids, draft_tokens, input_batch.has_structured_output_reqs + ) + + def set_consumed_draft_tokens( + self, input_batch: InputBatch, draft_tokens: torch.Tensor + ) -> None: + """Record the drafts that this step's inputs were built from. + + Args: + input_batch: The prepared step. + draft_tokens: The persistent per-request draft buffer, read on the + current stream after the step's input ids were built from it. + """ + if not self.track_consumed_drafts: + return + num_drafts = input_batch.num_draft_tokens_per_req + if num_drafts is None or not input_batch.has_structured_output_reqs: + self._set(input_batch.req_ids, draft_tokens[:0, :0], False) + return + self._set( + input_batch.req_ids, + draft_tokens[input_batch.idx_mapping, : int(num_drafts.max())], + True, + ) + + def _set(self, req_ids: list[str], draft_tokens: torch.Tensor, copy: bool) -> None: + self.req_ids = req_ids self.num_draft_tokens = draft_tokens.shape[1] - if not input_batch.has_structured_output_reqs: + if not copy: # No draft token validation needs to be performed by # the scheduler for this batch. self.draft_tokens_np = None @@ -93,7 +139,8 @@ def get_draft_tokens(self) -> DraftTokenIds | None: self.copy_event.synchronize() draft_token_ids = self.draft_tokens_np.tolist() else: - # This case only happens when async scheduling is disabled. + # The batch carried no structured-output request, so its drafts + # need no grammar validation. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] return DraftTokenIds(self.req_ids, draft_token_ids) From 1c4a92571f47bbca5e1718c0214eaea5d14f9cfe Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:32:59 +0000 Subject: [PATCH 293/365] fix(cudagraph): measure sampled graphs after the warmup cache is flushed The CUDA graph memory profiler captures a few FULL graphs and extrapolates the rest. Each sample read free memory before torch.cuda.graph() entered, but torch.cuda.graph() empties the allocator cache on entry, so the blocks cached by the eager warmup were released inside the sample and subtracted from the graph's own memory. On DeepSeek-V4.1 the warmup caches the same 116 MiB of decode metadata that each captured graph retained before #918; the two cancelled, the per-graph sample read about zero and the estimate reserved 1.8 GiB for graphs that took 5.3 GiB. With the full-capture debug accounting the same error went the other way: samples of -116 MiB were floored while the measured total kept them, reporting 5.1 GiB for 1.75 GiB. Flush the cache before and after each sampled capture, so a sample is the device memory the graph adds: new graph pool segments plus memory outside the allocator such as the instantiated graph. An unsampled graph is charged what a sampled graph keeps for its lifetime: the tensors it retains, even when the pool had room for them, and memory outside the allocator. Pool growth for activations is left out of that per-graph cost because later, smaller graphs reuse it. The two largest graphs alone cannot tell a fixed per-graph cost from one that scales with the batch: scaling the second sample by token count charged a 1-token graph 1/120 of it. Sample the smallest FULL graph too and interpolate each unsampled graph linearly in its token count between the second-largest and the smallest sample. Model runner V1 had the same bias whenever blocks were cached before a sample, so flush around its samples as well. Co-Authored-By: Claude Opus 5.5 --- ...gpu_model_runner_v2_cudagraph_profiling.py | 372 ++++++++++++++++-- vllm/v1/worker/gpu/cudagraph_utils.py | 225 ++++++++--- vllm/v1/worker/gpu_model_runner.py | 26 +- 3 files changed, 529 insertions(+), 94 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index 65f16fdc11db..50e44e7283cf 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py @@ -14,6 +14,7 @@ import gc from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock from weakref import ref import pytest @@ -44,8 +45,8 @@ def __init__( else: self._capture_descs = {CUDAGraphMode.FULL: descs} if needs_capture else {} # Profiling hooks set by profile_cudagraph_memory. - self._max_full_descs_to_capture: int | None = None - self._capture_mem_samples: list[int] | None = None + self._sample_full_descs = False + self._capture_mem_samples: list[Any] | None = None self.use_breakable_cg = False self.graphs: dict[Any, Any] = {} self.graph_capture_resources: dict[Any, list[Any]] = {} @@ -82,8 +83,9 @@ def _make_profiling_runner( num_full_descs: int = 3, piecewise_only: bool = False, captured_bytes: int = 7 << 30, - mem_samples: list[int] | None = None, + mem_samples: list[tuple[int, int, int]] | None = None, ) -> Any: + """Build a runner whose capture records ``(desc index, growth, cost)``.""" runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) runner.compilation_config = SimpleNamespace(cudagraph_mode=cudagraph_mode) runner.cudagraph_manager = _FakeCudaGraphManager( @@ -101,9 +103,14 @@ def _capture_model(*, profile_only: bool = False) -> int: events.append("capture") runner.pool_during_capture = runner.cudagraph_manager.pool # Simulate the manager's per-FULL-graph memory sampling. - samples = runner.cudagraph_manager._capture_mem_samples - if samples is not None: - samples.extend(mem_samples or []) + manager = runner.cudagraph_manager + samples = manager._capture_mem_samples + if samples is not None and mem_samples: + descs = manager._capture_descs[CUDAGraphMode.FULL] + samples.extend( + cgu._FullGraphMemorySample(descs[index], growth, cost) + for index, growth, cost in mem_samples + ) return captured_bytes runner.capture_model = _capture_model @@ -173,33 +180,52 @@ def test_profile_cudagraph_memory_no_graphs_tears_down(monkeypatch): def test_profile_cudagraph_memory_samples_and_extrapolates(monkeypatch): _patch_module(monkeypatch) - gib = 1 << 30 - # Measured delta 1000 MiB includes the sampled FULL graphs (100 + 20 MiB). - # The sampled graphs have sizes 3 and 2. Scale the second sample by the - # remaining graph's 1/2 token ratio: 100 + 20 + 10 = 130 MiB. + mib = 1 << 20 + # The measured delta includes the sampled FULL graphs of 4, 3 and 1 + # tokens. The unsampled 2-token graph costs the per-graph samples + # interpolated between 1 token (10 MiB) and 3 tokens (30 MiB). runner = _make_profiling_runner( CUDAGraphMode.FULL, - num_full_descs=3, - captured_bytes=1000 * gib, - mem_samples=[100 * gib, 20 * gib], + num_full_descs=4, + captured_bytes=1000 * mib, + mem_samples=[ + (0, 100 * mib, 100 * mib), + (1, 30 * mib, 30 * mib), + (3, 0, 10 * mib), + ], ) result = cgu.profile_cudagraph_memory( runner, lambda: runner.events.append("prepare") ) - assert result == (1000 - (100 + 20) + (100 + 20 + 10)) * gib + assert result == (1000 + 20) * mib # Bootstrap, capture, and teardown run in order. assert runner.events == ["init", "prepare", "capture", "teardown"] # Capture must use a throwaway pool, not the persistent global pool. assert runner.pool_during_capture == THROWAWAY_POOL assert runner.cudagraph_manager.pool == GLOBAL_POOL - # FULL capture must be limited to the largest few graphs. - assert ( - runner.cudagraph_manager._max_full_descs_to_capture - == cgu._FULL_GRAPH_PROFILING_SAMPLES + # FULL capture must be limited to a sample of the graphs. + assert runner.cudagraph_manager._sample_full_descs + + +def test_profile_cudagraph_memory_debug_accounting_captures_every_full_graph( + monkeypatch, +): + _patch_module(monkeypatch) + monkeypatch.setattr(cgu, "_DEBUG_GRAPH_MEMORY_ACCOUNTING", True) + mib = 1 << 20 + runner = _make_profiling_runner( + CUDAGraphMode.FULL, + num_full_descs=2, + captured_bytes=1000 * mib, + mem_samples=[(0, 100 * mib, 100 * mib), (1, 0, 50 * mib)], ) + # Every graph was captured, so the measured delta is the estimate. + assert cgu.profile_cudagraph_memory(runner) == 1000 * mib + assert not runner.cudagraph_manager._sample_full_descs + def test_profile_cudagraph_memory_piecewise_only_returns_measured(monkeypatch): _patch_module(monkeypatch) @@ -349,27 +375,253 @@ def test_model_runner_delegates_to_cudagraph_utils(monkeypatch): assert runner.profile_cudagraph_memory(prepare) == (runner, prepare) +def _full_descs(*token_counts: int) -> list[cgu.BatchExecutionDescriptor]: + return [ + cgu.BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens, num_tokens) + for num_tokens in token_counts + ] + + +def test_profiling_samples_two_largest_and_smallest_full_graphs(): + descs = _full_descs(40, 32, 24, 16, 8) + + assert cgu._profiling_full_descs(descs) == [descs[0], descs[1], descs[-1]] + assert cgu._profiling_full_descs(descs[:3]) == descs[:3] + + def test_extrapolate_full_graph_memory(): mib = 1 << 20 - descs = [ - cgu.BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens, None) - for num_tokens in (40, 32, 16, 8) - ] + descs = _full_descs(40, 32, 24, 16, 8) + + def sample(index: int, growth: int, cost: int | None = None) -> Any: + return cgu._FullGraphMemorySample( + descs[index], growth * mib, (growth if cost is None else cost) * mib + ) + # No samples (e.g. no FULL graphs): nothing to add. assert cgu._extrapolate_full_graph_memory([], []) == 0 - # A single graph costs exactly its sample. - assert cgu._extrapolate_full_graph_memory([100 * mib], descs[:1]) == 100 * mib - # Preserve sampled costs and scale the rest by their token counts. + # Sampled graphs cost exactly their measured growth. + assert cgu._extrapolate_full_graph_memory([sample(0, 100)], descs[:1]) == ( + 100 * mib + ) + # Unsampled graphs interpolate the per-graph samples by token count. + assert ( + cgu._extrapolate_full_graph_memory( + [sample(0, 100), sample(1, 32), sample(4, 8)], descs + ) + == (100 + 32 + 8 + 24 + 16) * mib + ) + # A fixed retained cost is charged to every unsampled graph, even when + # the sampled graphs found room for it in the pool, and one-time pool + # growth of a sampled graph is not. assert ( - cgu._extrapolate_full_graph_memory([100 * mib, 20 * mib], descs) - == (100 + 20 + 10 + 5) * mib + cgu._extrapolate_full_graph_memory( + [sample(0, 300), sample(1, 80, 122), sample(4, 6, 122)], descs + ) + == (300 + 80 + 6 + 2 * 122) * mib ) - # Per-graph cost is floored to account for driver overhead. + # The largest graph's sample is not a per-graph reference, and the + # per-graph cost is floored for driver overhead. assert ( - cgu._extrapolate_full_graph_memory([100 * mib, 0], descs) == (100 + 3 * 1) * mib + cgu._extrapolate_full_graph_memory( + [sample(0, 300), sample(1, 0), sample(4, 0)], descs + ) + == (300 + 2 * 1) * mib ) +_MIB = 1 << 20 +# Decode page metadata that a buggy attention backend retains in every FULL +# graph (the DeepSeek-V4.1 per-graph copy fixed in vLLM #918). +_METADATA_BYTES = 116 * _MIB +_GRAPH_EXEC_BYTES = 6 * _MIB + + +class _FakeDevice: + """One device as the graph memory sampler reads it, in bytes. + + Eager blocks stay cached after they are freed until ``empty_cache()``. + The private graph pool reuses its free blocks but never returns memory, + and every captured graph adds ``_GRAPH_EXEC_BYTES`` outside the Torch + allocator, like an instantiated CUDA graph. + """ + + def __init__(self) -> None: + self.eager_allocated = 0 + self.eager_reserved = 0 + self.pool_allocated = 0 + self.pool_reserved = 0 + self.native = 0 + self.capturing = False + # Cached bytes that torch.cuda.graph's own flush found on entry. + self.flushed_by_capture: list[int] = [] + + def get_memory_info(self) -> tuple[int, int]: + total = 96 << 30 + used = self.eager_reserved + self.pool_reserved + self.native + return total - used, total + + def memory_allocated(self) -> int: + return self.eager_allocated + self.pool_allocated + + def memory_reserved(self) -> int: + return self.eager_reserved + self.pool_reserved + + def empty_cache(self) -> None: + self.eager_reserved = self.eager_allocated + + def allocate(self, nbytes: int) -> None: + if self.capturing: + self.pool_allocated += nbytes + self.pool_reserved = max(self.pool_reserved, self.pool_allocated) + else: + self.eager_allocated += nbytes + self.eager_reserved = max(self.eager_reserved, self.eager_allocated) + + def free(self, nbytes: int) -> None: + if self.capturing: + self.pool_allocated -= nbytes + else: + self.eager_allocated -= nbytes + + @contextlib.contextmanager + def cuda_graph(self, *_args, **_kwargs): + # Like torch.cuda.graph, flush the allocator cache before capturing. + self.flushed_by_capture.append(self.eager_reserved - self.eager_allocated) + self.empty_cache() + self.capturing = True + try: + yield + finally: + self.capturing = False + self.native += _GRAPH_EXEC_BYTES + + def create_forward_fn(self, retain_metadata: bool): + """Decode forward: activations scale with tokens, metadata is fixed.""" + + def create_forward_fn(desc, warmup): + def forward_fn(_mode): + activations = 2 * _MIB * desc.num_tokens + self.allocate(activations + _METADATA_BYTES) + self.free(activations) + if not (self.capturing and retain_metadata): + self.free(_METADATA_BYTES) + + return forward_fn + + return create_forward_fn + + def capture_model(self, manager: cgu.CudaGraphManager, create_forward_fn): + """Mirror GPUModelRunner.capture_model's measured delta.""" + self.empty_cache() + free_before = self.get_memory_info()[0] + manager.capture(create_forward_fn) + return free_before - self.get_memory_info()[0] + + +class _CapturingGraphManager(cgu.CudaGraphManager): + """Runs the real capture loop over FULL decode graphs of the given sizes.""" + + def __init__(self, token_counts: list[int]) -> None: + self.device = torch.device("cpu") + self.pool: Any = GLOBAL_POOL + self.ubatch_runner = None + self.use_breakable_cg = False + self._capture_descs = { + CUDAGraphMode.FULL: _full_descs(*sorted(token_counts, reverse=True)) + } + self._sample_full_descs = False + self._capture_mem_samples: list[Any] | None = None + self.graphs: dict[Any, Any] = {} + self.graph_capture_resources: dict[Any, list[Any]] = {} + self._graphs_captured = False + + def needs_capture(self) -> bool: + return True + + def _capture_stream(self, desc): + return None + + +def _patch_capture(monkeypatch, device: _FakeDevice) -> None: + accelerator = cgu.torch.accelerator + monkeypatch.setattr(accelerator, "synchronize", lambda: None) + for name in ("get_memory_info", "memory_allocated", "memory_reserved"): + monkeypatch.setattr(accelerator, name, getattr(device, name)) + monkeypatch.setattr(accelerator, "empty_cache", device.empty_cache) + monkeypatch.setattr(cgu.torch.cuda, "graph", device.cuda_graph) + monkeypatch.setattr(cgu.torch.cuda, "CUDAGraph", MagicMock) + monkeypatch.setattr(cgu, "graph_capture", lambda **_: contextlib.nullcontext()) + monkeypatch.setattr(cgu, "get_offloader", MagicMock) + monkeypatch.setattr(cgu, "set_graph_pool_id", lambda _pool: None) + monkeypatch.setattr(cgu, "is_global_first_rank", lambda: False) + + +# Decode graph sizes shaped like DeepSeek-V4.1's DSpark verification graphs. +_DECODE_TOKENS = [*range(128, 0, -8), *range(7, 0, -1)] + + +@pytest.mark.parametrize("retain_metadata", [False, True]) +def test_full_graph_sample_excludes_released_warmup_cache(monkeypatch, retain_metadata): + """The baseline is read after the warmup's cached blocks are flushed. + + ``torch.cuda.graph`` empties the allocator cache when it is entered. A + baseline read before that flush subtracts the eager warmup's cached + activations and metadata from the graph's own memory. + """ + device = _FakeDevice() + _patch_capture(monkeypatch, device) + manager = _CapturingGraphManager([128, 120, 1]) + samples: list[Any] = [] + manager._capture_mem_samples = samples + + manager.capture(device.create_forward_fn(retain_metadata)) + + # Nothing is left for the capture's own flush to release in the sample. + assert device.flushed_by_capture == [0, 0, 0] + # With retained metadata, the 120-token graph needs 100 MiB more pool and + # the 1-token graph fits in blocks freed by earlier graphs; each graph + # still keeps its own metadata. + retained = _METADATA_BYTES if retain_metadata else 0 + pool_growth = [256 * _MIB + _METADATA_BYTES, 100 * _MIB if retained else 0, 0] + assert [sample.growth for sample in samples] == [ + pool + _GRAPH_EXEC_BYTES for pool in pool_growth + ] + assert [sample.cost for sample in samples] == [retained + _GRAPH_EXEC_BYTES] * 3 + + +@pytest.mark.parametrize("retain_metadata", [False, True]) +def test_profile_estimate_covers_every_full_graph(monkeypatch, retain_metadata): + """Three sampled graphs must account for memory each graph retains. + + Compare the profiled estimate with a capture of every graph. A per-graph + retained buffer must shrink the KV cache instead of running out of memory + after the real capture. + """ + _patch_module(monkeypatch) + full_device = _FakeDevice() + _patch_capture(monkeypatch, full_device) + all_graphs_bytes = full_device.capture_model( + _CapturingGraphManager(_DECODE_TOKENS), + full_device.create_forward_fn(retain_metadata), + ) + + device = _FakeDevice() + _patch_capture(monkeypatch, device) + runner = _make_profiling_runner(CUDAGraphMode.FULL) + runner.cudagraph_manager = _CapturingGraphManager(_DECODE_TOKENS) + runner.capture_model = lambda *, profile_only: device.capture_model( + runner.cudagraph_manager, device.create_forward_fn(retain_metadata) + ) + + estimate = cgu.profile_cudagraph_memory(runner) + + assert all_graphs_bytes <= estimate <= all_graphs_bytes * 21 // 20 + assert device.flushed_by_capture == [0, 0, 0] + if not retain_metadata: + assert estimate == all_graphs_bytes + + def test_profile_cudagraph_memory_clears_captured_graphs(monkeypatch): _patch_module(monkeypatch) runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE) @@ -920,3 +1172,65 @@ def _partial_init() -> None: assert not hasattr(runner, "kv_cache_config") assert runner.cache_config.num_gpu_blocks is None assert runner._mamba_bufs is None + + +def test_legacy_profile_samples_exclude_cached_blocks(monkeypatch): + """Model runner V1 samples each graph with the allocator cache empty.""" + from vllm.v1.worker import gpu_model_runner as legacy_runner_module + + device = _FakeDevice() + _patch_capture(monkeypatch, device) + runner = legacy_runner_module.GPUModelRunner.__new__( + legacy_runner_module.GPUModelRunner + ) + runner.vllm_config = object() + runner.lora_config = None + runner.device = torch.device("cpu") + runner.max_model_len = 1024 + runner.max_num_tokens = 128 + runner.cudagraph_dispatcher = SimpleNamespace( + get_capture_descs=lambda: [(CUDAGraphMode.FULL, _full_descs(128, 120, 1))], + cudagraph_keys={}, + keys_initialized=True, + ) + runner._init_minimal_kv_cache_for_profiling = lambda: None + runner._cleanup_profiling_kv_cache = lambda: None + runner._create_encoder_cudagraph_manager = lambda: None + runner._freeze_gc = contextlib.nullcontext + runner.maybe_remove_all_loras = lambda _config: None + create_forward_fn = device.create_forward_fn(retain_metadata=True) + + def warmup_and_capture(desc, cudagraph_runtime_mode, profile_seq_lens=None): + forward_fn = create_forward_fn(desc, warmup=True) + forward_fn(CUDAGraphMode.NONE) + with device.cuda_graph(): + forward_fn(cudagraph_runtime_mode) + # The dummy run's eager outputs stay cached after the capture. + device.allocate(64 * _MIB) + device.free(64 * _MIB) + + runner._warmup_and_capture = warmup_and_capture + monkeypatch.setattr( + legacy_runner_module, + "set_current_vllm_config", + lambda _config: contextlib.nullcontext(), + ) + monkeypatch.setattr( + legacy_runner_module, "graph_capture", lambda **_: contextlib.nullcontext() + ) + monkeypatch.setattr( + legacy_runner_module, "set_cudagraph_capturing_enabled", lambda _: None + ) + monkeypatch.setattr( + legacy_runner_module, + "current_platform", + SimpleNamespace(graph_pool_handle=object, is_rocm=lambda: False), + ) + + estimate = runner.profile_cudagraph_memory() + + # The largest graph costs its pool and instantiated graph; the next one + # the pool growth for its retained metadata, charged to both others. + first_graph = 256 * _MIB + _METADATA_BYTES + _GRAPH_EXEC_BYTES + per_graph = 100 * _MIB + _GRAPH_EXEC_BYTES + assert estimate == first_graph + 2 * per_graph diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index ea5d19cddf19..b4897e282249 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -3,8 +3,9 @@ import gc import itertools import os +from bisect import bisect_left from collections import defaultdict -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence from contextlib import ExitStack from dataclasses import dataclass, replace from itertools import groupby, product @@ -351,11 +352,11 @@ def __init__( self._graphs_captured = False - # Profiling hooks, set only by profile_cudagraph_memory() below: cap - # FULL-mode capture at the N largest descriptors and record each - # captured FULL graph's memory delta for extrapolation. - self._max_full_descs_to_capture: int | None = None - self._capture_mem_samples: list[int] | None = None + # Profiling hooks, set only by profile_cudagraph_memory() below: + # capture a sample of the FULL descriptors and record each captured + # FULL graph's memory for extrapolation. + self._sample_full_descs = False + self._capture_mem_samples: list[_FullGraphMemorySample] | None = None self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {} self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} @@ -700,14 +701,10 @@ def capture( continue descs = self._capture_descs[mode] - if ( - mode == CUDAGraphMode.FULL - and self._max_full_descs_to_capture is not None - ): - # Profiling only: capture a sample of the largest FULL - # graphs; the total cost is extrapolated from their - # per-graph memory deltas. - descs = descs[: self._max_full_descs_to_capture] + if mode == CUDAGraphMode.FULL and self._sample_full_descs: + # Profiling only: the total FULL cost is extrapolated from + # the memory of a few sampled graphs. + descs = _profiling_full_descs(descs) if is_global_first_rank(): descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})") for desc in descs: @@ -757,9 +754,9 @@ def capture( set_graph_pool_id(self.pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) - if self._capture_mem_samples is not None: - torch.accelerator.synchronize() - free_before = torch.accelerator.get_memory_info()[0] + mem_samples = self._capture_mem_samples + if mem_samples is not None: + memory_before = _flushed_device_memory() with ( collect_cuda_graph_capture_resources() as resources, torch.cuda.graph( @@ -770,10 +767,10 @@ def capture( # Join the offloader copy stream because the last layer # can leave a prefetch pending at capture end. get_offloader().join_after_forward() - if self._capture_mem_samples is not None: - torch.accelerator.synchronize() - free_after = torch.accelerator.get_memory_info()[0] - self._capture_mem_samples.append(free_before - free_after) + if mem_samples is not None: + mem_samples.append( + _measure_full_graph_memory(desc, memory_before) + ) self.graphs[desc] = graph self.graph_capture_resources[desc] = resources compilation_counter.num_cudagraph_captured += 1 @@ -1145,13 +1142,77 @@ def prepare_inputs_to_capture( # CUDA graph memory profiling # --------------------------------------------------------------------------- -# Number of FULL graphs captured during profiling; the total FULL capture -# cost is extrapolated from this sample to avoid a second full capture. -_FULL_GRAPH_PROFILING_SAMPLES = 2 +# Number of FULL graphs captured during profiling: the two largest and the +# smallest. The total FULL capture cost is extrapolated from this sample to +# avoid a second full capture. +_FULL_GRAPH_PROFILING_SAMPLES = 3 # Floor for the extrapolated per-graph cost (driver overhead per graph). _MIN_PER_GRAPH_BYTES = 1 << 20 +class _DeviceMemory(NamedTuple): + free: int + allocated: int + reserved: int + + +class _FullGraphMemorySample(NamedTuple): + """Device memory measured around the capture of one FULL graph.""" + + desc: BatchExecutionDescriptor + # Memory the capture added: new graph pool segments and memory outside + # the Torch allocator, such as the instantiated graph. + growth: int + # Memory the graph keeps while it lives: the tensors it retains, even in + # pool blocks freed by earlier captures, and memory outside the Torch + # allocator. Pool growth for activations is left out because graphs + # captured later reuse it. + cost: int + + +def _flushed_device_memory() -> _DeviceMemory: + """Read device memory around a sampled capture with the cache empty. + + ``torch.cuda.graph`` empties the allocator cache on entry, so blocks + cached by the eager warmup would otherwise be released inside the sample + and cancel out the graph's own memory. Flushing after the capture too + keeps transient cached blocks out of the sample. + """ + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + return _DeviceMemory( + free=torch.accelerator.get_memory_info()[0], + allocated=torch.accelerator.memory_allocated(), + reserved=torch.accelerator.memory_reserved(), + ) + + +def _measure_full_graph_memory( + desc: BatchExecutionDescriptor, before: _DeviceMemory +) -> _FullGraphMemorySample: + after = _flushed_device_memory() + growth = before.free - after.free + outside_allocator = growth - (after.reserved - before.reserved) + retained = after.allocated - before.allocated + return _FullGraphMemorySample( + desc, growth, max(outside_allocator, 0) + max(retained, 0) + ) + + +def _profiling_full_descs( + descs: Sequence[BatchExecutionDescriptor], +) -> list[BatchExecutionDescriptor]: + """Select the FULL graphs captured while profiling, largest first. + + The first sample also covers activations that smaller graphs reuse from + the pool. The next largest and the smallest graph then bound the + per-graph cost at both ends of the token range. + """ + if len(descs) <= _FULL_GRAPH_PROFILING_SAMPLES: + return list(descs) + return [*descs[: _FULL_GRAPH_PROFILING_SAMPLES - 1], descs[-1]] + + def _profiling_cudagraph_managers(runner: "GPUModelRunner") -> list[CudaGraphManager]: managers: list[CudaGraphManager] = [] if isinstance(runner.cudagraph_manager, CudaGraphManager): @@ -1178,13 +1239,13 @@ def profile_cudagraph_memory( graph capture. Bootstraps a minimal KV cache, runs ``capture_model()`` once, then releases everything so the real init/capture path starts clean. - FULL graphs bake in KV cache pointers, so only the largest few are - captured (into a throwaway pool) and their total cost is extrapolated. - PIECEWISE, encoder and speculator graphs are measured in full. All - profiling captures are discarded afterwards: replaying graphs recorded - against the throwaway profiling state is unsafe (e.g. inductor graph - partition reclaims the storages of earlier cudagraph recordings once the - real capture records new ones, leading to use-after-free crashes). + FULL graphs bake in KV cache pointers, so only the two largest and the + smallest are captured (into a throwaway pool) and their total cost is + extrapolated. PIECEWISE, encoder and speculator graphs are measured in + full. All profiling captures are discarded afterwards: replaying graphs + recorded against the throwaway profiling state is unsafe (e.g. inductor + graph partition reclaims the storages of earlier cudagraph recordings once + the real capture records new ones, leading to use-after-free crashes). """ runner.cudagraph_native_memory_profile = None if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: @@ -1251,12 +1312,8 @@ def profile_cudagraph_memory( for name, value in vars(speculator).items() if isinstance(value, CudaGraphManager) ] - manager._max_full_descs_to_capture = ( - None - if _DEBUG_GRAPH_MEMORY_ACCOUNTING - else _FULL_GRAPH_PROFILING_SAMPLES - ) - mem_samples: list[int] = [] + manager._sample_full_descs = not _DEBUG_GRAPH_MEMORY_ACCOUNTING + mem_samples: list[_FullGraphMemorySample] = [] manager._capture_mem_samples = mem_samples native_before_capture = _non_torch_memory_bytes() @@ -1274,8 +1331,14 @@ def profile_cudagraph_memory( # they share the global pool at runtime, so the overlap is not # double-counted. full_descs = manager._capture_descs.get(CUDAGraphMode.FULL, []) + sampled_growth = sum(sample.growth for sample in mem_samples) full_estimate = _extrapolate_full_graph_memory(mem_samples, full_descs) - return max(measured - sum(mem_samples) + full_estimate, 0) + estimate = max(measured - sampled_growth + full_estimate, 0) + if mem_samples: + _log_full_graph_estimate( + measured, mem_samples, full_descs, full_estimate, estimate + ) + return estimate finally: compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers @@ -1324,31 +1387,79 @@ def _non_torch_memory_bytes() -> int: def _extrapolate_full_graph_memory( - mem_samples: list[int], - graph_descs: list[BatchExecutionDescriptor], + mem_samples: Sequence[_FullGraphMemorySample], + graph_descs: Sequence[BatchExecutionDescriptor], ) -> int: - """Project unsampled FULL graph costs from their descriptor token counts.""" - if not mem_samples or not graph_descs: - return 0 - assert len(mem_samples) <= len(graph_descs) - - estimate = mem_samples[0] + sum( - max(sample, _MIN_PER_GRAPH_BYTES) for sample in mem_samples[1:] - ) - if len(mem_samples) == len(graph_descs) or len(mem_samples) < 2: + """Estimate the memory of all FULL graphs from the sampled captures. + + Sampled graphs count their measured growth. Each unsampled graph costs + what the sampled graphs keep, interpolated linearly in its token count, + which covers both memory that scales with the batch and a fixed cost per + graph, such as the instantiated graph and the capture resources it + retains. The first (largest) sample only bounds the pool activations + that smaller graphs reuse, so it is not a per-graph reference. + """ + estimate = sum(sample.growth for sample in mem_samples) + sampled = {id(sample.desc) for sample in mem_samples} + unsampled = [desc for desc in graph_descs if id(desc) not in sampled] + if not unsampled: return estimate - reference_desc = graph_descs[len(mem_samples) - 1] - reference_cost = max(mem_samples[-1], _MIN_PER_GRAPH_BYTES) - reference_tokens = reference_desc.num_tokens - for desc in graph_descs[len(mem_samples) :]: - scaled_cost = ( - reference_cost * desc.num_tokens + reference_tokens - 1 - ) // reference_tokens - estimate += max(scaled_cost, _MIN_PER_GRAPH_BYTES) + reference: dict[int, int] = {} + for sample in mem_samples[1:]: + num_tokens = sample.desc.num_tokens + reference[num_tokens] = max(sample.cost, reference.get(num_tokens, 0)) + points = sorted(reference.items()) + for desc in unsampled: + cost = _interpolate_graph_cost(points, desc.num_tokens) + estimate += max(cost, _MIN_PER_GRAPH_BYTES) return estimate +def _interpolate_graph_cost(points: list[tuple[int, int]], num_tokens: int) -> int: + """Interpolate sorted ``(num_tokens, cost)`` points, clamped at both ends.""" + if not points: + return 0 + index = bisect_left(points, num_tokens, key=lambda point: point[0]) + if index == len(points): + return points[-1][1] + high_tokens, high_cost = points[index] + if index == 0 or high_tokens == num_tokens: + return high_cost + low_tokens, low_cost = points[index - 1] + weighted = low_cost * (high_tokens - num_tokens) + high_cost * ( + num_tokens - low_tokens + ) + return -(-weighted // (high_tokens - low_tokens)) + + +def _log_full_graph_estimate( + measured: int, + mem_samples: Sequence[_FullGraphMemorySample], + graph_descs: Sequence[BatchExecutionDescriptor], + full_estimate: int, + estimate: int, +) -> None: + mib = 1 << 20 + sampled_growth = sum(sample.growth for sample in mem_samples) + logger.info( + "CUDA graph memory profile: capture took %.2f GiB with %d of %d FULL " + "graphs (tokens: grown/kept MiB %s); the other %d FULL graphs add " + "%.2f GiB; estimate %.2f GiB.", + measured / (1 << 30), + len(mem_samples), + len(graph_descs), + ", ".join( + f"{sample.desc.num_tokens}: {sample.growth / mib:+.1f}/" + f"{sample.cost / mib:.1f}" + for sample in mem_samples + ), + len(graph_descs) - len(mem_samples), + (full_estimate - sampled_growth) / (1 << 30), + estimate / (1 << 30), + ) + + def _init_minimal_kv_cache_for_profiling( runner: "GPUModelRunner", *, num_blocks: int | None = None ) -> None: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 2f3baa327880..b22d303afca0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -264,6 +264,18 @@ def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: return module.get_parameter(parameter_name) +def _flushed_free_memory() -> int: + """Read free device memory around a graph memory sample, cache empty. + + ``torch.cuda.graph`` empties the allocator cache on entry. Blocks cached + before the sample would otherwise be released inside it and subtracted + from the graphs' memory, and blocks cached at its end would be added. + """ + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + return torch.accelerator.get_memory_info()[0] + + AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] # list when ubatching is enabled PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict @@ -6898,7 +6910,7 @@ def profile_cudagraph_memory( mem_samples: list[int] = [] for i, desc in enumerate(profile_descs): - mem_before = torch.accelerator.get_memory_info()[0] + mem_before = _flushed_free_memory() self._warmup_and_capture( desc, cudagraph_runtime_mode=mode, @@ -6911,9 +6923,7 @@ def profile_cudagraph_memory( else None ), ) - torch.accelerator.synchronize() - free_after = torch.accelerator.get_memory_info()[0] - mem_samples.append(mem_before - free_after) + mem_samples.append(mem_before - _flushed_free_memory()) first_capture = mem_samples[0] # Use at least 1 MiB per graph for driver overhead @@ -6934,11 +6944,11 @@ def profile_cudagraph_memory( ) if encoder_cudagraph_manager is not None: - mem_before = torch.accelerator.get_memory_info()[0] + mem_before = _flushed_free_memory() encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool) - torch.accelerator.synchronize() - free_after = torch.accelerator.get_memory_info()[0] - encoder_memory_estimate = max(mem_before - free_after, 0) + encoder_memory_estimate = max( + mem_before - _flushed_free_memory(), 0 + ) logger.debug( "Estimated encoder CUDA graph memory: %.2f MiB for %d graphs", From e7fcc19982e31fc334e7822cc0901448b04adb5f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:32:59 +0000 Subject: [PATCH 294/365] docs(release): describe the CUDA graph memory estimate fix Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-930.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-930.json diff --git a/.lil/changes/vllm-930.json b/.lil/changes/vllm-930.json new file mode 100644 index 000000000000..f7249de4306b --- /dev/null +++ b/.lil/changes/vllm-930.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-930", + "category": "fix", + "summary": "The CUDA graph memory estimate counts the memory every captured graph keeps", + "models": [ + "all" + ], + "compatibility": "No user action required. With --kv-cache-memory-bytes the estimate is not used and nothing changes.", + "details": [ + "When --gpu-memory-utilization sizes the KV cache, vLLM captures a few CUDA graphs to estimate the memory of all of them. Each sample was read before torch.cuda.graph() emptied the allocator cache, so memory cached by the eager warmup was subtracted from the graph's own memory. DeepSeek-V4.1 graphs that each kept a private buffer before vLLM #918 were estimated at 1.8 GiB while they used 5.3 GiB, and the oversized KV cache could run out of memory after startup.", + "Samples are now read with the cache emptied, count the memory a graph keeps, and include the smallest graph. A graph that keeps memory now shrinks the KV cache instead of running out of memory later. For current models the KV cache is expected to change by at most about 0.1 GiB per GPU; startup profiles one more graph." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 930 + ], + "requires": [] +} From d58d3777e9594b6f5c8e0442d5b89e792ee2ec67 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:45:00 +0000 Subject: [PATCH 295/365] docs(release): describe the structured-output draft hand-off fix Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-931.json | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 .lil/changes/vllm-931.json diff --git a/.lil/changes/vllm-931.json b/.lil/changes/vllm-931.json new file mode 100644 index 000000000000..bedafd1a6809 --- /dev/null +++ b/.lil/changes/vllm-931.json @@ -0,0 +1,19 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-931", + "category": "fix", + "summary": "JSON-schema and other grammar-constrained requests with MTP or other speculative decoding no longer fail with \"grammar rejected tokens\" after sitting out a step", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "With speculative decoding and async scheduling, a grammar-constrained request can be left out of a step: during a prefill turn while other prompts prefill, or during the isolated step of a restored checkpoint. Its next step then checked the model's draft tokens against grammar rows that allowed every token. The request could emit tokens outside its schema and fail with HTTP 500 (\"Failed to advance FSM ... grammar rejected tokens\"), issue #726.", + "The scheduler now receives exactly the draft tokens each step verifies, so these requests stay constrained and keep their speculative speed-up.", + "A draft token whose grammar row cannot be built is now rejected instead of accepted, and the server logs a one-time warning when this happens." + ], + "requires": [], + "pull_requests": [ + 931 + ] +} From eec23823da4e578a45d6cf18789505d2eeada3a5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:39:31 +0000 Subject: [PATCH 296/365] fix(parser): return GLM reasoning and content with generated whitespace GLM-4.7/5.x chat templates render reasoning verbatim between and and add no whitespace around content or tool calls. The shared parser engine trimmed trailing reasoning whitespace, stripped content that precedes tool calls and dropped whitespace-only content, so a client that returns the message can never re-render the tokens the model generated. The next request then misses the response checkpoint of the hybrid model and prefills the whole previous response again. The GLM parser now returns every generated character before the first tool call, in streaming and non-streaming output. Whitespace after a tool call still never becomes content, because a message places its content before its calls. Other parsers keep dropping whitespace-only content. Co-Authored-By: Claude Opus 5.5 --- tests/parser/engine/trace_builder.py | 3 +- tests/parser/test_glm47_exact_history.py | 225 ++++++++++++++++++ .../test_glm47_moe_tool_parser.py | 6 +- .../tool_parsers/test_glm4_moe_tool_parser.py | 3 +- vllm/parser/abstract_parser.py | 6 +- vllm/parser/engine/adapters.py | 6 + vllm/parser/engine/parser_engine.py | 17 +- vllm/parser/glm47_moe.py | 8 + vllm/tool_parsers/abstract_tool_parser.py | 5 + 9 files changed, 273 insertions(+), 6 deletions(-) create mode 100644 tests/parser/test_glm47_exact_history.py diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 36e347fc572f..f98e3d16cfaa 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -872,7 +872,8 @@ def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample: vocab=_GLM47_MOE_VOCAB, segments=_glm47_moe_segments(scenario), expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "", - expected_content=_qwen3_expected_content(scenario), + # GLM keeps whitespace-only content before calls; templates render it. + expected_content=scenario.content, expected_tool_calls=_expected_tc(scenario), tools=_expected_tools(scenario), ) diff --git a/tests/parser/test_glm47_exact_history.py b/tests/parser/test_glm47_exact_history.py new file mode 100644 index 000000000000..faed6f6061b7 --- /dev/null +++ b/tests/parser/test_glm47_exact_history.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM parser output must re-render to exactly the generated text. + +GLM-4.7/5.x chat templates emit an assistant turn as reasoning between +```` and ````, then content, then tool calls, adding no +whitespace of their own. When the parser returns every generated character, +a template that renders content verbatim reproduces the model's tokens, so +the next request can reuse the cached state of the whole response instead of +prefilling it again. +""" + +import json + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.parser_manager import ParserManager + +GLM_MODEL = "zai-org/GLM-4.7" +QWEN_MODEL = "Qwen/Qwen3-32B" + +TOOLS = [ + { + "type": "function", + "function": { + "name": "read", + "parameters": { + "type": "object", + "properties": { + "filePath": {"type": "string"}, + "limit": {"type": "integer"}, + }, + }, + }, + }, + { + "type": "function", + "function": { + "name": "bash", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + }, + }, + }, +] + +READ = ( + "readfilePath/repo/a.py" + "limit200" +) +BASH = ( + "bashcommand" + "ls -la\n" +) + +# Raw generations that follow the prompt's trailing . +OUTPUTS = { + "answer_trailing_newline": "Plan the fix.Use move_to_end.\n", + "answer_code_block": "Plan.Fixed:\n\n```python\nx = 1\n```\n", + "answer_leading_newlines": "Plan.\n\nThe answer.", + "reasoning_trailing_newlines": "Plan the fix.\n\nDone.", + "empty_reasoning": "Done.\n", + "content_before_tool": "Read it.Let me read the file.\n" + READ, + "whitespace_only_before_tool": "Read it.\n" + READ, + "two_tools": "Check both." + READ + BASH, +} + + +@pytest.fixture(scope="module") +def glm_tokenizer(): + from vllm.tokenizers import get_tokenizer + + return get_tokenizer(GLM_MODEL) + + +@pytest.fixture(scope="module") +def glm_parser_cls(): + return ParserManager.get_parser( + tool_parser_name="glm47", + reasoning_parser_name="glm45", + enable_auto_tools=True, + ) + + +def _request() -> ChatCompletionRequest: + return ChatCompletionRequest.model_validate( + { + "model": "glm", + "messages": [{"role": "user", "content": "hi"}], + "tools": TOOLS, + "tool_choice": "auto", + } + ) + + +def _generated_ids(tokenizer, raw: str) -> list[int]: + stop = "<|observation|>" if "" in raw else "<|user|>" + ids = tokenizer.encode(raw, add_special_tokens=False) + return ids + [tokenizer.convert_tokens_to_ids(stop)] + + +def _render_turn(reasoning, content, calls) -> str: + """The assistant turn after the prompt's , as a GLM template + that renders content verbatim emits it.""" + text = (reasoning or "") + "" + (content or "") + for name, arguments in calls: + text += "" + name + for key, value in json.loads(arguments).items(): + if not isinstance(value, str): + value = json.dumps(value, ensure_ascii=False) + text += f"{key}{value}" + text += "" + return text + + +def _parse(parser_cls, tokenizer, raw): + request = _request() + ids = _generated_ids(tokenizer, raw) + parser = parser_cls(tokenizer, request.tools) + reasoning, content, calls = parser.parse( + tokenizer.decode(ids, skip_special_tokens=True), + request, + enable_auto_tools=True, + model_output_token_ids=ids, + ) + return reasoning, content, [(c.name, c.arguments) for c in calls or []] + + +def _stream(parser_cls, tokenizer, raw, tokens_per_delta): + request = _request() + ids = _generated_ids(tokenizer, raw) + parser = parser_cls(tokenizer, request.tools) + reasoning, content = "", "" + calls: dict[int, list[str]] = {} + previous = "" + for start in range(0, len(ids), tokens_per_delta): + end = min(start + tokens_per_delta, len(ids)) + text = tokenizer.decode(ids[:end], skip_special_tokens=True) + delta = parser.parse_delta( + delta_text=text[len(previous) :], + delta_token_ids=ids[start:end], + request=request, + prompt_token_ids=[tokenizer.convert_tokens_to_ids("")], + finished=end == len(ids), + ) + previous = text + if delta is None: + continue + reasoning += delta.reasoning or "" + content += delta.content or "" + for call in delta.tool_calls or []: + slot = calls.setdefault(call.index, ["", ""]) + if call.function is not None: + slot[0] += call.function.name or "" + slot[1] += call.function.arguments or "" + return ( + reasoning or None, + content or None, + [tuple(calls[index]) for index in sorted(calls)], + ) + + +@pytest.mark.parametrize("name", OUTPUTS) +def test_parse_output_rerenders_generated_text(glm_parser_cls, glm_tokenizer, name): + raw = OUTPUTS[name] + assert _render_turn(*_parse(glm_parser_cls, glm_tokenizer, raw)) == raw + + +@pytest.mark.parametrize("tokens_per_delta", [1, 2, 3, 4]) +@pytest.mark.parametrize("name", OUTPUTS) +def test_streamed_output_matches_parse( + glm_parser_cls, glm_tokenizer, name, tokens_per_delta +): + raw = OUTPUTS[name] + reasoning, content, calls = _stream( + glm_parser_cls, glm_tokenizer, raw, tokens_per_delta + ) + parsed_reasoning, parsed_content, parsed_calls = _parse( + glm_parser_cls, glm_tokenizer, raw + ) + assert (reasoning, content) == (parsed_reasoning, parsed_content) + assert [name for name, _ in calls] == [name for name, _ in parsed_calls] + assert [json.loads(args) for _, args in calls] == [ + json.loads(args) for _, args in parsed_calls + ] + assert _render_turn(reasoning, content, calls) == raw + + +@pytest.mark.parametrize("tokens_per_delta", [None, 1, 3]) +def test_whitespace_after_tool_calls_is_not_content( + glm_parser_cls, glm_tokenizer, tokens_per_delta +): + # Content precedes all calls in a message, so text between or after + # calls cannot be placed faithfully; it must not leak into content. + raw = "Check." + READ + "\n" + BASH + "\n" + if tokens_per_delta is None: + _, content, calls = _parse(glm_parser_cls, glm_tokenizer, raw) + else: + _, content, calls = _stream( + glm_parser_cls, glm_tokenizer, raw, tokens_per_delta + ) + assert content is None + assert [name for name, _ in calls] == ["read", "bash"] + + +def test_other_parsers_still_drop_whitespace_only_content(): + from vllm.tokenizers import get_tokenizer + + tokenizer = get_tokenizer(QWEN_MODEL) + parser_cls = ParserManager.get_parser( + tool_parser_name="qwen3_xml", + reasoning_parser_name="qwen3", + enable_auto_tools=True, + ) + request = _request() + raw = ( + "Plan.\n\n\n\n\n" + "\n/repo/a.py\n\n\n" + ) + parser = parser_cls(tokenizer, request.tools) + _, content, calls = parser.parse(raw, request, enable_auto_tools=True) + assert content is None + assert [call.name for call in calls or []] == ["read"] diff --git a/tests/tool_parsers/test_glm47_moe_tool_parser.py b/tests/tool_parsers/test_glm47_moe_tool_parser.py index dc4b726a5383..6b2b81139eec 100644 --- a/tests/tool_parsers/test_glm47_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm47_moe_tool_parser.py @@ -182,10 +182,12 @@ def test_empty_content_none(self, glm47_tool_parser, mock_request): r = glm47_tool_parser.extract_tool_calls(out, request=mock_request) assert r.content is None - def test_whitespace_content_none(self, glm47_tool_parser, mock_request): + def test_whitespace_content_preserved(self, glm47_tool_parser, mock_request): + # GLM templates render content verbatim before the first call, so + # whitespace the model generated there is kept for the next turn. out = " \n get_current_date" r = glm47_tool_parser.extract_tool_calls(out, request=mock_request) - assert r.content is None + assert r.content == " \n " def _reset(parser): diff --git a/tests/tool_parsers/test_glm4_moe_tool_parser.py b/tests/tool_parsers/test_glm4_moe_tool_parser.py index ca110adac0d2..96f6edffb4a9 100644 --- a/tests/tool_parsers/test_glm4_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm4_moe_tool_parser.py @@ -126,7 +126,8 @@ def test_extract_tool_calls_with_glm45_newline_format(): extracted = parser.extract_tool_calls(model_output, request=_request(tools)) assert extracted.tools_called - assert extracted.content == "I'll check it." + # Content is returned exactly as generated; templates normalize it. + assert extracted.content == "I'll check it. " assert len(extracted.tool_calls) == 1 tool_call = extracted.tool_calls[0] assert tool_call.function.name == "get_current_weather" diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f099d7752ae5..c631d823a2df 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -526,7 +526,11 @@ def _extract_tool_calls( for tc in tool_call_info.tool_calls ) content = tool_call_info.content - if content and content.strip() == "": + if ( + content + and content.strip() == "" + and tool_parser.drops_whitespace_only_content + ): content = None else: # No tool calls. diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index 4f4aa9f5a1e0..3e2488f16cad 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -190,6 +190,12 @@ def skip_reasoning_parsing(self) -> bool: def skip_reasoning_parsing(self, value: bool) -> None: self._parser_engine.skip_reasoning_parsing = value + @property + def drops_whitespace_only_content(self) -> bool: + # Non-streaming output must match what streaming emits for the grammar. + config = self._parser_engine.parser_engine_config + return config.drop_whitespace_only_content_before_tools + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 9e3431c25890..22a16175e111 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -790,6 +790,9 @@ def _events_to_delta( reasoning_parts: list[str] = [] carried_deferred = self._deferred_content + # Content always precedes the tool calls in a message, so only + # whitespace from before the first call can be kept as content. + before_first_tool = not self._tool_slots seen_tool_event = False suppress = self._suppress_tool_calls for event in events: @@ -825,6 +828,18 @@ def _events_to_delta( if len(tool_call_deltas) > 1: tool_call_deltas = self._coalesce_tool_call_deltas(tool_call_deltas) + if ( + carried_deferred + and before_first_tool + and self._tool_slots + and not self._drop_ws_only_content_before_tools + ): + # Whitespace deferred while no tool call existed precedes the + # first one; emit it now rather than behind the call's deltas. + content_parts.insert(0, carried_deferred) + self._deferred_content = self._deferred_content[len(carried_deferred) :] + carried_deferred = "" + if self._deferred_content and (not seen_tool_event or not tool_call_deltas): # Deferred content carried in from a previous delta precedes this # delta's content; content deferred during this call (text after @@ -845,7 +860,7 @@ def _events_to_delta( if stripped: self._content_has_nonws = True elif self._tool_slots: - if self._drop_ws_only_content_before_tools: + if self._drop_ws_only_content_before_tools or not before_first_tool: content_str = "" elif not finished: self._deferred_content = content_str diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 81eafe4264f9..35e0464653b5 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -220,6 +220,14 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: stream_arg_deltas=True, tool_args_json=False, validate_tool_names=True, + # GLM templates put reasoning verbatim between and and + # add no whitespace around content or tool calls. Returning the text + # exactly as generated lets a whitespace-preserving template re-render + # the assistant turn token for token, so the next request can reuse + # the response checkpoint instead of prefilling the whole response. + strip_trailing_reasoning_whitespace=False, + drop_whitespace_only_content_before_tools=False, + strip_content_whitespace_with_tools=False, ) diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 3ee35365e4a8..bf9ed5ab1a4b 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -61,6 +61,11 @@ class ToolParser: structural_tag_model: str | None = None engine_based_streaming: bool = False + @property + def drops_whitespace_only_content(self) -> bool: + """Whether whitespace-only content before tool calls becomes None.""" + return True + def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) if ( From b5dedaf16b387fffd86cd3bd3891edb903159bb3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:54:48 +0000 Subject: [PATCH 297/365] docs(release): record GLM parser whitespace preservation Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-932.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-932.json diff --git a/.lil/changes/vllm-932.json b/.lil/changes/vllm-932.json new file mode 100644 index 000000000000..4ee679eaa4ff --- /dev/null +++ b/.lil/changes/vllm-932.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-932", + "category": "fix", + "summary": "GLM responses keep the whitespace the model generated in reasoning and before tool calls, so the next turn can reuse the cached response", + "models": ["GLM-5.3-Flash"], + "compatibility": "With the glm45/glm47 parsers, reasoning can now end in the newline the model generated, content before tool calls keeps its surrounding whitespace, and whitespace-only content before a tool call is returned instead of null. Nothing changes when the model generates no whitespace there. Other parsers are unchanged.", + "details": [ + "A follow-up request reuses a hybrid model's response checkpoint only when the returned assistant message re-renders to the tokens the model generated. The parser trimmed reasoning, stripped content before tool calls and dropped whitespace-only content, which no chat template can restore, so the whole previous response was prefilled again.", + "Together with the runtime's GLM chat template (docker-110), a CPU reproduction with the deployed GLM-5.3 tokenizer reuses about 99.7% of the next prompt instead of 79% when the previous answer or reasoning ended in a newline or content preceded a tool call on its own line." + ], + "pull_requests": [932], + "requires": [] +} From c9181eb0503dfa7625f13f98e045e4daaea88682 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 21:58:34 +0000 Subject: [PATCH 298/365] test(worker): stub the draft hand-off in the logits-only runner test execute_model now records the drafts each step consumes; the fake runner in test_boundary_logits_only_dispatches_pending_cache_tasks needs the handler and draft buffer it reads. Co-Authored-By: Claude Opus 5.5 --- tests/v1/engine/test_structured_output_draft_handoff.py | 5 ++--- tests/v1/worker/test_gpu_model_runner_v2.py | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/v1/engine/test_structured_output_draft_handoff.py b/tests/v1/engine/test_structured_output_draft_handoff.py index 20fd9e8a7d54..d72b7d048c9a 100644 --- a/tests/v1/engine/test_structured_output_draft_handoff.py +++ b/tests/v1/engine/test_structured_output_draft_handoff.py @@ -328,9 +328,8 @@ def test_skipped_step_keeps_draft_rows_constrained(opt_model_path, consumed_hand def test_deferred_step_without_inflight_batch_samples_immediately(): - """A deferred step may have no earlier batch to wait for: a request that - skipped a step has its output processed already, and only its drafts still - need the back-fill.""" + """A deferred step must not depend on an earlier batch being in flight: + with nothing queued, its drafts are handed back and it is sampled at once.""" done: Future = Future() done.set_result(None) scheduler_output = SchedulerOutput.make_empty() diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 748a2e68bfb4..2292f3f5393a 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -425,6 +425,8 @@ def test_boundary_logits_only_dispatches_pending_cache_tasks( gather_batch_req_state=lambda *args: (SimpleNamespace(num_tokens=1), 1), prepare_inputs=lambda *args: input_batch, prepare_attn=lambda *args: pytest.fail("logits-only must skip attention"), + draft_tokens_handler=SimpleNamespace(set_consumed_draft_tokens=lambda *a: None), + req_states=SimpleNamespace(draft_tokens=None), ) monkeypatch.setattr( model_runner_module, From 3fbef7291851a4a336002ff8f50741226dfae237 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 22:22:14 +0000 Subject: [PATCH 299/365] revert: the DeepSeek-V4.1 compressor-ring change (#927) had no effect Reverts #927 (port of #926). CircularBufferSpec groups use uses_slot_mapping=False, so the runner calls the slot-mapping kernel with mapping disabled for ring groups and pads the whole group; the new mapping path is never taken. Measured on DS4.1 DSpark (decode-vs-prefill KL by position over 2K tokens): .046/.053/.068/.085/.127 without and .037/.053/.064/.097/.152 with the change, i.e. the same drift. The release note of beta 6e8298a1 claimed a fix that does not happen. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-926.json | 17 ----------- tests/v1/worker/test_gpu_block_table.py | 38 ------------------------- vllm/v1/worker/gpu/block_table.py | 15 +--------- vllm/v1/worker/gpu/model_runner.py | 4 --- 4 files changed, 1 insertion(+), 73 deletions(-) delete mode 100644 .lil/changes/vllm-926.json diff --git a/.lil/changes/vllm-926.json b/.lil/changes/vllm-926.json deleted file mode 100644 index 5fdc9511ddda..000000000000 --- a/.lil/changes/vllm-926.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "schema": "local-inference-release-change/v1", - "id": "vllm-926", - "category": "fix", - "summary": "Keep DeepSeek-V4.1 compressor state across decode steps past the ring size", - "models": [ - "DeepSeek-V4.1-Flash" - ], - "compatibility": "No user action required.", - "details": [ - "Circular-buffer KV groups map every position into their one block, so the compressor state ring is written past its first 8–16 tokens and decode no longer drifts from the model as generation lengthens." - ], - "pull_requests": [ - 926 - ], - "requires": [] -} diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5d144b5daafc..1aa1da0eafef 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -378,41 +378,3 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() - - -def test_block_tables_circular_group_maps_positions_into_its_ring(): - """A circular-buffer group holds one block per request as a ring, so positions at - or past its block size still map into that block (position modulo the block size). - A request without a ring block keeps padding slots; positional groups are - unchanged.""" - device = torch.device("cuda") - block_tables = BlockTables( - block_sizes=[16, 8], - max_num_reqs=2, - max_num_batched_tokens=64, - max_num_blocks_per_group=[4, 1], - device=device, - kernel_block_sizes=[16, 8], - slot_mapping_circular=[False, True], - ) - block_tables.append_block_ids( - req_index=0, new_block_ids=([3, 4, 5], [9]), overwrite=True - ) - block_tables.append_block_ids(req_index=1, new_block_ids=([6], []), overwrite=True) - block_tables.apply_staged_writes() - - positions = list(range(3, 40)) + [5, 6] - idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device) - query_start_loc = torch.tensor([0, 37, 39], dtype=torch.int32, device=device) - slot_mappings = block_tables.compute_slot_mappings( - idx_mapping, - query_start_loc, - torch.tensor(positions, dtype=torch.int64, device=device), - num_tokens_padded=len(positions), - ) - torch.accelerator.synchronize() - - ring = [9 * 8 + p % 8 for p in positions[:37]] - assert slot_mappings[1].tolist() == ring + [-1, -1] - paged = [[3, 4, 5][p // 16] * 16 + p % 16 for p in positions[:37]] - assert slot_mappings[0].tolist() == paged + [6 * 16 + 5, 6 * 16 + 6] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 5aa9374eae60..2c40aa28a459 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -28,7 +28,6 @@ def __init__( cp_interleave: int = 1, slot_mapping_enabled: list[bool] | None = None, group_cp_sizes: list[int] | None = None, - slot_mapping_circular: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -51,12 +50,6 @@ def __init__( slot_mapping_enabled = [True] * self.num_kv_cache_groups assert len(slot_mapping_enabled) == self.num_kv_cache_groups self._slot_mapping_enabled = slot_mapping_enabled - # A circular group keeps one block per request as a ring: every position maps - # into that block at position modulo the block size. - if slot_mapping_circular is None: - slot_mapping_circular = [False] * self.num_kv_cache_groups - assert len(slot_mapping_circular) == self.num_kv_cache_groups - self._slot_mapping_circular = slot_mapping_circular self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -121,9 +114,6 @@ def init_block_table_layout_tensors(self) -> None: self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) - self.slot_mapping_circular = torch.tensor( - self._slot_mapping_circular, dtype=torch.bool, device=self.device - ) self.group_cp_sizes = torch.tensor( self.group_cp_sizes_list, dtype=torch.int32, device=self.device ) @@ -238,7 +228,6 @@ def compute_slot_mappings( self.block_sizes_tensor, self.kernel_block_sizes_tensor, self.slot_mapping_enabled, - self.slot_mapping_circular, self.group_cp_sizes, slot_mappings, slot_mappings.stride(0), @@ -319,7 +308,6 @@ def _compute_slot_mappings_kernel( block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] - slot_mapping_circular, # [num_kv_cache_groups] group_cp_sizes, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, @@ -351,7 +339,6 @@ def _compute_slot_mappings_kernel( kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) - circular = tl.load(slot_mapping_circular + group_id) group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) @@ -379,7 +366,7 @@ def _compute_slot_mappings_kernel( local_positions = virtual_block_indices * kv_block_size + local_offsets block_indices = tl.where( - mapping_enabled & ~circular, local_positions // kernel_block_size, 0 + mapping_enabled, local_positions // kernel_block_size, 0 ) block_offsets = local_positions % kernel_block_size valid_block = token_mask & (block_indices < num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 55a2dfa07462..4a478494054c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -68,7 +68,6 @@ from vllm.v1.core.boundary_checkpoint import NUM_BOUNDARY_CHECKPOINT_SLOTS from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( - CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -618,7 +617,6 @@ def initialize_kv_cache( block_sizes = [] max_num_blocks_per_group = [] slot_mapping_enabled = [] - slot_mapping_circular = [] group_cp_sizes = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec @@ -627,7 +625,6 @@ def initialize_kv_cache( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) slot_mapping_enabled.append(layer_spec.uses_slot_mapping) - slot_mapping_circular.append(isinstance(layer_spec, CircularBufferSpec)) group_cp_sizes.append( 1 if getattr(layer_spec, "dcp_replicated", False) else self.dcp_size ) @@ -710,7 +707,6 @@ def initialize_kv_cache( device=self.device, kernel_block_sizes=self.kernel_block_sizes, slot_mapping_enabled=slot_mapping_enabled, - slot_mapping_circular=slot_mapping_circular, cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, From 0aa6de75ae55f8765386a82b9d4d02c0931d08de Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 22:22:44 +0000 Subject: [PATCH 300/365] docs(release): correct the DeepSeek-V4.1 drift note Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-933.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-933.json diff --git a/.lil/changes/vllm-933.json b/.lil/changes/vllm-933.json new file mode 100644 index 000000000000..7b6342bed0fe --- /dev/null +++ b/.lil/changes/vllm-933.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-933", + "category": "fix", + "summary": "Correction: DeepSeek-V4.1 decode drift is not fixed; the compressor-ring change from beta 6e8298a1 had no effect and is removed", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required. Output is the same as with the previous betas.", + "details": [ + "The note \"Keep DeepSeek-V4.1 compressor state across decode steps past the ring size\" (vllm #926/#927) was wrong: ring groups do not use the slot mapping that change edited, so it never ran. Measured decode-vs-prefill divergence over 2K tokens was the same with and without it. The drift remains under investigation." + ], + "pull_requests": [ + 933 + ], + "requires": [] +} From 296220595ba5fccfe137915b2dd4a2947a5cb8e2 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 22:24:08 +0000 Subject: [PATCH 301/365] fix(tokenizers): max_token_id is vocab_size - 1, not vocab_size get_cached_tokenizer took max(max id in get_vocab(), vocab_size). vocab_size counts tokens, so this admitted one id past the vocabulary whenever vocab_size is not smaller than the highest id, e.g. with fastokens (VLLM_USE_FASTOKENS=1), whose vocab_size includes added tokens: on DeepSeek-V4.x the prompt validation then accepted token 129280, one row past the embedding table. Co-Authored-By: Claude Opus 5.5 --- tests/tokenizers_/test_hf_max_token_id.py | 50 +++++++++++++++++++++++ vllm/tokenizers/hf.py | 7 +++- 2 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 tests/tokenizers_/test_hf_max_token_id.py diff --git a/tests/tokenizers_/test_hf_max_token_id.py b/tests/tokenizers_/test_hf_max_token_id.py new file mode 100644 index 000000000000..2795b1055be1 --- /dev/null +++ b/tests/tokenizers_/test_hf_max_token_id.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""max_token_id is the highest valid id, never the vocabulary size.""" + +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from transformers import PreTrainedTokenizerFast + +from vllm.tokenizers.hf import get_cached_tokenizer + + +def _tokenizer(cls=PreTrainedTokenizerFast) -> PreTrainedTokenizerFast: + backend = Tokenizer(WordLevel({"a": 0, "b": 1, "c": 2}, unk_token="a")) + tokenizer = cls(tokenizer_object=backend) + tokenizer.add_special_tokens({"additional_special_tokens": [""]}) # id 3 + return tokenizer + + +class _CountsAddedTokens(PreTrainedTokenizerFast): + """A backend whose vocab_size includes added tokens, like fastokens.""" + + @property + def vocab_size(self) -> int: + return len(self.get_vocab()) + + +class _HidesSpecialTokens(PreTrainedTokenizerFast): + """Special ids beyond get_vocab() that vocab_size counts (QwenTokenizer).""" + + @property + def vocab_size(self) -> int: + return 6 + + def get_vocab(self) -> dict[str, int]: + return {"a": 0, "b": 1, "c": 2} + + +def test_max_token_id_is_the_highest_id(): + assert get_cached_tokenizer(_tokenizer()).max_token_id == 3 + + +def test_vocab_size_counting_added_tokens_admits_no_extra_id(): + tokenizer = _tokenizer(_CountsAddedTokens) + assert tokenizer.vocab_size == 4 + assert get_cached_tokenizer(tokenizer).max_token_id == 3 + + +def test_ids_only_vocab_size_knows_about_are_still_covered(): + tokenizer = _tokenizer(_HidesSpecialTokens) + assert get_cached_tokenizer(tokenizer).max_token_id == 5 diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py index 3d35952c5f4c..dd8cbbc4ae7e 100644 --- a/vllm/tokenizers/hf.py +++ b/vllm/tokenizers/hf.py @@ -126,10 +126,13 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: # Some tokenizers (e.g., QwenTokenizer) have special tokens that # are added and included in the implementation of the vocab_size # property, but not in get_vocab(); if there is an implementation - # of vocab size, we should take the greater value. + # of vocab size, we should take the greater value. vocab_size counts + # tokens, so its highest id is vocab_size - 1: taking vocab_size itself + # admitted one id past the vocabulary when a backend (e.g. fastokens) + # counts the added tokens in vocab_size. if hasattr(tokenizer, "vocab_size"): with contextlib.suppress(NotImplementedError): - max_token_id = max(max_token_id, tokenizer.vocab_size) + max_token_id = max(max_token_id, tokenizer.vocab_size - 1) class CachedTokenizer(tokenizer.__class__): # type: ignore @property From f9cf8e2d0e9514be9aa9702f22c5be0a40611722 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 22:24:20 +0000 Subject: [PATCH 302/365] docs(release): record the max_token_id fix Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-934.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-934.json diff --git a/.lil/changes/vllm-934.json b/.lil/changes/vllm-934.json new file mode 100644 index 000000000000..219670edbd56 --- /dev/null +++ b/.lil/changes/vllm-934.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-934", + "category": "fix", + "summary": "Prompts with a token id one past the vocabulary are rejected again when the tokenizer counts added tokens in its size", + "models": [ + "all" + ], + "compatibility": "No user action required.", + "details": [ + "The largest accepted prompt token id was taken from the tokenizer's vocabulary size instead of size minus one. With VLLM_USE_FASTOKENS=1, whose size includes added tokens, DeepSeek-V4.x accepted token 129280, one past the embedding table; such prompts now get a validation error." + ], + "pull_requests": [ + 934 + ], + "requires": [] +} From 7b2fc4e1db400a9d93e5194b56f47480f804a8c1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 22:54:27 +0000 Subject: [PATCH 303/365] test(parser): cover GLM exact-history turns through the serving API Drive OpenAIServingChat end to end with a server-supplied GLM chat template, the glm45/glm47 parsers, return_token_ids, streaming and non-streaming, and send the answer back like a client: the next prompt must extend the previous prompt and every generated token but the stop token. Seeded random GLM turns also round-trip through parse and parse_delta with 1-5 tokens per delta. Co-Authored-By: Claude Opus 5.5 --- .../test_glm_exact_history_serving.py | 308 ++++++++++++++++++ tests/parser/test_glm47_exact_history.py | 22 ++ 2 files changed, 330 insertions(+) create mode 100644 tests/entrypoints/openai/chat_completion/test_glm_exact_history_serving.py diff --git a/tests/entrypoints/openai/chat_completion/test_glm_exact_history_serving.py b/tests/entrypoints/openai/chat_completion/test_glm_exact_history_serving.py new file mode 100644 index 000000000000..2f0182e70108 --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_glm_exact_history_serving.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A returned GLM assistant turn re-renders to the tokens the model generated. + +Drives OpenAIServingChat end to end (server-supplied chat template, glm45 and +glm47 parsers, streaming and non-streaming, return_token_ids) with a fake +engine, then sends the message back like an OpenAI client. The next prompt +must start with the previous prompt and every generated token except the +stop token, which is where a hybrid model's response checkpoint ends. +""" + +import asyncio +import json +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from vllm.config import MultiModalConfig +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.models.serving import ( + BaseModelPath, + OpenAIServingModels, +) +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.renderers.hf import HfRenderer +from vllm.renderers.online_renderer import OnlineRenderer +from vllm.sampling_params import RequestOutputKind +from vllm.tokenizers.registry import cached_tokenizer_from_config +from vllm.v1.engine.async_llm import AsyncLLM + +MODEL = "zai-org/GLM-4.7" +# The GLM-5.x layout with assistant content rendered as generated, as the +# runtime's GLM template does. +TEMPLATE = ( + "[gMASK]" + "{%- for m in messages -%}" + "{%- if m.role == 'system' -%}<|system|>{{ m.content }}" + "{%- elif m.role == 'user' -%}<|user|>{{ m.content }}" + "{%- elif m.role == 'assistant' -%}<|assistant|>" + "{%- if m.reasoning_content is string -%}" + "{{ '' + m.reasoning_content + '' }}" + "{%- else -%}{%- endif -%}" + "{{ m.content or '' }}" + "{%- for tc in m.tool_calls or [] -%}{{ tc.function.name }}" + "{%- for k, v in tc.function.arguments.items() -%}" + "{{ k }}" + "{{ v if v is string else v | tojson }}" + "{%- endfor -%}" + "{%- endfor -%}" + "{%- elif m.role == 'tool' -%}" + "{%- if loop.first or messages[loop.index0 - 1].role != 'tool' -%}" + "<|observation|>{%- endif -%}" + "{{ m.content }}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}<|assistant|>{%- endif -%}" +) +TOOLS = [ + { + "type": "function", + "function": { + "name": "read_file", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string"}, + "limit": {"type": "integer"}, + }, + }, + }, + } +] +READ = ( + "read_filepathsrc/cache.py" + "limit50" +) +# Generated text after the prompt's ; each ends a turn differently. +GENERATIONS = { + "answer_trailing_newline": "Plan.```python\nx = 1\n```\n", + "reasoning_trailing_newline": "Plan.\n\nDone.", + "content_before_tool": "Read first.Let me read it.\n" + READ, + "whitespace_before_tool": "Read first.\n" + READ, +} + + +@dataclass +class _HFConfig: + model_type: str = "glm4_moe" + + +@dataclass +class _ModelConfig: + task = "generate" + runner_type = "generate" + model = MODEL + tokenizer = MODEL + trust_remote_code = False + tokenizer_mode = "auto" + max_model_len = 8192 + revision = None + code_revision = None + tokenizer_revision = None + multimodal_config = MultiModalConfig() + hf_config = _HFConfig() + hf_text_config = _HFConfig() + logits_processors: list[str] | None = None + diff_sampling_param: dict | None = None + allowed_local_media_path: str = "" + allowed_media_domains: list[str] | None = None + encoder_config = None + generation_config: str = "auto" + override_generation_config: dict[str, Any] = field(default_factory=dict) + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) + skip_tokenizer_init: bool = False + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + enable_prompt_embeds: bool = False + + def get_diff_sampling_param(self): + return self.diff_sampling_param or {} + + +@dataclass +class _ParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class _VllmConfig: + model_config: _ModelConfig + parallel_config: _ParallelConfig + + +class _FakeEngine: + """Replays one canned generation per request and records the prompts.""" + + def __init__(self, tokenizer) -> None: + self.tokenizer = tokenizer + self.generation = "" + self.prompts: list[list[int]] = [] + self.outputs: list[list[int]] = [] + + async def generate(self, prompt, sampling_params, request_id, *args, **kwargs): + tok = self.tokenizer + prompt_ids = list(prompt["prompt_token_ids"]) + stop = "<|observation|>" if "" in self.generation else "<|user|>" + ids = tok.encode(self.generation, add_special_tokens=False) + ids.append(tok.convert_tokens_to_ids(stop)) + self.prompts.append(prompt_ids) + self.outputs.append(ids) + + def text(upto: int) -> str: + # The stop token finishes the request; its text is not emitted. + return tok.decode(ids[: min(upto, len(ids) - 1)], skip_special_tokens=False) + + def output(token_ids, delta_text, finished) -> RequestOutput: + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=prompt_ids, + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=delta_text, + token_ids=token_ids, + cumulative_logprob=None, + logprobs=None, + finish_reason="stop" if finished else None, + stop_reason=ids[-1] if finished else None, + ) + ], + finished=finished, + ) + + if sampling_params.output_kind != RequestOutputKind.DELTA: + yield output(ids, text(len(ids)), True) + return + previous = "" + for start in range(0, len(ids), 3): + end = min(start + 3, len(ids)) + current = text(end) + yield output(ids[start:end], current[len(previous) :], end == len(ids)) + previous = current + + +@pytest.fixture(scope="module") +def tokenizer(): + return cached_tokenizer_from_config(_ModelConfig()) + + +def _serving(tokenizer) -> tuple[OpenAIServingChat, _FakeEngine]: + model_config = _ModelConfig() + engine = MagicMock(spec=AsyncLLM) + engine.errored = False + engine.model_config = model_config + engine.input_processor = MagicMock() + engine.renderer = HfRenderer( + _VllmConfig(model_config, _ParallelConfig()), tokenizer + ) + fake = _FakeEngine(tokenizer) + engine.generate = fake.generate + models = OpenAIServingModels( + engine_client=engine, + base_model_paths=[BaseModelPath(name=MODEL, model_path=MODEL)], + ) + online = OnlineRenderer( + model_config=model_config, + renderer=engine.renderer, + request_logger=None, + chat_template=TEMPLATE, + chat_template_content_format="string", + enable_auto_tools=True, + tool_parser="glm47", + reasoning_parser="glm45", + ) + serving = OpenAIServingChat( + engine, + models, + response_role="assistant", + online_renderer=online, + chat_template=TEMPLATE, + chat_template_content_format="string", + request_logger=None, + reasoning_parser="glm45", + tool_parser="glm47", + enable_auto_tools=True, + ) + return serving, fake + + +async def _chat(serving: OpenAIServingChat, body: dict) -> dict: + response = await serving.create_chat_completion(ChatCompletionRequest(**body)) + assert not hasattr(response, "error"), response + if not body.get("stream"): + return json.loads(response.model_dump_json())["choices"][0]["message"] + message: dict[str, Any] = {"content": "", "reasoning": ""} + calls: dict[int, dict] = {} + async for line in response: + payload = line.strip().removeprefix("data: ") + if not payload or payload == "[DONE]": + continue + for choice in json.loads(payload)["choices"]: + delta = choice["delta"] + message["content"] += delta.get("content") or "" + message["reasoning"] += delta.get("reasoning") or "" + for call in delta.get("tool_calls") or []: + slot = calls.setdefault( + call["index"], {"id": None, "name": "", "arguments": ""} + ) + slot["id"] = call.get("id") or slot["id"] + function = call.get("function") or {} + slot["name"] += function.get("name") or "" + slot["arguments"] += function.get("arguments") or "" + message["tool_calls"] = [ + { + "id": call["id"], + "type": "function", + "function": {"name": call["name"], "arguments": call["arguments"]}, + } + for _, call in sorted(calls.items()) + ] + return message + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("name", GENERATIONS) +def test_next_prompt_extends_the_generated_turn(tokenizer, name, stream): + asyncio.run(_two_turns(tokenizer, name, stream)) + + +async def _two_turns(tokenizer, name: str, stream: bool) -> None: + serving, engine = _serving(tokenizer) + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are a coding agent."}, + {"role": "user", "content": "Why is the new cache entry evicted?"}, + ] + body: dict[str, Any] = { + "model": MODEL, + "messages": messages, + "tools": TOOLS, + "stream": stream, + "return_token_ids": True, + } + engine.generation = GENERATIONS[name] + message = await _chat(serving, body) + + returned: dict[str, Any] = {"role": "assistant", "content": message["content"]} + if message.get("reasoning"): + returned["reasoning_content"] = message["reasoning"] + if message.get("tool_calls"): + returned["tool_calls"] = message["tool_calls"] + follow = [ + {"role": "tool", "tool_call_id": call["id"], "content": "x = 1"} + for call in message["tool_calls"] + ] + else: + follow = [{"role": "user", "content": "And in get()?"}] + body["messages"] = [*messages, returned, *follow] + await _chat(serving, body) + + first_prompt, first_output = engine.prompts[0], engine.outputs[0] + generated = first_prompt + first_output[:-1] + assert engine.prompts[1][: len(generated)] == generated diff --git a/tests/parser/test_glm47_exact_history.py b/tests/parser/test_glm47_exact_history.py index faed6f6061b7..2a32f2551a2b 100644 --- a/tests/parser/test_glm47_exact_history.py +++ b/tests/parser/test_glm47_exact_history.py @@ -11,6 +11,7 @@ """ import json +import random import pytest @@ -205,6 +206,27 @@ def test_whitespace_after_tool_calls_is_not_content( assert [name for name, _ in calls] == ["read", "bash"] +def _random_turn(rng: random.Random) -> str: + """A GLM turn whose every character the message format can represent.""" + pieces = ["Plan", " it", ".", "\n", "\n\n", " ", "```python\nx = 1\n```", "`y`"] + + def text() -> str: + return "".join(rng.choice(pieces) for _ in range(rng.randint(0, 6))) + + turn = (text() + "" if rng.random() < 0.9 else "") + text() + return turn + "".join(rng.choice([READ, BASH]) for _ in range(rng.randint(0, 2))) + + +@pytest.mark.parametrize("seed", range(40)) +def test_random_turns_round_trip(glm_parser_cls, glm_tokenizer, seed): + rng = random.Random(seed) + raw = _random_turn(rng) + parsed = _parse(glm_parser_cls, glm_tokenizer, raw) + assert _render_turn(*parsed) == raw + streamed = _stream(glm_parser_cls, glm_tokenizer, raw, rng.randint(1, 5)) + assert _render_turn(*streamed) == raw + + def test_other_parsers_still_drop_whitespace_only_content(): from vllm.tokenizers import get_tokenizer From 0e4d5b7c6030a807d168c2ed82bb03c403fb422e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 00:38:33 +0000 Subject: [PATCH 304/365] docs(release): restore the published vllm-926 fragment Published change fragments are immutable: the release changelog refuses a channel whose earlier fragments were deleted, and the revert in #933 removed vllm-926.json. The correction lives in vllm-933. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-926.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .lil/changes/vllm-926.json diff --git a/.lil/changes/vllm-926.json b/.lil/changes/vllm-926.json new file mode 100644 index 000000000000..5fdc9511ddda --- /dev/null +++ b/.lil/changes/vllm-926.json @@ -0,0 +1,17 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-926", + "category": "fix", + "summary": "Keep DeepSeek-V4.1 compressor state across decode steps past the ring size", + "models": [ + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required.", + "details": [ + "Circular-buffer KV groups map every position into their one block, so the compressor state ring is written past its first 8–16 tokens and decode no longer drifts from the model as generation lengthens." + ], + "pull_requests": [ + 926 + ], + "requires": [] +} From 94fdf0707c39d1d1340b03237021f7914661f258 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 09:16:12 +0000 Subject: [PATCH 305/365] Run Kimi-K3 B12X KDA prefill against the recurrent state pool Expose --kda-prefill-backend b12x through prepared state-pool bindings and caller-owned scratch. Avoid recurrent-state gather/scatter. Auto selection and the state/cache contract remain unchanged. Validation: KDA implementation AST matches the qualified source revision 091cdc82ef. The source commit records seven SM120 correctness cases. No GPU test was run in this audit environment. Co-authored-by: Martin Vit Co-authored-by: Claude Opus 5.5 Co-authored-by: Codex --- .lil/changes/vllm-kimi-b12x-kda-prefill.json | 18 + tests/models/kimi_k3/test_kda.py | 192 +++++++++ vllm/models/kimi_k3/nvidia/kda.py | 411 ++++++++++++++++++- 3 files changed, 611 insertions(+), 10 deletions(-) create mode 100644 .lil/changes/vllm-kimi-b12x-kda-prefill.json diff --git a/.lil/changes/vllm-kimi-b12x-kda-prefill.json b/.lil/changes/vllm-kimi-b12x-kda-prefill.json new file mode 100644 index 000000000000..69589a449c8d --- /dev/null +++ b/.lil/changes/vllm-kimi-b12x-kda-prefill.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-kimi-b12x-kda-prefill", + "category": "feature", + "summary": "Run Kimi-K3 KDA prefill directly against the recurrent state pool", + "models": [ + "Kimi-K3" + ], + "compatibility": "Select --kda-prefill-backend b12x; auto backend selection is unchanged.", + "details": [ + "Prepares caller-owned scratch before execution and avoids gathering and scattering recurrent states." + ], + "pull_requests": [], + "requires": [], + "authors": [ + "voipmonitor" + ] +} diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 67282a095005..4824ffad899d 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -29,6 +29,7 @@ _flashinfer_kda_prefill, _flashkda_prefill, _store_cache_checkpoints_kernel, + is_b12x_kda_prefill_supported, is_flashinfer_fused_kda_decode_supported, is_flashinfer_recurrent_kda_prefill_supported, is_flashkda_supported, @@ -50,6 +51,7 @@ ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.utils.b12x import get_b12x_kda_prefill from vllm.utils.flashinfer import flashinfer_fused_kda_decode from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -1366,6 +1368,10 @@ def _require_kda_prefill_backend( supported = is_flashinfer_recurrent_kda_prefill_supported( 128, torch.bfloat16, state_dtype, lower_bound ) + elif backend == "b12x": + supported = is_b12x_kda_prefill_supported( + 128, torch.bfloat16, state_dtype, lower_bound + ) else: assert backend == "flashkda" supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound) @@ -1411,6 +1417,23 @@ def _run_kda_prefill_backend( seq_order=seq_order, ) + if backend == "b12x": + return _b12x_kda_prefill( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + has_initial_state=torch.ones( + initial_state.shape[0], dtype=torch.bool, device=q.device + ), + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + assert backend == "flashkda" import vllm._flashkda_C # noqa: F401 @@ -1439,6 +1462,110 @@ def _run_kda_prefill_backend( ) +def _b12x_kda_prefill( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + lower_bound: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run b12x prefill with the layer's slot conventions. + + Slot ``NULL_BLOCK_ID`` is the null slot. Request ``r`` owns slot ``r + 1`` + of a state pool, reads it only when ``has_initial_state[r]`` is set, and + receives its final state there; no checkpoint slot is written. + """ + from b12x.preparation import PreparationSession, PreparedCall, require_prepared + + api = get_b12x_kda_prefill() + assert api is not None + assert NULL_BLOCK_ID == 0 + num_seqs = cu_seqlens.numel() - 1 + tokens, heads = q.shape[1], q.shape[2] + pool = torch.zeros( + (num_seqs + 1, *initial_state.shape[1:]), + dtype=initial_state.dtype, + device=q.device, + ) + pool[1:] = initial_state + final_indices = torch.arange(1, num_seqs + 1, dtype=torch.int32, device=q.device) + initial_indices = final_indices.masked_fill(~has_initial_state, NULL_BLOCK_ID) + tensors = { + "q": q[0], + "k": k[0], + "v": v[0], + "raw_g": raw_g[0], + "raw_beta": raw_beta[0], + "A_log": A_log, + "dt_bias": dt_bias, + "recurrent_state": pool, + "cu_seqlens": cu_seqlens, + "initial_state_indices": initial_indices, + "final_state_indices": final_indices, + "checkpoint_state_indices": torch.full_like(final_indices, NULL_BLOCK_ID), + "checkpoint_offsets": torch.zeros_like(final_indices), + "num_seqs": torch.tensor([num_seqs], dtype=torch.int32, device=q.device), + "num_tokens": torch.tensor([tokens], dtype=torch.int32, device=q.device), + "output": torch.empty_like(v[0]), + } + plan = api.plan( + api.Caps( + device=q.device, + max_tokens=tokens, + max_seqs=num_seqs, + max_state_slots=num_seqs + 1, + heads=heads, + null_state_index=NULL_BLOCK_ID, + ), + invocation=api.invocation_from_tensors( + A_log=A_log, dt_bias=dt_bias, initial_state_indices=initial_indices + ), + ) + original_pool = pool.clone() + + def prepare_call(state): + (spec,) = state.layout.scratch_specs() + binding = state.bind( + scratch=torch.empty(spec.shape, dtype=spec.dtype, device=q.device), + **tensors, + ) + + def reset(): + pool.copy_(original_pool) + + return PreparedCall( + run=lambda: state.run(binding, lower_bound=lower_bound), + reset=reset, + restore=reset, + ) + + session = PreparationSession(device=q.device, autotune=False) + result = session.prepare( + (plan.request(name="kimi-k3-kda-prefill", prepare_call=prepare_call),) + ) + try: + (spec,) = require_prepared(plan, "sequence.kda_prefill").layout.scratch_specs() + pool.copy_(original_pool) + binding = api.bind( + plan, + scratch=torch.empty(spec.shape, dtype=spec.dtype, device=q.device), + **tensors, + ) + api.run(binding, lower_bound=lower_bound) + torch.accelerator.synchronize() + finally: + result.close() + session.close() + return tensors["output"].unsqueeze(0), pool[1:].clone() + + def _kda_prefill_reference( inputs: SimpleNamespace, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1477,6 +1604,7 @@ def _kda_prefill_reference( [ pytest.param("flashkda", torch.float32, id="flashkda"), pytest.param("flashinfer", torch.bfloat16, id="flashinfer"), + pytest.param("b12x", torch.float32, id="b12x"), ], ) @torch.inference_mode() @@ -1524,6 +1652,7 @@ def test_kda_prefill_near_collinear_keys_remain_finite( pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"), pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"), pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"), + pytest.param("b12x", torch.float32, 0.01, id="b12x-fp32"), ], ) @torch.inference_mode() @@ -1545,6 +1674,69 @@ def test_kda_prefill_correctness( assert_close("ht", expected_state, actual_state, tolerance) +@torch.inference_mode() +def test_b12x_kda_prefill_matches_triton_serving_batch(): + """Compare b12x with Triton on a TP9-shaped 4,096-token prefill batch. + + Eleven local heads, ragged requests including a one-token row, and one + request without a computed prefix; b12x reads and writes state slots in + place while Triton consumes gathered initial states. + """ + lower_bound = -5.0 + _require_kda_prefill_backend("b12x", torch.float32, lower_bound) + torch.manual_seed(7) + H, D = 11, 128 + lengths = [1, 700, 1531, 1864] + T = sum(lengths) + q, k, v, raw_g = [ + torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) + ] + raw_beta = torch.randn(1, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = ( + torch.randn(len(lengths), H, D, D, dtype=torch.float32, device=DEVICE) * 0.1 + ) + has_initial_state = torch.tensor( + [True, False, True, True], dtype=torch.bool, device=DEVICE + ) + cu_seqlens = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEVICE) + cu_seqlens[1:] = torch.tensor(lengths, device=DEVICE).cumsum(0) + + b12x_out, b12x_state = _b12x_kda_prefill( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + has_initial_state=has_initial_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + triton_out, triton_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=dt_bias.reshape(-1), + lower_bound=lower_bound, + initial_state=initial_state * has_initial_state[:, None, None, None], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + assert torch.isfinite(b12x_out).all() + assert torch.isfinite(b12x_state).all() + assert_close("o", triton_out, b12x_out, 0.01) + assert_close("ht", triton_state, b12x_state, 0.01) + + @torch.inference_mode() def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): if not is_flashinfer_recurrent_kda_prefill_supported( diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index af7a25806b03..e43cbf77dfc0 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import replace +from typing import Any import torch from einops import rearrange @@ -60,6 +61,13 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.triton_utils import tl, triton +from vllm.utils.b12x import ( + B12xPreparationUnit, + B12xWorkload, + PreparationResourceUnavailableError, + get_b12x_kda_prefill, + set_b12x_preparation_provider, +) from vllm.utils.flashinfer import ( flashinfer_fused_kda_decode, flashinfer_recurrent_kda, @@ -328,6 +336,24 @@ def is_flashkda_supported( ) +def is_b12x_kda_prefill_supported( + head_dim: int, + input_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + api = get_b12x_kda_prefill() + if api is None or not current_platform.is_cuda(): + return False + return ( + head_dim == 128 + and input_dtype == torch.bfloat16 + and recurrent_state_dtype == torch.float32 + and lower_bound is not None + and api.is_supported(torch.device(current_platform.current_device())) + ) + + def is_flashinfer_recurrent_kda_prefill_supported( head_dim: int, input_dtype: torch.dtype, @@ -501,8 +527,26 @@ def resolve_kda_prefill_backend( recurrent_state_dtype: torch.dtype, lower_bound: float | None, ) -> str: - if backend not in ("auto", "triton", "flashkda", "flashinfer"): + """Resolve the packed KDA prefill implementation for one server. + + ``auto`` never selects ``b12x``; that backend must be requested by name. + """ + if backend not in ("auto", "triton", "flashkda", "flashinfer", "b12x"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") + if backend == "b12x": + if not is_b12x_kda_prefill_supported( + head_dim, + input_dtype, + recurrent_state_dtype, + lower_bound, + ): + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package on a " + "supported CUDA device, bfloat16 activations, head_dim=128, " + "float32 recurrent state, and a bounded KDA gate." + ) + logger.info_once("Using b12x KDA prefill backend.") + return "b12x" flashinfer_supported = is_flashinfer_recurrent_kda_prefill_supported( head_dim, input_dtype, @@ -872,6 +916,9 @@ def __init__( (1, T, H, D), self.model_config.dtype, ) + self._b12x_prefill_api: Any | None = None + self._b12x_prefill_plan = None + self._initialize_b12x_kda_prefill(vllm_config) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None @@ -911,11 +958,327 @@ def __init__( self.gemm_rs_ar = gemm_rs_ar else: gemm_rs_ar.warn_incompatible_projection() + self._b12x_preparation_prefix = prefix + if self._b12x_prefill_api is not None: + set_b12x_preparation_provider(self, self) compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + super().bind_kv_cache(kv_cache) + self._b12x_prefill_plan = None + if self._b12x_prefill_api is not None: + self._b12x_prefill_plan = self._b12x_kda_prefill_declaration( + self.kv_cache[1].shape[0] + ) + + def unbind_kv_cache(self) -> None: + # Only bindings to this state-pool generation become invalid. + self._b12x_prefill_plan = None + super().unbind_kv_cache() + + def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: + if self.kda_prefill_backend != "b12x": + return + api = get_b12x_kda_prefill() + if api is None: + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package." + ) + device = torch.device(current_platform.current_device()) + scheduler_config = vllm_config.scheduler_config + self._b12x_prefill_api = api + self._b12x_prefill_max_tokens = int(scheduler_config.max_num_batched_tokens) + self._b12x_prefill_max_seqs = int(scheduler_config.max_num_seqs) + max_seqs = self._b12x_prefill_max_seqs + # Device-side metadata the op reads by address; one set per layer. + for name, fill in ( + ("_b12x_prefill_cu_seqlens", None), + ("_b12x_prefill_initial_indices", None), + ("_b12x_prefill_null_indices", NULL_BLOCK_ID), + ("_b12x_prefill_zero_offsets", None), + ): + size = max_seqs + 1 if name == "_b12x_prefill_cu_seqlens" else max_seqs + buffer = torch.zeros(size, dtype=torch.int32, device=device) + if fill is not None: + buffer.fill_(fill) + self.register_buffer(name, buffer, persistent=False) + for name in ("_b12x_prefill_num_seqs", "_b12x_prefill_num_tokens"): + self.register_buffer( + name, + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + + def _b12x_kda_prefill_declaration(self, max_state_slots: int): + api = self._b12x_prefill_api + if api is None: + raise RuntimeError("b12x KDA prefill was not initialized") + return api.plan( + api.Caps( + device=current_platform.current_device(), + max_tokens=self._b12x_prefill_max_tokens, + max_seqs=self._b12x_prefill_max_seqs, + max_state_slots=max_state_slots, + heads=self.local_num_heads, + head_dim=self.head_dim, + model_dtype=self.model_config.dtype, + state_dtype=self.get_state_dtype()[1], + qk_l2norm=True, + # Prefill chunks end on recurrent block boundaries, as with + # Triton; no mid-chunk checkpoint slot is written. + checkpoint_export=False, + null_state_index=NULL_BLOCK_ID, + ), + invocation=api.invocation_from_tensors( + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state_indices=self._b12x_prefill_initial_indices, + ), + ) + + def _b12x_prefill_scratch_spec(self) -> tuple[tuple[int, ...], torch.dtype]: + plan = self._b12x_prefill_plan + if plan is None: + raise PreparationResourceUnavailableError( + "b12x KDA prefill is not prepared for the current KV generation" + ) + (spec,) = plan.scratch_specs() + return spec.shape, spec.dtype + + def get_b12x_preparation_units( + self, + layer: torch.nn.Module, + workload: B12xWorkload, + ) -> tuple[B12xPreparationUnit, ...]: + if layer is not self: + raise ValueError("KDA preparation owner mismatch") + if self._b12x_prefill_plan is None: + return () + request = self._b12x_prefill_plan.request( + name=f"{self._b12x_preparation_prefix}.kda.prefill", + prepare_call=self._prepare_b12x_kda_prefill, + benchmark_call=self._benchmark_b12x_kda_prefill, + ) + return ( + B12xPreparationUnit( + name="KDA prefill", + key=(self._b12x_preparation_prefix, "kda-prefill"), + requests=(request,), + stage="state", + ), + ) + + @staticmethod + def _benchmark_values(tensor: torch.Tensor) -> None: + values = torch.arange( + tensor.numel(), dtype=tensor.dtype, device=tensor.device + ).reshape_as(tensor) + tensor.copy_(values.div_(max(values.numel(), 1))) + + @staticmethod + def _b12x_trial_tensor(source: torch.Tensor) -> torch.Tensor: + """Allocate a trial tensor with the source's layout and alignment.""" + storage_size = 1 + sum( + (size - 1) * stride + for size, stride in zip(source.shape, source.stride()) + if size + ) + source_pointer = source.data_ptr() + source_alignment = ( + min(16, source_pointer & -source_pointer) if source_pointer else 16 + ) + storage = torch.empty( + storage_size + 16, dtype=source.dtype, device=source.device + ) + for storage_offset in range(16): + trial = storage.as_strided( + source.shape, source.stride(), storage_offset=storage_offset + ) + pointer = trial.data_ptr() + alignment = min(16, pointer & -pointer) if pointer else 16 + if alignment == source_alignment: + return trial + return storage.as_strided(source.shape, source.stride()) + + def _prepare_b12x_kda_prefill(self, state): + return self._b12x_kda_prefill_call(state, benchmark=False) + + def _benchmark_b12x_kda_prefill(self, state): + return self._b12x_kda_prefill_call(state, benchmark=True) + + def _b12x_kda_prefill_call(self, state, *, benchmark: bool): + from b12x.preparation import PreparedCall + + specs = tuple(state.layout.scratch_specs()) + if len(specs) != 1: + raise RuntimeError("b12x KDA prefill requires exactly one scratch buffer") + slots = self.kv_cache[1] + slot = 1 if slots.shape[0] > 1 else 0 + # Preparation owns its scratch and synthetic activations; serving + # binds the live request tensors and workspace-manager scratch. + scratch = torch.empty(specs[0].shape, dtype=specs[0].dtype, device=slots.device) + shape = (self._b12x_prefill_max_tokens, self.local_num_heads, self.head_dim) + q, k, v, raw_g, output = ( + torch.empty(shape, dtype=self.model_config.dtype, device=slots.device) + for _ in range(5) + ) + raw_beta = torch.empty( + shape[:2], dtype=self.model_config.dtype, device=slots.device + ) + if benchmark: + cu_seqlens = self._b12x_trial_tensor(self._b12x_prefill_cu_seqlens) + indices = self._b12x_trial_tensor(self._b12x_prefill_initial_indices) + checkpoint_indices = self._b12x_trial_tensor( + self._b12x_prefill_null_indices + ) + offsets = self._b12x_trial_tensor(self._b12x_prefill_zero_offsets) + num_seqs = self._b12x_trial_tensor(self._b12x_prefill_num_seqs) + num_tokens = self._b12x_trial_tensor(self._b12x_prefill_num_tokens) + saved_state = slots[slot : slot + 1].clone() + else: + cu_seqlens = self._b12x_prefill_cu_seqlens + indices = self._b12x_prefill_initial_indices + checkpoint_indices = self._b12x_prefill_null_indices + offsets = self._b12x_prefill_zero_offsets + num_seqs = self._b12x_prefill_num_seqs + num_tokens = self._b12x_prefill_num_tokens + saved_state = None + + def produce(): + for tensor in (q, k, v, raw_g, raw_beta): + self._benchmark_values(tensor) + cu_seqlens.zero_() + cu_seqlens[1:].fill_(self._b12x_prefill_max_tokens) + indices.fill_(slot) + if benchmark: + checkpoint_indices.fill_(NULL_BLOCK_ID) + offsets.zero_() + num_seqs.fill_(1) + num_tokens.fill_(self._b12x_prefill_max_tokens) + + def reset(): + if saved_state is not None: + slots[slot : slot + 1].copy_(saved_state) + + binding = state.bind( + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias.view(-1, self.head_dim), + recurrent_state=slots, + cu_seqlens=cu_seqlens, + initial_state_indices=indices, + final_state_indices=indices, + checkpoint_state_indices=checkpoint_indices, + checkpoint_offsets=offsets, + num_seqs=num_seqs, + num_tokens=num_tokens, + output=output, + ) + return PreparedCall( + run=lambda: state.run( + binding, + lower_bound=self.gate_lower_bound, + max_live_tokens=self._b12x_prefill_max_tokens, + max_live_seqs=1, + ), + produce=produce, + reset=reset, + restore=reset if saved_state is not None else None, + ) + + def _run_b12x_kda_prefill( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + cu_seqlens: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + recurrent_state: torch.Tensor, + output: torch.Tensor | None, + ) -> torch.Tensor: + """Advance the recurrent-state pool in place and return the output. + + Without a caller ``output``, the rows land in workspace storage + reserved together with the scratch. + """ + api = self._b12x_prefill_api + plan = self._b12x_prefill_plan + if api is None or plan is None: + raise PreparationResourceUnavailableError( + "b12x KDA prefill is not prepared for the current KV generation" + ) + num_tokens, num_requests = int(q.shape[0]), int(state_indices.shape[0]) + if ( + num_tokens > self._b12x_prefill_max_tokens + or num_requests > self._b12x_prefill_max_seqs + ): + raise ValueError( + "b12x KDA prefill capacity exceeded: " + f"tokens={num_tokens}/{self._b12x_prefill_max_tokens}, " + f"requests={num_requests}/{self._b12x_prefill_max_seqs}" + ) + manager = current_workspace_manager() + if output is None: + scratch, output = manager.get_simultaneous( + self._b12x_prefill_scratch_spec(), + ( + ( + self._b12x_prefill_max_tokens, + self.local_num_heads, + self.head_dim, + ), + self.model_config.dtype, + ), + ) + output = output[:num_tokens] + else: + (scratch,) = manager.get_simultaneous(self._b12x_prefill_scratch_spec()) + # A request without a computed prefix starts from the zero state. + initial_indices = self._b12x_prefill_initial_indices[:num_requests] + initial_indices.copy_(state_indices) + initial_indices.masked_fill_(~has_initial_state[:num_requests], NULL_BLOCK_ID) + self._b12x_prefill_num_seqs.fill_(num_requests) + self._b12x_prefill_num_tokens.fill_(num_tokens) + binding = api.bind( + plan, + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias.view(-1, self.head_dim), + recurrent_state=recurrent_state, + cu_seqlens=cu_seqlens[: num_requests + 1], + initial_state_indices=initial_indices, + final_state_indices=state_indices, + checkpoint_state_indices=self._b12x_prefill_null_indices[:num_requests], + checkpoint_offsets=self._b12x_prefill_zero_offsets[:num_requests], + num_seqs=self._b12x_prefill_num_seqs, + num_tokens=self._b12x_prefill_num_tokens, + output=output, + ) + api.run( + binding, + lower_bound=self.gate_lower_bound, + max_live_tokens=num_tokens, + max_live_seqs=num_requests, + ) + return output + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: spec = super().get_kv_cache_spec(vllm_config) assert isinstance(spec, MambaSpec) @@ -1297,12 +1660,38 @@ def _prefill_conv( assert non_spec_state_indices_tensor is not None assert has_initial_state is not None - initial_state = gather_initial_states( - recurrent_state, - non_spec_state_indices_tensor, - has_initial_state, + # The b12x op reads and writes the state pool in place. + initial_state = ( + None + if self.kda_prefill_backend == "b12x" + else gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) ) - if self.kda_prefill_backend == "flashkda": + if self.kda_prefill_backend == "b12x": + assert self.gate_lower_bound is not None + assert non_spec_query_start_loc is not None + b12x_out = None + if non_spec_out is not None: + b12x_out = non_spec_out[0] + elif not has_spec_decode: + b12x_out = core_attn_out[0, : q_ns.shape[1]] + core_attn_out_non_spec = self._run_b12x_kda_prefill( + q=q_ns[0], + k=k_ns[0], + v=v_ns[0], + raw_g=g1_ns[0], + raw_beta=beta_ns[0], + cu_seqlens=non_spec_query_start_loc, + state_indices=non_spec_state_indices_tensor, + has_initial_state=has_initial_state, + recurrent_state=recurrent_state, + output=b12x_out, + ).unsqueeze(0) + elif self.kda_prefill_backend == "flashkda": + assert initial_state is not None assert self.gate_lower_bound is not None assert self._flashkda_buffer_specs is not None workspace_out, final_state, checkpoint_state, workspace = ( @@ -1395,6 +1784,7 @@ def _prefill_conv( workspace=workspace, ) elif self.kda_prefill_backend == "flashinfer": + assert initial_state is not None assert self.gate_lower_bound is not None assert m.flashinfer_prefill_query_start_loc is not None if q_ns.shape[1] > initial_state.shape[0]: @@ -1444,9 +1834,10 @@ def _prefill_conv( cu_seqlens=non_spec_query_start_loc, out=non_spec_out, ) - recurrent_state[non_spec_state_indices_tensor] = ( - last_recurrent_state.to(recurrent_state.dtype) - ) + if self.kda_prefill_backend != "b12x": + recurrent_state[non_spec_state_indices_tensor] = ( + last_recurrent_state.to(recurrent_state.dtype) + ) else: # Pure non-speculative decode. assert non_spec_state_indices_tensor is not None @@ -1490,7 +1881,7 @@ def _prefill_conv( ) elif core_attn_out_non_spec is not None: if ( - self.kda_prefill_backend not in ("flashkda", "flashinfer") + self.kda_prefill_backend not in ("flashkda", "flashinfer", "b12x") or m.num_prefills == 0 ): # TODO: decode kernels write directly to core_attn_out From 185c10804080f760b07ba4319d02245b7676ff88 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 00:46:43 +0000 Subject: [PATCH 306/365] Preserve selectable FP32 dense MLA partials in canonical serving Retain BF16 defaults and qualify both partial dtypes through capacity and graph tests. GPU qualification pending. Co-authored-by: Codex --- .../changes/vllm-dense-mla-partial-precision.json | 15 +++++++++++++++ tests/v1/attention/test_b12x.py | 4 +++- vllm/envs.py | 5 +++++ vllm/v1/attention/backends/mla/b12x_mla.py | 6 ++++++ 4 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 .lil/changes/vllm-dense-mla-partial-precision.json diff --git a/.lil/changes/vllm-dense-mla-partial-precision.json b/.lil/changes/vllm-dense-mla-partial-precision.json new file mode 100644 index 000000000000..e2dd270a6214 --- /dev/null +++ b/.lil/changes/vllm-dense-mla-partial-precision.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-dense-mla-partial-precision", + "category": "feature", + "summary": "Select BF16 or FP32 split partials for B12X dense MLA", + "models": ["Kimi-K3"], + "compatibility": "BF16 MLA partials remain the default. VLLM_K3_DENSE_MLA_PARTIAL_DTYPE=fp32 requires the B12X FP32-partials companion.", + "details": [ + "Passes the declared partial dtype to prepared B12X MLA plans.", + "Covers both partial dtypes in CUDA Graph replay with capacity changes and page addresses beyond the signed 32-bit offset boundary." + ], + "pull_requests": [], + "authors": ["voipmonitor"], + "requires": ["b12x-dense-mla-fp32-partials"] +} diff --git a/tests/v1/attention/test_b12x.py b/tests/v1/attention/test_b12x.py index cc7e105c8429..1a581903537a 100644 --- a/tests/v1/attention/test_b12x.py +++ b/tests/v1/attention/test_b12x.py @@ -41,12 +41,14 @@ ) @pytest.mark.parametrize("parallel_drafting", [False, True]) @pytest.mark.parametrize("page_padding", [0, 192]) +@pytest.mark.parametrize("partial_dtype", ["bf16", "fp32"]) @torch.inference_mode() def test_b12x_dense_mla_prepared_capacity_replay_and_high_pages( - dtype, dcp_size, parallel_drafting, page_padding, monkeypatch + dtype, dcp_size, parallel_drafting, page_padding, partial_dtype, monkeypatch ): """Real prepared kernels consume high page IDs and mutable graph inputs.""" _require_b12x_paged_attention() + monkeypatch.setenv("VLLM_K3_DENSE_MLA_PARTIAL_DTYPE", partial_dtype) from b12x._lib.runtime_control import kernel_resolution_guard from b12x.attention import dense_mla from b12x.preparation import PreparationSession diff --git a/vllm/envs.py b/vllm/envs.py index 5a50d57da8fb..76fbe397aa49 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -172,6 +172,7 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_MLA_DISABLE: bool = False + VLLM_K3_DENSE_MLA_PARTIAL_DTYPE: Literal["bf16", "fp32"] = "bf16" VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" VLLM_CUDART_SO_PATH: str | None = None @@ -1499,6 +1500,10 @@ def _resolve_rust_cli_path() -> str | None: ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), + # Keep dense MLA split partials in FP32 until the merge when requested. + "VLLM_K3_DENSE_MLA_PARTIAL_DTYPE": env_with_choices( + "VLLM_K3_DENSE_MLA_PARTIAL_DTYPE", "bf16", ["bf16", "fp32"] + ), # If set, vLLM will pick up the provided Flash Attention MLA # Number of GPUs per worker in Ray, if it is set to be a fraction, # it allows ray to schedule multiple actors on a single GPU, diff --git a/vllm/v1/attention/backends/mla/b12x_mla.py b/vllm/v1/attention/backends/mla/b12x_mla.py index e9929439fdfd..0b8bb159bb14 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla.py +++ b/vllm/v1/attention/backends/mla/b12x_mla.py @@ -7,6 +7,7 @@ import torch +from vllm import envs from vllm.config import get_current_vllm_config from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group @@ -342,6 +343,11 @@ def get_b12x_preparation_units(self, layer, workload: B12xWorkload): v_head_dim=self.kv_lora_rank, physical_record_width=cache.shape[2], use_cuda_graph=True, + partial_dtype=( + torch.float32 + if envs.VLLM_K3_DENSE_MLA_PARTIAL_DTYPE == "fp32" + else torch.bfloat16 + ), ) def descriptor(shape, strides, dtype): From 66432ae736ececf9fc94c2d0f98a0c3394934f21 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 01:48:45 +0000 Subject: [PATCH 307/365] Describe bounded checkpoint loading in container release notes --- .lil/changes/vllm-841.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-841.json diff --git a/.lil/changes/vllm-841.json b/.lil/changes/vllm-841.json new file mode 100644 index 000000000000..eda6d1d13557 --- /dev/null +++ b/.lil/changes/vllm-841.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-841", + "category": "fix", + "summary": "Bound InstantTensor staging and retain deferred checkpoint tensor ownership", + "models": ["Kimi-K3"], + "compatibility": "No user action required.", + "details": [ + "Limits temporary checkpoint loading memory and preserves tensors retained by deferred layer processing.", + "CPU staging does not place inference weights on the CPU." + ], + "pull_requests": [841], + "requires": [] +} From 47b5694266e40be8ac3d87837b9b98b023bbf5bf Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 01:48:45 +0000 Subject: [PATCH 308/365] Describe native wheel reuse and autotune cache recovery --- .lil/changes/vllm-846.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-846.json diff --git a/.lil/changes/vllm-846.json b/.lil/changes/vllm-846.json new file mode 100644 index 000000000000..daa1459e127d --- /dev/null +++ b/.lil/changes/vllm-846.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-846", + "category": "performance", + "summary": "Reuse verified native wheel payloads and recover incomplete FlashInfer autotune caches", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "Python-only wheel rebuilds reuse native binaries only when native source, dependency and ABI signatures agree.", + "Autotune cache writes are atomic; an incomplete cache is discarded and regenerated." + ], + "pull_requests": [846], + "requires": [] +} From ac33bf51e4c455e7fac639ff908c2b50c19cdeb5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 01:53:36 +0000 Subject: [PATCH 309/365] fix(release): drop Python staged by other branches from the shared wheel build cache The native build cache (/src/vllm/build) is keyed by the dependency recipe, which canonical and beta share. setuptools stages Python modules under build/lib.* and bdist.*, and bdist_wheel packages everything staged there, so each channel's wheel shipped modules that exist only on the other branch: 8 canonical-only modules in beta 39ac37fa42 and 7 beta-only modules in canonical 502d6cb5ac. Remove the staging directories before building; compiled objects stay cached under build/temp.*. Co-Authored-By: Claude Opus 5.5 --- .../jovian_wheel_release/build_vllm_wheel.sh | 8 ++++++ .../test_source_packaging.py | 26 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/tools/jovian_wheel_release/build_vllm_wheel.sh b/tools/jovian_wheel_release/build_vllm_wheel.sh index 4a49bf3d0e78..782c51ba7943 100755 --- a/tools/jovian_wheel_release/build_vllm_wheel.sh +++ b/tools/jovian_wheel_release/build_vllm_wheel.sh @@ -38,6 +38,14 @@ if ! test -f "${dense_source}/kernel_selector.h" \ done < <(find "${source_root}/build" -type f -name CMakeCache.txt -print0) fi +# Every branch with the same dependency recipe shares the native build cache. +# setuptools stages Python modules under build/lib.* and bdist.*, and the wheel +# packages everything staged there, so a module that exists only on another +# branch would ship in this wheel. Compiled objects stay cached in build/temp.*; +# CMake installs them into the staging directory again. +find "${source_root}/build" -mindepth 1 -maxdepth 1 \ + \( -name 'lib.*' -o -name 'bdist.*' \) -exec rm -rf -- {} + + env -u PYTHONPATH \ VLLM_TARGET_DEVICE=cuda \ CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120a -DFETCHCONTENT_BASE_DIR=/tmp/vllm-fetchcontent -DVLLM_BUILD_CUTLASS_SCALED_MM_C2X=${VLLM_BUILD_CUTLASS_SCALED_MM_C2X:?}" \ diff --git a/tools/jovian_wheel_release/test_source_packaging.py b/tools/jovian_wheel_release/test_source_packaging.py index 1a07a3f05b76..fb5a3cc5fdf2 100644 --- a/tools/jovian_wheel_release/test_source_packaging.py +++ b/tools/jovian_wheel_release/test_source_packaging.py @@ -61,3 +61,29 @@ def git(*args): git("add", ".") git("commit", "--quiet", "-m", "Patched dependency fixture") assert git("rev-parse", "HEAD:cmake/external_projects") != identity + + +def test_wheel_build_drops_python_staged_by_other_branches(tmp_path): + """The shared native cache keeps compiled objects, not staged modules.""" + script = ( + Path(__file__).resolve().parents[2] + / "tools/jovian_wheel_release/build_vllm_wheel.sh" + ).read_text() + start = script.index('find "${source_root}/build" -mindepth 1 -maxdepth 1') + cleanup = script[start : script.index("\n\n", start)] + assert script.index(cleanup) < script.index("uv build") + build = tmp_path / "build" + for staged in ( + "lib.linux-x86_64-cpython-312/vllm/other_branch.py", + "bdist.linux-x86_64/wheel/vllm/other_branch.py", + "temp.linux-x86_64-cpython-312/_C.o", + ): + (build / staged).parent.mkdir(parents=True, exist_ok=True) + (build / staged).write_text("staged by an earlier build\n") + subprocess.run( + ["bash", "-euo", "pipefail", "-c", cleanup], + env={**os.environ, "source_root": str(tmp_path)}, + check=True, + ) + assert [path.name for path in build.iterdir()] == ["temp.linux-x86_64-cpython-312"] + assert (build / "temp.linux-x86_64-cpython-312/_C.o").is_file() From 7f82f89d65a7925d0de00d405359b21faba7e473 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 01:54:00 +0000 Subject: [PATCH 310/365] docs(release): describe the wheel staging fix Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-939.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-939.json diff --git a/.lil/changes/vllm-939.json b/.lil/changes/vllm-939.json new file mode 100644 index 000000000000..1c874e6ef783 --- /dev/null +++ b/.lil/changes/vllm-939.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-939", + "category": "fix", + "summary": "The vLLM wheel contains only the Python modules of its own source commit", + "models": ["all"], + "compatibility": "No user action required. Nothing imported the stray modules, so serving is unchanged.", + "details": [ + "Beta and canonical wheels share one native build cache, and Python modules that a build of the other branch had staged there were packaged too: beta 39ac37fa42 carried 8 canonical-only modules and canonical 502d6cb5ac 7 beta-only ones.", + "The wheel build now removes the staged Python before building. Compiled objects stay cached, so build time is unchanged." + ], + "pull_requests": [939], + "requires": [] +} From 6842dfcab3b060ebc4314cd6717f91bf30073f17 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 05:16:29 +0000 Subject: [PATCH 311/365] fix(b12x): release a prepared plan batch with one release_many call B12X master (15c7f6be) reclaims unused programs after every release call: a device synchronize, a sweep of all program caches and a full gc.collect(). B12xPreparedBatch released its plans one release() at a time, so profiling teardown repeated that for every plan. With b12x b4b12bcf, DeepSeek-V4.1 TP4 and MiMo-V2.6 TP2 did not finish starting within 60 minutes; both workers sat in _reclaim_programs -> gc.collect() under _release_b12x_profile_state. Release the batch with one session.release_many() when the session has it; older B12X keeps the per-plan path. Co-Authored-By: Claude Opus 5.5 --- tests/model_executor/test_b12x_warmup.py | 13 +++++++++++++ vllm/model_executor/warmup/b12x_prepare.py | 8 ++++++++ 2 files changed, 21 insertions(+) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index e28a574c72db..e86c4a05f053 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -345,6 +345,19 @@ def release(plan): assert released == ["second", "first", "pool"] +def test_prepared_batch_releases_all_plans_in_one_session_call() -> None: + """Program reclamation runs once per batch, not once per released plan.""" + calls = [] + session = SimpleNamespace( + release=lambda plan: calls.append(("release", plan)), + release_many=lambda plans: calls.append(("release_many", tuple(plans))), + ) + batch = b12x_prepare.B12xPreparedBatch(session, ("first", "second")) + batch.release() + batch.release() + assert calls == [("release_many", ("first", "second"))] + + def test_collect_units_filters_by_stage_and_tunes_eager_shapes() -> None: record: list[tuple] = [] weights, state, tower = torch.nn.Module(), torch.nn.Module(), torch.nn.Module() diff --git a/vllm/model_executor/warmup/b12x_prepare.py b/vllm/model_executor/warmup/b12x_prepare.py index 6bcf4c99f0cd..f5b97f91350c 100644 --- a/vllm/model_executor/warmup/b12x_prepare.py +++ b/vllm/model_executor/warmup/b12x_prepare.py @@ -599,6 +599,14 @@ def __init__(self, session, plans): def release(self) -> None: plans, self.plans = self.plans, () + if not plans: + return + # B12X reclaims unused programs and collects garbage after each release + # call; per-plan releases repeat that for every plan of the batch. + release_many = getattr(self.session, "release_many", None) + if release_many is not None: + release_many(plans) + return with ExitStack() as stack: for plan in plans: stack.callback(self.session.release, plan) From 7182eac522129c2f63114618e6d0c1b2c306802c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 05:24:14 +0000 Subject: [PATCH 312/365] docs(release): describe the batched B12X profiling release Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-941.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-941.json diff --git a/.lil/changes/vllm-941.json b/.lil/changes/vllm-941.json new file mode 100644 index 000000000000..b7755a6535d0 --- /dev/null +++ b/.lil/changes/vllm-941.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-941", + "category": "fix", + "summary": "Startup releases its B12X profiling plans in one batch, so program reclamation and garbage collection run once instead of once per plan", + "models": ["all"], + "compatibility": "No user action required. With a B12X release that has no batched release, plans are released one at a time as before.", + "details": [ + "B12X now reclaims unused programs and runs a full garbage collection after every release call. vLLM released each profiling plan separately, so a start with thousands of plans spent that time once per plan: DeepSeek-V4.1 TP4 and MiMo-V2.6 TP2 did not finish starting within an hour, and GLM-5.3 and Qwen3.8 started minutes later.", + "The prepared batch now goes to B12X in one release_many call. DeepSeek-V4.1 TP4 starts in about 6 minutes again, with the same KV cache size and decode speed." + ], + "pull_requests": [941], + "requires": [] +} From 39bdccc5271f1259bdaf2039555204819982e45e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 12:10:52 +0000 Subject: [PATCH 313/365] docs(release): the full #926 port supersedes the vllm-933 note Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-943.json | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .lil/changes/vllm-943.json diff --git a/.lil/changes/vllm-943.json b/.lil/changes/vllm-943.json new file mode 100644 index 000000000000..3f220725912f --- /dev/null +++ b/.lil/changes/vllm-943.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-943", + "category": "fix", + "summary": "DeepSeek-V4.1 decode no longer drifts from the model as generation lengthens: the complete compressor-ring fix from vllm #926", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "No user action required. Supersedes the vllm-933 note: that measurement covered only the first half of #926, which cannot run on its own.", + "details": [ + "Circular-buffer KV groups now get a worker slot mapping, so DeepSeek-V4.1 writes its compressor state ring for every real decode row.", + "8 prompts x 2,048 greedy tokens on RTX PRO 6000: KL(prefill||decode) at decode steps 1536-2048 fell from 0.156 to 0.018 and top-1 agreement rose from 0.92 to 0.97. Divergence now stays flat with length. Decode speed, DSpark acceptance and KV size are unchanged." + ], + "authors": ["MadeBy561"], + "pull_requests": [943, 926], + "requires": [] +} From cd03c96de79a20332e3f7f66dcafd8b5ca896731 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 12:30:03 +0000 Subject: [PATCH 314/365] docs(release): state who the DeepSeek-V4.1 decode drift affected and what it did vllm-943 is not published yet. Name the affected hardware (SM12x) and images, what was wrong in plain terms, and the measured consequence, so the release note answers what the fix changes for a user. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-943.json | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.lil/changes/vllm-943.json b/.lil/changes/vllm-943.json index 3f220725912f..8ca90b5cd1ad 100644 --- a/.lil/changes/vllm-943.json +++ b/.lil/changes/vllm-943.json @@ -2,12 +2,13 @@ "schema": "local-inference-release-change/v1", "id": "vllm-943", "category": "fix", - "summary": "DeepSeek-V4.1 decode no longer drifts from the model as generation lengthens: the complete compressor-ring fix from vllm #926", + "summary": "DeepSeek-V4.1 on RTX PRO 6000, RTX 5090 and DGX Spark no longer drifts away from the model as a generation gets longer", "models": ["DeepSeek-V4.1-Flash"], - "compatibility": "No user action required. Supersedes the vllm-933 note: that measurement covered only the first half of #926, which cannot run on its own.", + "compatibility": "No user action required. Affected every earlier DeepSeek-V4.1 image, beta and canonical, on SM12x GPUs (RTX PRO 6000, RTX 5090, DGX Spark). DeepSeek-V4.1 on other GPUs runs a different implementation and was not affected, and no other model uses this path. Supersedes the vllm-933 note, which measured only the first half of vllm #926.", "details": [ - "Circular-buffer KV groups now get a worker slot mapping, so DeepSeek-V4.1 writes its compressor state ring for every real decode row.", - "8 prompts x 2,048 greedy tokens on RTX PRO 6000: KL(prefill||decode) at decode steps 1536-2048 fell from 0.156 to 0.018 and top-1 agreement rose from 0.92 to 0.97. Divergence now stays flat with length. Decode speed, DSpark acceptance and KV size are unchanged." + "What was wrong: DeepSeek-V4.1 packs groups of tokens into compressed KV entries that its sparse attention reads across the context. During decode, the compressor's partial state for newly generated tokens was never stored, because circular-buffer KV groups got no worker slot mapping. Compressed entries covering generated text were therefore built from stale state. Prompt processing (prefill) was correct.", + "Consequence: the longer the answer, the more generation diverged from what the model computes for the same text. Over 8 greedy 2,048-token generations, the divergence of the next-token distribution from a prefill of the same text (KL) grew from 0.038 over the first 256 tokens to 0.156 over tokens 1,536-2,048, where the top-1 token differed at 7.6% of positions. Longer outputs, such as long reasoning, were likely affected more; we did not measure beyond 2,048 tokens.", + "With the fix the divergence is 0.018-0.030 at every length, and the top-1 token differs at 3-5% of positions, the residual numeric difference between decode and prefill. Decode speed, DSpark acceptance and KV cache size are unchanged." ], "authors": ["MadeBy561"], "pull_requests": [943, 926], From 0f2584c158449f1ff4bfada83dd1ea50d7a10b9e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 13:00:16 +0000 Subject: [PATCH 315/365] docs(release): long prompts, multi-turn chats and images under the DeepSeek-V4.1 ring fix Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-943-long-context.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-943-long-context.json diff --git a/.lil/changes/vllm-943-long-context.json b/.lil/changes/vllm-943-long-context.json new file mode 100644 index 000000000000..92355e051e86 --- /dev/null +++ b/.lil/changes/vllm-943-long-context.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-943-long-context", + "category": "fix", + "summary": "More on the DeepSeek-V4.1 compressor-ring fix (vllm-943): long prompts, multi-turn chats and images", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "No user action required. Corrects vllm-943's statement that prompt processing was correct: long prompts were affected slightly at chunk boundaries.", + "details": [ + "The ring was not written during prompt processing either. A long prompt is processed in chunks (4,096 tokens per step in the ds41-flash profile), and only the compression windows that span a chunk boundary were built from stale state. That is a small part of a long prompt; it was not measured separately.", + "Generated text feeds compressed entries once it leaves the 128-token sliding window, which is why the drift grew with answer length. In multi-turn chats and agent sessions that reuse earlier answers from the prefix cache, those stale entries stayed in the context of every later turn. Long sessions where the model wrote much of the context were therefore likely affected the most; this was not measured.", + "Images are processed with the prompt, so they were affected like prompt text. Answers about them drifted like any other output." + ], + "authors": ["MadeBy561"], + "pull_requests": [943, 926], + "requires": [] +} From 0b1bccb858f7ccaf6cea039623f28e0579fca21d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 19:45:54 +0000 Subject: [PATCH 316/365] fix(scheduler): let decodes preempt in a prefill turn that scheduled nothing With prefill compute sharing, any waiting request counts as runnable prefill. While those requests cannot be admitted (every running slot is taken, or restores are waiting for GPU capacity), each contended step is a prefill turn, and running decodes get only the leftover pass, which may not preempt. When a decode needed a KV block from a full pool, that pass stopped without scheduling or preempting anything. The empty step left the fairness controller unchanged, so every later step was the same. Generation stopped with idle GPUs and no log lines until clients aborted the running requests (mrweiner: GLM-5.3-Flash TP2 with LMCache restore admission, four slots, KV cache at 98%). A prefill turn that scheduled nothing has no admitted work to protect, so its decode pass may now preempt, as in a decode turn. Co-Authored-By: Claude Opus 5.5 --- .../changes/vllm-restore-admission-stall.json | 16 +++++ tests/v1/core/test_boundary_admission.py | 62 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 6 +- 3 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 .lil/changes/vllm-restore-admission-stall.json diff --git a/.lil/changes/vllm-restore-admission-stall.json b/.lil/changes/vllm-restore-admission-stall.json new file mode 100644 index 000000000000..70ad3f20ec92 --- /dev/null +++ b/.lil/changes/vllm-restore-admission-stall.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-restore-admission-stall", + "category": "fix", + "summary": "Generation no longer stops for good when the KV cache fills while requests are waiting, such as LMCache restores waiting for GPU capacity", + "models": ["GLM-5.3-Flash", "DeepSeek-V4.1-Flash", "MiMo-V2.6-Flash"], + "compatibility": "No user action required. Affects servers with prefill compute sharing (prefill-compute-share), which the GLM-5.3-Flash, DeepSeek-V4.1-Flash and MiMo-V2.6-Flash profiles enable. Restore admission (vllm-restore-admission) made it common, because restores that wait for capacity stay queued instead of being admitted to recompute.", + "details": [ + "What was wrong: with compute sharing, any waiting request counts as prefill work, so while the waiting requests could not be admitted (all running slots taken, or restores waiting for capacity) every model step was a prefill turn that admitted nothing. Running decodes were served only as leftover work, which may not preempt. When a decode needed a KV block and the cache was full, nothing was preempted, the step was empty, and every later step was the same.", + "Consequence: the engine kept stepping with idle GPUs, KV cache usage near 100%, prompt and generation at 0 tok/s, no preemptions and no log lines, until clients aborted the running requests. VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S could not help: with every slot taken the waiting restores were never considered, and a restore that gave up could not be admitted to recompute either.", + "Now a prefill turn that schedules no work lets decodes preempt as a decode turn does. Once a running slot frees, waiting restores are considered again, and a restore still refused after VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S recomputes its prompt and logs a warning, as before." + ], + "pull_requests": [946], + "authors": ["Local Inference Lab"], + "requires": [] +} diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index 7d73321e1bad..cb5ab27cb504 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -952,6 +952,68 @@ def poll(req): assert not cache.external_boundary_wait_expired(waiter.request_id) +def run_contended_steps(scheduler, steps): + """Schedule steps, charging each contended one to its compute class.""" + outputs = [] + for _ in range(steps): + output = scheduler.schedule() + if output.compute_service_class is not None: + scheduler.record_compute_time( + output.compute_service_class, 0.01, contended=True + ) + outputs.append(output) + return outputs + + +@pytest.mark.parametrize("slots", ["full", "free"]) +def test_decodes_on_a_full_pool_preempt_while_restores_wait(slots): + """A prefill turn that schedules nothing still lets decodes preempt. + + Waiting restores keep compute sharing contended, so every step after the + first decode quantum is a prefill turn, whether the restores wait behind + full running slots or for GPU capacity. When a decode then needs a block + from a full pool, that turn must preempt like a decode turn; otherwise + every later step is empty and nothing changes. + """ + from tests.v1.core.utils import create_requests + + scheduler, cache, connector = import_scheduler( + max_num_seqs=2 if slots == "full" else 3, + num_blocks=17, + **compute_share_fixture_options(0.4), + ) + first, second = create_requests( + num_requests=2, num_tokens=60, max_tokens=64, req_ids=["first", "second"] + ) + for req in (first, second): + scheduler.add_request(req) + assert scheduler.schedule().num_scheduled_tokens == {"first": 60, "second": 60} + assert cache.block_pool.get_num_free_blocks() == 0 + restores = [request(name, name) for name in ("restore-a", "restore-b")] + polled = [] + + def poll(req): + polled.append(req.request_id) + if req not in restores: + return True + assert reserve_import(cache, req, prefix=128) is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + scheduler.connector = connector + for req in restores: + scheduler.add_request(req) + # Four decode steps fill the last block; the fifth needs a new one. + outputs = run_contended_steps(scheduler, 8) + assert all(output.num_scheduled_tokens for output in outputs) + assert second.num_preemptions == 1 + assert first.status == RequestStatus.RUNNING + assert first.num_computed_tokens == 68 + # With a slot free the restores compete, and their wait is bounded, again. + assert {"restore-a", "restore-b"} <= set(polled) + assert cache._boundary_import_waiters.keys() == {"restore-a", "restore-b"} + + def test_reset_with_running_requests_releases_ready_unadmitted_import(): """A cache reset drops a finished restore; its request then recomputes.""" scheduler, cache, connector = import_scheduler() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 66bd89742b79..e9feff11a4a1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2033,11 +2033,13 @@ def schedule_running_requests( # A prefill turn gives prefills first use of the model-step capacity. # Existing decodes then consume only capacity genuinely left over. New # requests admitted above are excluded by running_req_ids_at_step_start, - # so no request can be scheduled twice in one step. + # so no request can be scheduled twice in one step. A turn that + # scheduled nothing has no admitted work to protect; its decodes preempt + # as in a decode turn, or a full pool would leave every step empty. if adaptive_prefill_turn and token_budget > 0: schedule_running_requests( "decode", - allow_preemption=False, + allow_preemption=not num_scheduled_tokens, enforce_lora_limit=True, ) From 1dc03737776e6a07f58b608c6fcd4ec7b0a01c5f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 23:14:36 +0000 Subject: [PATCH 317/365] feat(ds41): load lossless X4T experts through B12X Register exact_mxfp4 loading for compressed expert scales while retaining native dense and draft precision and file-backed Engram tables. Reuse serialized scale scratch and reject unsupported parallel execution modes. Preserve EXL3 prepared-weight validation and validate the SwiGLU clamp. Validation: 122 X4T, EXL3, and loader-registry tests pass with B12X #446. New code passes mypy; the five existing model.py findings reproduce on the KK base. Full-model serving evidence is linked in the companion B12X PR. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/quantization/test_ds41_exact_mxfp4.py | 217 +++++++++++++++++ vllm/config/load.py | 2 + vllm/config/model.py | 1 + vllm/model_executor/layers/fused_moe/b12x.py | 26 +- .../layers/quantization/__init__.py | 3 + vllm/model_executor/model_loader/__init__.py | 3 + .../model_loader/exact_mxfp4_loader.py | 59 +++++ vllm/models/deepseek_v4_1/exact_mxfp4.py | 229 ++++++++++++++++++ vllm/models/deepseek_v4_1/nvidia/model.py | 5 +- 9 files changed, 535 insertions(+), 10 deletions(-) create mode 100644 tests/quantization/test_ds41_exact_mxfp4.py create mode 100644 vllm/model_executor/model_loader/exact_mxfp4_loader.py create mode 100644 vllm/models/deepseek_v4_1/exact_mxfp4.py diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_ds41_exact_mxfp4.py new file mode 100644 index 000000000000..2de6bcd4daa3 --- /dev/null +++ b/tests/quantization/test_ds41_exact_mxfp4.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file + +from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader +from vllm.models.deepseek_v4_1.exact_mxfp4 import ( + DeepseekV41ExactMXFP4Config, + ExactMXFP4MoEMethod, +) + + +@pytest.fixture +def moe(): + return FusedMoEConfig( + num_experts=384, + num_local_experts=384, + num_logical_experts=384, + experts_per_token=6, + hidden_dim=5120, + intermediate_size=576, + in_dtype=torch.bfloat16, + device="cpu", + activation=MoEActivation.SILU, + swiglu_limit=10.0, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + ) + + +@pytest.fixture +def owner(): + return DeepseekV41ExactMXFP4Config.from_config( + {"format_version": 1, "checkpoint_root": "/x4t"} + ) + + +def test_exact_config_keeps_native_dense_precision_and_empty_expert_handles(moe, owner): + assert owner.is_scale_e8m0 and owner.weight_block_size == [32, 32] + assert owner.is_checkpoint_fp8_serialized and owner.scale_scratch is None + method = ExactMXFP4MoEMethod(moe, owner) + layer = torch.nn.Module() + layer.layer_name = "model.layers.17.ffn.experts" + method.create_weights(layer, 384, 5120, 576, torch.bfloat16) + assert not list(layer.named_parameters()) and not layer.state_dict() + assert method.layer_index == 17 and method.local_intermediate == 576 + assert method.get_fused_moe_quant_config(layer).weight_quant_dtype == "mxfp4" + with pytest.raises(ValueError, match="format_version"): + DeepseekV41ExactMXFP4Config.from_config({}) + + +def test_exact_config_and_loader_are_registered(): + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.model_executor.model_loader import get_model_loader + + assert get_quantization_config("exact_mxfp4") is DeepseekV41ExactMXFP4Config + assert isinstance( + get_model_loader(LoadConfig(load_format="exact_mxfp4")), ExactMXFP4ModelLoader + ) + + +@pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) +def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( + monkeypatch, moe, owner, mode +): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.models.deepseek_v4_1 import exact_mxfp4 + + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace( + pipeline_parallel_size=2 if mode == "pipeline" else 1, + use_ubatching=mode == "ubatching", + ), + load_config=SimpleNamespace( + load_format="safetensors" if mode == "wrong_loader" else "exact_mxfp4" + ), + ) + monkeypatch.setattr(exact_mxfp4, "get_current_vllm_config", lambda: config) + native = object() + monkeypatch.setattr( + exact_mxfp4.DeepseekV41FP8Config, + "get_quant_method", + lambda *args, **kwargs: native, + ) + layer = Mock(spec=RoutedExperts) + layer.moe_config = moe + if mode == "normal": + assert isinstance( + owner.get_quant_method(layer, "model.layers.39.ffn.experts"), + ExactMXFP4MoEMethod, + ) + elif mode == "wrong_loader": + with pytest.raises(ValueError, match="load-format"): + owner.get_quant_method(layer, "model.layers.39.ffn.experts") + else: + with pytest.raises(NotImplementedError, match="PP1 without ubatching"): + owner.get_quant_method(layer, "model.layers.39.ffn.experts") + assert owner.get_quant_method(layer, "model.layers.40.ffn.experts") is native + assert owner.get_quant_method(torch.nn.Module(), "model.layers.0.attn") is native + + +@pytest.mark.parametrize( + "field,value", [("dp_size", 2), ("use_ep", True), ("enable_eplb", True)] +) +def test_exact_experts_reject_non_tp_execution(moe, owner, field, value): + parallel = replace(moe.moe_parallel_config, **{field: value}) + with pytest.raises(NotImplementedError, match="TP without EP/DP"): + ExactMXFP4MoEMethod(replace(moe, moe_parallel_config=parallel), owner) + + +def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( + monkeypatch, moe, owner +): + from b12x.moe import fused_moe + + from vllm.model_executor.layers.fused_moe import b12x + + layer = torch.nn.Module() + layer.activation = moe.activation + layer.apply_router_weight_on_input = False + quant = ExactMXFP4MoEMethod(moe, owner).get_fused_moe_quant_config(layer) + backend = b12x.B12xExperts(moe, quant) + prepared = Mock(spec=fused_moe.PreparedExperts) + activation = fused_moe.ActivationSpec( + mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16, swiglu_limit=10.0 + ) + prepared.plan = SimpleNamespace( + source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), + activation=activation, + ) + prepared.num_experts = moe.num_experts + prepared.hidden_size = moe.hidden_dim + prepared.intermediate_size = moe.intermediate_size_per_partition + monkeypatch.setattr( + b12x, "_register_b12x_moe_output_collective", lambda *a, **k: None + ) + backend.install_prepared_experts(layer, prepared) + assert layer._b12x_prepared_experts is prepared + assert backend._prepared_experts is prepared + prepared.plan.activation = replace(activation, swiglu_limit=None) + with pytest.raises(ValueError, match="activation"): + backend.install_prepared_experts(layer, prepared) + prepared.plan.activation = activation + prepared.hidden_size += 128 + with pytest.raises(ValueError, match="geometry"): + backend.install_prepared_experts(layer, prepared) + prepared.hidden_size = moe.hidden_dim + prepared.plan.source = fused_moe.PackedSource(format="modelopt_nvfp4") + with pytest.raises(TypeError, match="encoding"): + backend.install_prepared_experts(layer, prepared) + + +def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): + from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + + tensor_dir = tmp_path / "tensors" + tensor_dir.mkdir() + name = "model-00001.safetensors" + tensors = { + "layers.0.ffn.experts.0.w1.weight": torch.ones(8, dtype=torch.uint8), + "layers.0.ffn.experts.0.w1.scale.exact_mxfp4_fixed": torch.ones( + 8, dtype=torch.uint8 + ), + "layers.40.ffn.experts.0.w1.weight": torch.full((8,), 42, dtype=torch.uint8), + "layers.1.engram.embedding.weight": torch.ones((8, 8), dtype=torch.uint8), + "norm.weight": torch.ones(8, dtype=torch.bfloat16), + } + save_file(tensors, tensor_dir / name) + common = {"schema": SCHEMA, "codec": CODEC, "family": "deepseek_v41"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": name}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: name for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace( + num_hidden_layers=40, + quantization_config={ + "quant_method": "exact_mxfp4", + "checkpoint_root": str(tmp_path), + }, + ) + ) + model = SimpleNamespace( + checkpoint_file_weight_filter=lambda n: ".engram.embedding." in n + ) + loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + result = dict(loader.get_all_weights(config, model)) + assert len(result) == 3 + assert result["layers.1.engram.embedding.weight"].device.type == "meta" + assert result[ + "layers.1.engram.embedding.weight" + ]._vllm_file_tensor_source.shape == (8, 8) + assert torch.equal( + result["layers.40.ffn.experts.0.w1.weight"], + tensors["layers.40.ffn.experts.0.w1.weight"], + ) + model.checkpoint_weight_name_prefixes = ("layers.40.",) + assert list(dict(loader.get_all_weights(config, model))) == [ + "layers.40.ffn.experts.0.w1.weight" + ] diff --git a/vllm/config/load.py b/vllm/config/load.py index 5e24a617ba85..6729f6820c45 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -35,6 +35,8 @@ class LoadConfig: back to the pytorch bin format if safetensors format is not available. - "pt" will load the weights in the pytorch bin format. - "safetensors" will load the weights in the safetensors format. + - "exact_mxfp4" reads losslessly compressed DS4.1 MXFP4 expert scales and + retains file-backed loading for Engram tables. - "instanttensor" will load the Safetensors weights on CUDA devices using InstantTensor, which enables distributed loading with pipelined prefetching and fast direct I/O. diff --git a/vllm/config/model.py b/vllm/config/model.py index 3cb01da228c3..30df9ca07aea 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1284,6 +1284,7 @@ def _verify_quantization(self) -> None: "mxfp8", "modelopt_mixed", "exl3", + "exact_mxfp4", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index af958dca03ba..7f9f129ebe4c 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -499,28 +499,36 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> None: - """Install EXL3 weights prepared through the common trellis API.""" + """Install canonical prepared experts without retaining source parameters.""" fused_moe = _require_b12x_fused_moe() - if ( - self._source_format != "exl3" - or not isinstance(prepared, fused_moe.PreparedExperts) - or not isinstance(prepared.plan.source, fused_moe.TrellisSource) - ): - raise TypeError("EXL3 installation requires B12X prepared trellis weights") + if not isinstance(prepared, fused_moe.PreparedExperts): + raise TypeError("Installation requires B12X PreparedExperts") + trellis = self._source_format == "exl3" and isinstance( + prepared.plan.source, fused_moe.TrellisSource + ) + packed = ( + self._source_format == "fp4_e8m0_k32" + and isinstance(prepared.plan.source, fused_moe.PackedSource) + and prepared.plan.source.format == "fp4_e8m0_k32" + ) + if not (trellis or packed): + raise TypeError("Prepared expert encoding does not match the backend") if ( prepared.num_experts != self.moe_config.num_experts or prepared.hidden_size != self.moe_config.hidden_dim or prepared.intermediate_size != self.moe_config.intermediate_size_per_partition or prepared.plan.activation.mode != "a16" - or prepared.plan.activation.rotation_dtype != torch.float16 + or (trellis and prepared.plan.activation.rotation_dtype != torch.float16) or prepared.plan.activation.io_dtype != self.moe_config.in_dtype or prepared.plan.activation.nonlinearity != _b12x_activation_name(layer.activation) + or prepared.plan.activation.swiglu_limit != self.moe_config.swiglu_limit or layer.apply_router_weight_on_input ): raise ValueError( - "Prepared EXL3 geometry, activation or routing does not match the layer" + "Prepared expert geometry, activation or routing " + "does not match the layer" ) self._apply_router_weight_on_input = False self._reuse_prepared_storage(layer, prepared) diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index c5e0f1446de1..95e79137d9d7 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -38,6 +38,7 @@ "gpt_oss_mxfp4", "deepseek_v4_fp8", "deepseek_v41_fp8", + "exact_mxfp4", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -132,6 +133,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: deepseek_config = DeepseekV4FP8Config + from vllm.models.deepseek_v4_1.exact_mxfp4 import DeepseekV41ExactMXFP4Config from vllm.models.deepseek_v4_1.quant_config import DeepseekV41FP8Config from .auto_awq import AutoAWQConfig @@ -182,6 +184,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "gpt_oss_mxfp4": GptOssMxfp4Config, "deepseek_v4_fp8": deepseek_config, "deepseek_v41_fp8": DeepseekV41FP8Config, + "exact_mxfp4": DeepseekV41ExactMXFP4Config, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index 6305c45f5d0b..d542fc862463 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -11,6 +11,7 @@ from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader from vllm.model_executor.model_loader.modelexpress_loader import ( ModelExpressModelLoader, ) @@ -36,6 +37,7 @@ "dummy", "fastsafetensors", "instanttensor", + "exact_mxfp4", "ipc_cache", "mistral", "modelexpress", @@ -53,6 +55,7 @@ "dummy": DummyModelLoader, "fastsafetensors": DefaultModelLoader, "instanttensor": DefaultModelLoader, + "exact_mxfp4": ExactMXFP4ModelLoader, "ipc_cache": IpcModelLoader, "mistral": DefaultModelLoader, "modelexpress": ModelExpressModelLoader, diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py new file mode 100644 index 000000000000..1a81c4b703f1 --- /dev/null +++ b/vllm/model_executor/model_loader/exact_mxfp4_loader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read retained native tensors from a lossless MXFP4 container. + +Compressed routed experts are owned by the model's quantization method. +Engram tensors retain their immutable file-range descriptors, avoiding full +CPU/GPU staging allocations for the host-mapped embedding tables. +""" + +import time +from pathlib import Path + +import regex as re +from safetensors import safe_open + +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.weight_utils import ( + file_source_tensor, + safetensors_file_sources, +) + + +class ExactMXFP4ModelLoader(DefaultModelLoader): + def _root(self, model_config): + from b12x.moe.checkpoints.exact_mxfp4 import checkpoint_contract + + quant = model_config.hf_config.quantization_config + if quant.get("quant_method") != "exact_mxfp4": + raise ValueError("exact_mxfp4 loading requires the X4T model config") + root = Path(quant["checkpoint_root"]) + if not root.is_absolute(): + raise ValueError("X4T checkpoint_root must be an absolute local path") + return root, checkpoint_contract(str(root.resolve())) + + def download_model(self, model_config): + self._root(model_config) + + def get_all_weights(self, model_config, model): + root, contract = self._root(model_config) + if getattr(model, "secondary_weights", ()): + raise NotImplementedError("X4T does not support secondary weight sources") + file_filter = getattr(model, "checkpoint_file_weight_filter", None) + prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) + layers = model_config.hf_config.num_hidden_layers + self.counter_before_loading_weights = time.perf_counter() + for filename in sorted(set(contract["source_names"].values())): + path = str(root / "tensors" / filename) + descriptors = safetensors_file_sources(path) + with safe_open(path, framework="pt", device="cpu") as handle: + for name in sorted(descriptors): + if prefixes is not None and not name.startswith(prefixes): + continue + match = re.match(r"layers\.(\d+)\.ffn\.experts\.", name) + if match and int(match.group(1)) < layers: + continue + if callable(file_filter) and file_filter(name): + yield name, file_source_tensor(descriptors[name]) + else: + yield name, handle.get_tensor(name) diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/exact_mxfp4.py new file mode 100644 index 000000000000..1e4536b4ae25 --- /dev/null +++ b/vllm/models/deepseek_v4_1/exact_mxfp4.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Losslessly compressed DS4.1 routed experts with native dense tensors.""" + +import regex as re +import torch + +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + FusedMoEQuantDesc, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +from .quant_config import DeepseekV41FP8Config + +logger = init_logger(__name__) + + +class DeepseekV41ExactMXFP4Config(DeepseekV41FP8Config): + """X4T main experts; unchanged native dense, shared and draft weights.""" + + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "exact_mxfp4" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if ( + user_quant in (None, cls.get_name()) + and hf_quant_cfg is not None + and hf_quant_cfg.get("quant_method") == cls.get_name() + and getattr(hf_config, "model_type", None) + in ("deepseek_v41", "deepseek_v41_text") + ): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("X4T requires format_version=1 and checkpoint_root") + result = cls( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[32, 32], + ) + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) + if match is None: + raise ValueError("DS4.1 routed experts require a numbered layer") + config = get_current_vllm_config() + if int(match.group(1)) < config.model_config.hf_config.num_hidden_layers: + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "X4T shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "exact_mxfp4": + raise ValueError("X4T requires --load-format exact_mxfp4") + return ExactMXFP4MoEMethod(layer.moe_config, self) + return super().get_quant_method(layer, prefix) + + +class ExactMXFP4MoEMethod(FusedMoEMethodBase): + """Expand scales for routed experts into serialized, model-owned scratch.""" + + def __init__(self, moe, owner): + super().__init__(moe) + self.owner = owner + parallel = moe.moe_parallel_config + if ( + parallel.use_ep + or parallel.ep_size != 1 + or parallel.dp_size != 1 + or parallel.use_all2all_kernels + or parallel.enable_eplb + ): + raise NotImplementedError("X4T DS4.1 supports TP without EP/DP") + if ( + moe.activation != MoEActivation.SILU + or moe.in_dtype != torch.bfloat16 + or moe.has_bias + ): + raise ValueError("X4T DS4.1 requires bias-free BF16 SwiGLU experts") + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) + if match is None or params_dtype != torch.bfloat16: + raise ValueError("X4T requires a numbered BF16 expert layer") + self.layer_index = int(match.group(1)) + self.num_experts, self.hidden_size = num_experts, hidden_size + self.local_intermediate = intermediate_size_per_partition + for name in ("w13_weight", "w2_weight"): + layer.register_buffer( + name, torch.empty(0, dtype=torch.uint8), persistent=False + ) + + def get_fused_moe_quant_config(self, layer): + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(dtype="mxfp4"), + _w2=FusedMoEQuantDesc(dtype="mxfp4"), + ) + + def process_weights_after_loading(self, layer): + from b12x.moe import fused_moe + from b12x.moe.checkpoints.exact_mxfp4 import read_exact_mxfp4_layer + + tp, rank = ( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + ) + device = layer.w13_weight.device + e, h, n = self.num_experts, self.hidden_size, self.local_intermediate + shapes = ((e, h // 32, 2 * n), (e, n // 32, h)) + if self.owner.scale_scratch is None: + # DS4.1 disallows ubatching: every layer consumes these scale grids + # before its successor may overwrite them on the same stream. + self.owner.scale_scratch = tuple( + torch.empty(s, dtype=torch.uint8, device=device) for s in shapes + ) + scratch = self.owner.scale_scratch + if any( + t.shape != shape or t.device != device for t, shape in zip(scratch, shapes) + ): + raise ValueError("X4T shared scale scratch geometry/device mismatch") + weights = read_exact_mxfp4_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) + plan = fused_moe.plan_weights( + source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), + activation=fused_moe.ActivationSpec( + mode="a16", + nonlinearity="silu", + io_dtype=torch.bfloat16, + swiglu_limit=self.moe.swiglu_limit, + ), + geometry=fused_moe.MoEGeometry( + num_experts=e, hidden_size=h, intermediate_size=n + ), + ) + prepared = fused_moe.prepare_weights(plan=plan, weights=weights) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + backend = B12xExperts(self.moe, self.moe_quant_config) + backend.install_prepared_experts(layer, prepared) + self.moe_kernel = mk.FusedMoEKernel( + MoEPrepareAndFinalizeNoDPEPModular(), backend + ) + logger.info( + "X4T lossless MXFP4 layer %d rank %d/%d: BF16 activations, " + "compressed scales, shared scratch %d bytes", + self.layer_index, + rank, + tp, + sum(t.numel() for t in scratch), + ) + + def apply( + self, + layer, + x, + topk_weights, + topk_ids, + shared_experts, + shared_experts_input, + workspace=None, + ): + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + workspace=workspace, + ) diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index f2323ae3c915..f4b283eec113 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -369,7 +369,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if vllm_config.lora_config is not None: raise ValueError("V4.1 native kernels do not support LoRA adapters") - if quant_config is None or quant_config.get_name() != "deepseek_v41_fp8": + if quant_config is None or quant_config.get_name() not in ( + "deepseek_v41_fp8", + "exact_mxfp4", + ): raise ValueError("V4.1 requires its native block32 quantization config") self.vocab_size = config.vocab_size self.hc_eps = config.hc_eps From 9ce5adbd56dd1c14907b62225014234413d36477 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 10:31:23 +0000 Subject: [PATCH 318/365] fix(executor): keep replies that answer a later RPC Ranks other than the output rank answer execute_model and sample_tokens only when they fail, so the gather of such an RPC watches them and, since #928, dropped any success it read from them as stale. It is not stale: with the batch queue, the engine sends a collective RPC such as wait_for_boundary_checkpoint_copies to every rank before it has gathered the replies of the single-rank RPCs ahead of it. The other ranks run both and answer the collective first; the gather of execute_model read and dropped those answers ("Dropped a reply from rank N to execute_model"), and the collective waited for them until VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS ran out and the engine stopped with TimeoutError. Keep such a reply per rank, in arrival order, and let the gather of the RPC it answers take it first. A failure from a watched rank still ends the wait at once. Reproduced on GLM-5.3-Flash TP4/DCP4/MTP3 with 8 parallel 30K-token prompts (max_tokens=1): the first burst stopped the engine after 318 s. Reported by thecrimea on Discord. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-950.json | 23 +++++++ .../test_multiproc_gather_responses.py | 67 +++++++++++++++++-- vllm/v1/executor/multiproc_executor.py | 52 ++++++++------ vllm/v1/executor/ray_executor_v2.py | 1 + 4 files changed, 117 insertions(+), 26 deletions(-) create mode 100644 .lil/changes/vllm-950.json diff --git a/.lil/changes/vllm-950.json b/.lil/changes/vllm-950.json new file mode 100644 index 000000000000..98030f8c5a15 --- /dev/null +++ b/.lil/changes/vllm-950.json @@ -0,0 +1,23 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-950", + "category": "fix", + "summary": "GLM-5.3-Flash and Qwen3.8-Flash-Next on two or more GPUs no longer stop with \"RPC call to wait_for_boundary_checkpoint_copies timed out\"", + "models": [ + "GLM-5.3-Flash", + "Qwen3.8-Flash-Next" + ], + "compatibility": "No user action required. Affected were beta images from karmic-kraken-beta-20260928-d7935130 on, serving a hybrid model (request-boundary checkpoints) with tensor parallelism. Canonical images were not affected. Models without recurrent layers, such as DeepSeek-V4.1-Flash and MiMo-V2.6-Flash, were not affected.", + "details": [ + "What was wrong: ranks other than the output rank answer execute_model and sample_tokens only when they fail, and vllm-928 made the engine watch them for such a failure. With async scheduling the engine can send the next request to every rank, such as waiting for the boundary checkpoint copies of a finished prompt, before it has read the earlier replies. The other ranks answer that request first, and the engine took their answers for stray replies and dropped them.", + "Consequence: the engine then waited for answers that never came. After about a minute it logged that an engine step was blocked, and after VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (300 s) it stopped with TimeoutError, failing the requests in flight. The GPUs were idle the whole time. Parallel long prompts triggered it most: 8 simultaneous 30K-token prompts on GLM-5.3-Flash TP4/DCP4 stopped the engine in the first burst. Such a stop is preceded by \"Dropped a reply from rank N to execute_model\" (or sample_tokens) warnings in the log.", + "Now a reply that answers a later request is kept for it, and a failure from any rank still stops the engine at once. Reported by thecrimea with a reproducer." + ], + "pull_requests": [ + 950 + ], + "authors": [ + "Local Inference Lab" + ], + "requires": [] +} diff --git a/tests/v1/executor/test_multiproc_gather_responses.py b/tests/v1/executor/test_multiproc_gather_responses.py index 5653520714a8..9a909c85529b 100644 --- a/tests/v1/executor/test_multiproc_gather_responses.py +++ b/tests/v1/executor/test_multiproc_gather_responses.py @@ -4,16 +4,22 @@ A rank that never replies (for example, stuck in a collective another rank failed out of) must not hide a failure another rank has already reported, -including a rank that does not own the reply of a single-rank RPC. +including a rank that does not own the reply of a single-rank RPC. A reply +read while gathering an earlier RPC is kept for the RPC it answers. """ import threading import time +from collections import deque import pytest from vllm.distributed.device_communicators.shm_broadcast import MessageQueue -from vllm.v1.executor.multiproc_executor import WorkerProc, _gather_responses +from vllm.v1.executor.multiproc_executor import ( + MultiprocExecutor, + WorkerProc, + _gather_responses, +) SUCCESS = WorkerProc.ResponseStatus.SUCCESS FAILURE = WorkerProc.ResponseStatus.FAILURE @@ -130,11 +136,60 @@ def test_a_single_rank_reply_is_read_while_watching(): assert quiet.checks >= 2 and quiet.reply is None -def test_an_unexpected_success_from_a_watched_rank_is_dropped(): - stale = _Queue((SUCCESS, "late")) +def test_a_success_from_a_watched_rank_is_held_for_the_rpc_it_answers(): + early = _Queue((SUCCESS, "copies 1")) output = _Queue((SUCCESS, "r0"), ready_after=1) - assert _gather_responses([output], (0,), None, "m", [(1, stale)]) == ["r0"] - assert stale.reply is None + held: dict = {} + assert _gather_responses([output], (0,), None, "m", [(1, early)], held) == ["r0"] + later = [_Queue((SUCCESS, "copies 0")), _Queue()] + assert _gather_responses(later, range(2), None, "copies", held=held) == [ + "copies 0", + "copies 1", + ] + assert not held[1] + + +class _Replies: + """One rank's reply queue, holding what the rank has already sent.""" + + def __init__(self, *replies) -> None: + self.replies = list(replies) + + def ready(self) -> bool: + return bool(self.replies) + + def dequeue(self, timeout=None): + assert self.replies + return self.replies.pop(0) + + def wait_for_message(self, timeout_ms: int) -> None: + time.sleep(timeout_ms / 1000) + + +class _Broadcast: + def enqueue(self, obj) -> None: + pass + + +def test_a_collective_rpc_sent_behind_a_single_rank_rpc_gets_every_reply(): + """The batch queue sends wait_for_boundary_checkpoint_copies to every rank + before it gathers execute_model, which only rank 0 answers. Ranks 1 and 2 + run both and answer the collective first. Dropping those replies left the + collective waiting out its timeout, which ended the engine (GLM TP4).""" + executor = object.__new__(MultiprocExecutor) + executor.rpc_broadcast_mq = _Broadcast() + executor.is_failed = False + executor.futures_queue = deque() + executor.held_replies = {} + executor.response_mqs = [ + _Replies((SUCCESS, "step"), (SUCCESS, "copies 0")), + _Replies((SUCCESS, "copies 1")), + _Replies((SUCCESS, "copies 2")), + ] + step = executor.collective_rpc("execute_model", non_block=True, unique_reply_rank=0) + copies = executor.collective_rpc("wait_for_boundary_checkpoint_copies", timeout=1) + assert copies == ["copies 0", "copies 1", "copies 2"] + assert step.result() == "step" def test_the_deadline_applies_while_watching(): diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 5f1f1483fb51..efb7132e486d 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -100,6 +100,7 @@ def _gather_responses( deadline: float | None, method: str, watched: Sequence[tuple[int, MessageQueue]] = (), + held: dict[int, deque[tuple[Any, Any]]] | None = None, ) -> list[Any]: """Collect one reply per queue, failing as soon as any rank reports one. @@ -109,15 +110,23 @@ def _gather_responses( in the next all-reduce, and the engine waits on TP0 with TP1's error in its queue. ``watched`` are the other ranks of a single-rank RPC such as ``execute_model``: they reply only when they fail, and the output rank is - then blocked in a collective the failed rank has left, so a message from + then blocked in a collective the failed rank has left, so a failure from one of them ends the wait with its error. + A success from a watched rank answers a later RPC that rank has already + run: with the batch queue, the engine sends a collective RPC such as + ``wait_for_boundary_checkpoint_copies`` before it has gathered the replies + of earlier single-rank RPCs. Such a reply is kept in ``held``, per rank in + arrival order, and the gather of the RPC it answers takes it first. + Every queue is checked without consuming anything, and a reply is read only once it has arrived, then with the RPC's own deadline. Between checks the gather waits on the oldest pending queue the way a blocking read does: spinning while replies are frequent, otherwise parked on its notification for at most RESPONSE_WAIT_SLICE_MS. """ + if held is None: + held = {} def remaining() -> float | None: if deadline is None: @@ -128,10 +137,13 @@ def remaining() -> float | None: return left if len(response_mqs) == 1 and not watched: - try: - status, result = response_mqs[0].dequeue(timeout=remaining()) - except TimeoutError as e: - raise TimeoutError(f"RPC call to {method} timed out.") from e + if replies := held.get(ranks[0]): + status, result = replies.popleft() + else: + try: + status, result = response_mqs[0].dequeue(timeout=remaining()) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e return [_check_response(ranks[0], status, result)] responses: list[Any] = [None] * len(response_mqs) @@ -143,14 +155,17 @@ def remaining() -> float | None: while pending: progressed = False for index in list(pending): - # ready() never consumes, so a reply is only read once it has - # arrived, and then with the RPC's own deadline. - if not response_mqs[index].ready(): - continue - try: - status, result = response_mqs[index].dequeue(timeout=remaining()) - except TimeoutError as e: - raise TimeoutError(f"RPC call to {method} timed out.") from e + if replies := held.get(ranks[index]): + status, result = replies.popleft() + else: + # ready() never consumes, so a reply is only read once it has + # arrived, and then with the RPC's own deadline. + if not response_mqs[index].ready(): + continue + try: + status, result = response_mqs[index].dequeue(timeout=remaining()) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e responses[index] = _check_response(ranks[index], status, result) pending.remove(index) progressed = True @@ -162,12 +177,7 @@ def remaining() -> float | None: except TimeoutError as e: raise TimeoutError(f"RPC call to {method} timed out.") from e _check_response(rank, status, result) - logger.warning( - "Dropped a reply from rank %d to %s, which only rank %d answers.", - rank, - method, - ranks[0], - ) + held.setdefault(rank, deque()).append((status, result)) if pending and not progressed: left = remaining() # raises once the deadline has passed if time.monotonic() - started >= interval * (warnings + 1): @@ -366,6 +376,8 @@ def _init_executor(self) -> None: response_mq.wait_until_ready() self.futures_queue = deque[FutureWrapper]() + # Replies gathered ahead of the RPC they answer, by rank. + self.held_replies: dict[int, deque[tuple[Any, Any]]] = {} self._post_init_executor() @@ -552,7 +564,7 @@ def _aggregate(outputs: Any) -> Any: def get_response(): responses = _gather_responses( - response_mqs, ranks, deadline, str(method), watched + response_mqs, ranks, deadline, str(method), watched, self.held_replies ) return responses[0] if output_rank is not None else responses diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py index 8f4b84ed0f13..dd80095d4459 100644 --- a/vllm/v1/executor/ray_executor_v2.py +++ b/vllm/v1/executor/ray_executor_v2.py @@ -477,6 +477,7 @@ def _init_executor(self) -> None: response_mq.wait_until_ready() self.futures_queue = deque[FutureWrapper]() + self.held_replies = {} self._post_init_executor() self.start_worker_monitor() From 16c52365e1c01e7b38468d2e0270cec648e36890 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 11:15:27 +0000 Subject: [PATCH 319/365] chore(release): wheels require CUTLASS DSL 4.7.1 with B12X B12X master moved to nvidia-cutlass-dsl 4.7.1 (b12x 4bacd509), and the runtime image keeps one CUTLASS DSL for B12X, FlashInfer and vLLM. The wheel normalizer accepted only the 4.6.2 foundation; it now requires 4.7.1, and the release lock names it. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-951.json | 14 ++++++++++++++ tools/jovian_wheel_release/normalize_wheel.py | 8 ++++---- tools/jovian_wheel_release/runtime.lock | 2 +- tools/jovian_wheel_release/test_normalize_wheel.py | 12 ++++++------ 4 files changed, 25 insertions(+), 11 deletions(-) create mode 100644 .lil/changes/vllm-951.json diff --git a/.lil/changes/vllm-951.json b/.lil/changes/vllm-951.json new file mode 100644 index 000000000000..430458d2b4d8 --- /dev/null +++ b/.lil/changes/vllm-951.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-951", + "category": "internal", + "summary": "vLLM wheels require CUTLASS DSL 4.7.1, the compiler B12X and FlashInfer now use", + "models": ["all"], + "compatibility": "No user action required. The first start after updating compiles B12X kernels again for the new compiler.", + "details": [ + "B12X moved to CUTLASS DSL 4.7.1. The runtime image keeps one compiler for B12X, FlashInfer and vLLM, so vLLM's wheel metadata now names 4.7.1 instead of 4.6.2." + ], + "pull_requests": [951], + "authors": ["Local Inference Lab"], + "requires": [] +} diff --git a/tools/jovian_wheel_release/normalize_wheel.py b/tools/jovian_wheel_release/normalize_wheel.py index fc32d53afad0..89b7f1f56b98 100755 --- a/tools/jovian_wheel_release/normalize_wheel.py +++ b/tools/jovian_wheel_release/normalize_wheel.py @@ -48,7 +48,7 @@ def rewrite_requirements( cutlass_dsl_version: str, ) -> bytes: """Return wheel metadata with dependencies matching the foundation ABI.""" - if cutlass_dsl_version != "4.6.2": + if cutlass_dsl_version != "4.7.1": raise ValueError(f"unreviewed CUTLASS DSL foundation: {cutlass_dsl_version}") message = BytesParser(policy=compat32).parsebytes(metadata) requirements = message.get_all("Requires-Dist", []) @@ -63,9 +63,9 @@ def rewrite_requirements( for value in requirements: requirement = Requirement(value) name = canonicalize_name(requirement.name) - # The SM120 foundation, B12X and FlashInfer share DSL 4.6.2. Upstream's - # 4.7.1 bump accompanies FA4, which does not support SM120. Keep the - # explicitly supported foundation contract; re-review future bumps. + # The SM120 foundation, B12X and FlashInfer share DSL 4.7.1 since B12X + # moved to it (b12x 4bacd509). Keep the explicitly supported + # foundation contract; re-review future bumps. if name == "nvidia-cutlass-dsl" and str(requirement.specifier) not in { "==4.6.2", "==4.7.1", diff --git a/tools/jovian_wheel_release/runtime.lock b/tools/jovian_wheel_release/runtime.lock index 15dafea52fd9..99ba2e4df497 100644 --- a/tools/jovian_wheel_release/runtime.lock +++ b/tools/jovian_wheel_release/runtime.lock @@ -15,5 +15,5 @@ pytorch.runtime-version=2.14.0a0+4fdf77b940.nv26.08 pytorch.commit=4fdf77b940a7e61d854b9be527884e38f32d335a torchvision.version=0.29.0a0+0bc41e67.nv26.8.63802676 flashinfer.requirement=0.6.18 -cutlass-dsl.version=4.6.2 +cutlass-dsl.version=4.7.1 cuda.arch-list=12.0a diff --git a/tools/jovian_wheel_release/test_normalize_wheel.py b/tools/jovian_wheel_release/test_normalize_wheel.py index 308c5094c9c8..4c148082185b 100644 --- a/tools/jovian_wheel_release/test_normalize_wheel.py +++ b/tools/jovian_wheel_release/test_normalize_wheel.py @@ -18,11 +18,11 @@ @pytest.mark.parametrize( "source_cutlass,target_cutlass", [ - ("4.6.2", "4.6.2"), - ("4.7.1", "4.6.2"), - ("4.8.0", "4.6.2"), + ("4.7.1", "4.7.1"), ("4.6.2", "4.7.1"), - ("4.6.2", "4.8.0"), + ("4.8.0", "4.7.1"), + ("4.7.1", "4.6.2"), + ("4.7.1", "4.8.0"), ], ) def test_rewrites_foundation_dependencies(source_cutlass, target_cutlass) -> None: @@ -47,7 +47,7 @@ def test_rewrites_foundation_dependencies(source_cutlass, target_cutlass) -> Non Requires-Dist: click>=8 """.replace(b"CUTLASS_VERSION", source_cutlass.encode()) - if source_cutlass == "4.8.0" or target_cutlass != "4.6.2": + if source_cutlass == "4.8.0" or target_cutlass != "4.7.1": with pytest.raises(ValueError, match="unreviewed CUTLASS DSL"): rewrite_requirements( metadata, @@ -69,7 +69,7 @@ def test_rewrites_foundation_dependencies(source_cutlass, target_cutlass) -> Non "torch==2.14.0a0+nv", "torchvision==0.29.0a0+nv", "flashinfer-python==0.6.18", - "nvidia-cutlass-dsl[cu13]==4.6.2", + "nvidia-cutlass-dsl[cu13]==4.7.1", 'fastsafetensors>=0.3.3; extra == "fastsafetensors"', "click>=8", ] From 371ca44f1e59b36abce0e048ed509d8f551ea255 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 10:38:16 +0000 Subject: [PATCH 320/365] Release dense MLA priming operands after synchronous preparation Keep compiled launchers and caller-owned serving workspaces alive without retaining temporary Q, output and metadata buffers for every layer and capacity. Add BF16 and FP8 lifetime regression checks. Co-authored-by: OpenAI Codex Signed-off-by: voipmonitor --- tests/v1/attention/test_b12x.py | 58 ++++++++++++++++++++++ vllm/v1/attention/backends/mla/b12x_mla.py | 4 +- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/v1/attention/test_b12x.py b/tests/v1/attention/test_b12x.py index cc7e105c8429..6ed5a4272f3f 100644 --- a/tests/v1/attention/test_b12x.py +++ b/tests/v1/attention/test_b12x.py @@ -1074,6 +1074,64 @@ def test_b12x_attention_runtime_page_size_comes_from_cache() -> None: _kv_page_size(key_cache, torch.empty((3, 128, 4, 128), device="meta")) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_b12x_dense_mla_preparation_releases_temporary_buffers( + dtype, monkeypatch +) -> None: + from vllm.v1.attention.backends.mla import b12x_mla + from vllm.v1.worker import workspace + + impl = object.__new__(b12x_mla.B12xMLAImpl) + impl.num_heads, impl.dcp_world_size = 6, 16 + impl.head_size, impl.kv_lora_rank = 576, 512 + impl.kv_cache_dtype = "fp8" if dtype == torch.float8_e4m3fn else "bfloat16" + impl.scale = 192**-0.5 + impl._plans = { + 8: SimpleNamespace(scratch_specs=lambda: [SimpleNamespace(nbytes=64)]) + } + monkeypatch.setattr( + workspace, + "_manager", + SimpleNamespace( + get_simultaneous=lambda *args: None, reserve_by_lane=lambda: None + ), + ) + monkeypatch.setattr( + b12x_mla, + "get_b12x_scratch_buffers", + lambda state: (torch.empty(64, dtype=torch.uint8),), + ) + references = [] + + def bind(**kwargs): + references.extend( + weakref.ref(kwargs[name]) + for name in ( + "scratch", + "q", + "output", + "cache_seqlens", + "cu_seqlens_q", + "page_table", + ) + ) + return kwargs + + state = SimpleNamespace( + bind=bind, + prime=lambda binding: None, + run=lambda binding: binding["output"].zero_(), + ) + layer = SimpleNamespace(kv_cache=torch.empty(2, 64, 576, dtype=dtype)) + call = impl._prepared_call(layer, state, 8) + call.invoke() + assert all(reference() is not None for reference in references) + published = SimpleNamespace(state=state, owners=call.owners) + del call + assert all(reference() is None for reference in references) + assert published.state is state + + def test_b12x_attention_preparation_releases_temporary_buffers() -> None: impl = object.__new__(B12xPagedAttentionImpl) impl._noncausal = None diff --git a/vllm/v1/attention/backends/mla/b12x_mla.py b/vllm/v1/attention/backends/mla/b12x_mla.py index e9929439fdfd..ad16478d240f 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla.py +++ b/vllm/v1/attention/backends/mla/b12x_mla.py @@ -439,10 +439,12 @@ def _prepared_call(self, layer, state, rows): sm_scale=self.scale, ) state.prime(binding) + # PreparationSession synchronizes the priming call before discarding it. + # Serving binds its own operands; the plan retains compiled launchers, + # not the temporary tensors used to prime each layer and capacity. return PreparedCall( run=lambda: state.run(binding), output=output, - owners=(scratch, q, output, lengths, cu, pages, scale, binding), ) def forward_mqa(self, q, kv_c_and_k_pe_cache, attn_metadata, layer): From c390adba0d35a4b91a83c9f8fd0d6b2475f7f5e4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 14:23:14 +0000 Subject: [PATCH 321/365] Load Kimi-K3 X4T experts through the lossless MXFP4 reader Preserve BF16 dense weights, select SiTU expert activation, recognize nested checkpoint metadata and bound CPU shard-copy threads. Shared scale scratch remains serialized; PP and ubatching are rejected. Port the implementation qualified on TP16 without changing native MXFP4 arithmetic. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/quantization/test_ds41_exact_mxfp4.py | 97 +++++++++++++++++++ vllm/config/model.py | 1 + .../layers/quantization/__init__.py | 3 + .../layers/quantization/kimi_x4t.py | 58 +++++++++++ .../model_loader/exact_mxfp4_loader.py | 14 ++- vllm/models/deepseek_v4_1/exact_mxfp4.py | 45 ++++++--- vllm/models/kimi_k3/nvidia/model.py | 2 +- 7 files changed, 199 insertions(+), 21 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/kimi_x4t.py diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_ds41_exact_mxfp4.py index 2de6bcd4daa3..5b855fe478d9 100644 --- a/tests/quantization/test_ds41_exact_mxfp4.py +++ b/tests/quantization/test_ds41_exact_mxfp4.py @@ -73,6 +73,32 @@ def test_exact_config_and_loader_are_registered(): ) +def test_x4t_cpu_shard_reader_bounds_threads_and_restores_on_failure( + monkeypatch, moe, owner +): + from b12x.moe.checkpoints import exact_mxfp4 as checkpoint + + from vllm.models.deepseek_v4_1 import exact_mxfp4 + from vllm.utils.torch_utils import set_default_torch_num_threads + + monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_rank", lambda: 0) + method = ExactMXFP4MoEMethod(moe, owner) + layer = torch.nn.Module() + layer.layer_name = "model.layers.17.ffn.experts" + method.create_weights(layer, 384, 5120, 576, torch.bfloat16) + + def fail_reader(*args, **kwargs): + assert torch.get_num_threads() == 1 + raise OSError("checkpoint read failed") + + monkeypatch.setattr(checkpoint, "read_exact_mxfp4_layer", fail_reader) + with set_default_torch_num_threads(4): + with pytest.raises(OSError, match="checkpoint read failed"): + method.process_weights_after_loading(layer) + assert torch.get_num_threads() == 4 + + @pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( monkeypatch, moe, owner, mode @@ -215,3 +241,74 @@ def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): assert list(dict(loader.get_all_weights(config, model))) == [ "layers.40.ffn.experts.0.w1.weight" ] + + +def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + from vllm.model_executor.layers.quantization import kimi_x4t + + owner = kimi_x4t.KimiX4TConfig.from_config( + {"format_version": 1, "checkpoint_root": "/x4t"} + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace(pipeline_parallel_size=1, use_ubatching=False), + load_config=SimpleNamespace(load_format="exact_mxfp4"), + ) + monkeypatch.setattr(kimi_x4t, "get_current_vllm_config", lambda: config) + experts = Mock(spec=RoutedExperts) + experts.moe_config = replace( + moe, + activation=MoEActivation.SITU, + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + swiglu_limit=None, + ) + assert isinstance( + owner.get_quant_method(experts, "model.layers.1.mlp.experts"), + ExactMXFP4MoEMethod, + ) + assert isinstance( + owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), + UnquantizedLinearMethod, + ) + config.parallel_config.use_ubatching = True + with pytest.raises(NotImplementedError, match="PP1 without ubatching"): + owner.get_quant_method(experts, "model.layers.1.mlp.experts") + + +def test_kimi_loader_skips_compressed_experts_with_nested_text_config(tmp_path): + from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + + tensor_dir = tmp_path / "tensors" + tensor_dir.mkdir() + filename = "model-00001.safetensors" + expert = "language_model.model.layers.1.block_sparse_moe.experts.0.w1." + dense = "language_model.model.layers.0.mlp.gate_proj.weight" + tensors = { + expert + "weight_packed": torch.ones(8, dtype=torch.uint8), + expert + "weight_scale.exact_mxfp4_fixed": torch.ones(8, dtype=torch.uint8), + dense: torch.ones((8, 8), dtype=torch.bfloat16), + } + save_file(tensors, tensor_dir / filename) + common = {"schema": SCHEMA, "codec": CODEC, "family": "kimi_k3"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": filename}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: filename for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace(), + hf_text_config=SimpleNamespace( + num_hidden_layers=93, + quantization_config={ + "quant_method": "kimi_x4t", + "checkpoint_root": str(tmp_path), + }, + ), + ) + loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + result = dict(loader.get_all_weights(config, SimpleNamespace())) + assert list(result) == [dense] + assert torch.equal(result[dense], tensors[dense]) diff --git a/vllm/config/model.py b/vllm/config/model.py index 30df9ca07aea..79b540db676d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1285,6 +1285,7 @@ def _verify_quantization(self) -> None: "modelopt_mixed", "exl3", "exact_mxfp4", + "kimi_x4t", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 95e79137d9d7..e8ca734d6a74 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -39,6 +39,7 @@ "deepseek_v4_fp8", "deepseek_v41_fp8", "exact_mxfp4", + "kimi_x4t", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -148,6 +149,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .fp_quant import FPQuantConfig from .humming import HummingConfig from .inc import INCConfig + from .kimi_x4t import KimiX4TConfig from .modelopt import ( ModelOptFp8Config, ModelOptMixedPrecisionConfig, @@ -185,6 +187,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "deepseek_v4_fp8": deepseek_config, "deepseek_v41_fp8": DeepseekV41FP8Config, "exact_mxfp4": DeepseekV41ExactMXFP4Config, + "kimi_x4t": KimiX4TConfig, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/kimi_x4t.py b/vllm/model_executor/layers/quantization/kimi_x4t.py new file mode 100644 index 000000000000..bfb5ab45b01b --- /dev/null +++ b/vllm/model_executor/layers/quantization/kimi_x4t.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 lossless MXFP4 expert scales with unchanged BF16 dense weights.""" + +import torch + +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config +from vllm.models.deepseek_v4_1.exact_mxfp4 import ExactMXFP4MoEMethod + + +class KimiX4TConfig(Mxfp4Config): + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "kimi_x4t" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if ( + user_quant in (None, cls.get_name()) + and hf_quant_cfg is not None + and hf_quant_cfg.get("quant_method") == cls.get_name() + and getattr(hf_config, "model_type", None) in ("kimi_k3", "kimi_linear") + ): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("Kimi X4T requires format_version=1 and checkpoint_root") + result = cls() + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + config = get_current_vllm_config() + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "X4T shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "exact_mxfp4": + raise ValueError("Kimi X4T requires --load-format exact_mxfp4") + return ExactMXFP4MoEMethod(layer.moe_config, self) + return super().get_quant_method(layer, prefix) diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py index 1a81c4b703f1..9fa45d9dc206 100644 --- a/vllm/model_executor/model_loader/exact_mxfp4_loader.py +++ b/vllm/model_executor/model_loader/exact_mxfp4_loader.py @@ -24,8 +24,10 @@ class ExactMXFP4ModelLoader(DefaultModelLoader): def _root(self, model_config): from b12x.moe.checkpoints.exact_mxfp4 import checkpoint_contract - quant = model_config.hf_config.quantization_config - if quant.get("quant_method") != "exact_mxfp4": + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + quant = getattr(model_config.hf_config, "quantization_config", None) + quant = quant or text_config.quantization_config + if quant.get("quant_method") not in ("exact_mxfp4", "kimi_x4t"): raise ValueError("exact_mxfp4 loading requires the X4T model config") root = Path(quant["checkpoint_root"]) if not root.is_absolute(): @@ -41,7 +43,8 @@ def get_all_weights(self, model_config, model): raise NotImplementedError("X4T does not support secondary weight sources") file_filter = getattr(model, "checkpoint_file_weight_filter", None) prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) - layers = model_config.hf_config.num_hidden_layers + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + layers = text_config.num_hidden_layers self.counter_before_loading_weights = time.perf_counter() for filename in sorted(set(contract["source_names"].values())): path = str(root / "tensors" / filename) @@ -50,7 +53,10 @@ def get_all_weights(self, model_config, model): for name in sorted(descriptors): if prefixes is not None and not name.startswith(prefixes): continue - match = re.match(r"layers\.(\d+)\.ffn\.experts\.", name) + match = re.search( + r"(?:^|\.)layers\.(\d+)\.(?:ffn|block_sparse_moe)\.experts\.", + name, + ) if match and int(match.group(1)) < layers: continue if callable(file_filter) and file_filter(name): diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/exact_mxfp4.py index 1e4536b4ae25..3efced1878f1 100644 --- a/vllm/models/deepseek_v4_1/exact_mxfp4.py +++ b/vllm/models/deepseek_v4_1/exact_mxfp4.py @@ -25,6 +25,7 @@ MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.utils.torch_utils import set_default_torch_num_threads from .quant_config import DeepseekV41FP8Config @@ -106,11 +107,18 @@ def __init__(self, moe, owner): ): raise NotImplementedError("X4T DS4.1 supports TP without EP/DP") if ( - moe.activation != MoEActivation.SILU + moe.activation not in (MoEActivation.SILU, MoEActivation.SITU) or moe.in_dtype != torch.bfloat16 or moe.has_bias + or ( + moe.activation == MoEActivation.SITU + and ( + moe.activation_situ_beta != 4.0 + or moe.activation_situ_linear_beta != 25.0 + ) + ) ): - raise ValueError("X4T DS4.1 requires bias-free BF16 SwiGLU experts") + raise ValueError("X4T requires bias-free BF16 SwiGLU or SiTU(4,25) experts") def create_weights( self, @@ -152,7 +160,7 @@ def process_weights_after_loading(self, layer): e, h, n = self.num_experts, self.hidden_size, self.local_intermediate shapes = ((e, h // 32, 2 * n), (e, n // 32, h)) if self.owner.scale_scratch is None: - # DS4.1 disallows ubatching: every layer consumes these scale grids + # X4T disallows ubatching: every layer consumes these scale grids # before its successor may overwrite them on the same stream. self.owner.scale_scratch = tuple( torch.empty(s, dtype=torch.uint8, device=device) for s in shapes @@ -162,23 +170,28 @@ def process_weights_after_loading(self, layer): t.shape != shape or t.device != device for t, shape in zip(scratch, shapes) ): raise ValueError("X4T shared scale scratch geometry/device mismatch") - weights = read_exact_mxfp4_layer( - self.owner.checkpoint_root, - self.layer_index, - num_experts=e, - hidden_size=h, - intermediate_size=n * tp, - tp_rank=rank, - tp_size=tp, - device=device, - w13_scale_scratch=scratch[0], - w2_scale_scratch=scratch[1], - ) + # Per-expert CPU slices are too small to amortize intra-op barriers. + # Restore the serving thread policy before kernel preparation. + with set_default_torch_num_threads(1): + weights = read_exact_mxfp4_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) plan = fused_moe.plan_weights( source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), activation=fused_moe.ActivationSpec( mode="a16", - nonlinearity="silu", + nonlinearity="situ" + if self.moe.activation == MoEActivation.SITU + else "silu", io_dtype=torch.bfloat16, swiglu_limit=self.moe.swiglu_limit, ), diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index f979c669e27f..8d313664a54b 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -766,7 +766,7 @@ def __init__( config, "min_moe_intermediate_per_partition", 256 ) native_b12x_mxfp4 = ( - vllm_config.model_config.quantization == "mxfp4" + vllm_config.model_config.quantization in ("mxfp4", "kimi_x4t") and vllm_config.kernel_config.moe_backend == "b12x" ) self._aligned_decode_projections = ( From 5f8a5755b6464543e4cbf5e5e6f31e49999ee050 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 14:29:51 +0000 Subject: [PATCH 322/365] Document transient MLA preparation ownership for runtime releases Co-authored-by: OpenAI Codex --- .../vllm-b12x-mla-transient-preparation.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .lil/changes/vllm-b12x-mla-transient-preparation.json diff --git a/.lil/changes/vllm-b12x-mla-transient-preparation.json b/.lil/changes/vllm-b12x-mla-transient-preparation.json new file mode 100644 index 000000000000..16b9e242d8a5 --- /dev/null +++ b/.lil/changes/vllm-b12x-mla-transient-preparation.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-b12x-mla-transient-preparation", + "category": "fix", + "summary": "Release temporary dense MLA operands after synchronized kernel preparation.", + "models": [ + "Kimi-K3" + ], + "compatibility": "Prepared serving rebinds caller-owned buffers. Kernel arithmetic, cache representation and public configuration are unchanged.", + "pull_requests": [], + "authors": [ + "voipmonitor" + ] +} From 120a909da7ebd8139dc6d300d5b615981ae2d5d7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 14:31:55 +0000 Subject: [PATCH 323/365] Declare Kimi X4T loader compatibility for runtime releases Co-authored-by: OpenAI Codex --- .lil/changes/vllm-kimi-x4t-experts.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-kimi-x4t-experts.json diff --git a/.lil/changes/vllm-kimi-x4t-experts.json b/.lil/changes/vllm-kimi-x4t-experts.json new file mode 100644 index 000000000000..753195799499 --- /dev/null +++ b/.lil/changes/vllm-kimi-x4t-experts.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-kimi-x4t-experts", + "category": "feature", + "summary": "Load lossless Kimi-K3 X4T experts with native dense precision.", + "models": [ + "Kimi-K3" + ], + "compatibility": "Requires B12X X4T checkpoint support, --quantization kimi_x4t and --load-format exact_mxfp4. Native MXFP4 and DS4.1 source precision are unchanged.", + "pull_requests": [ + 949 + ], + "authors": [ + "voipmonitor" + ] +} From da9fc45f696379d880503cc5356e7fea9d39bcd5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 16:31:52 +0000 Subject: [PATCH 324/365] fix(release): install CUTLASS DSL 4.7.1 in the wheel build environment The wheel builder asserts that its build venv has the CUTLASS DSL version the wheel will require (CUTLASS_DSL_VERSION from runtime.lock). #951 moved that to 4.7.1, but the build venv only saw the NGC builder's 4.6.2 through its system site packages, and wheel run 36725514670 stopped at the assertion. Add the five 4.7.1 packages to build-requirements.lock (the x86_64 CPython 3.12 wheels; the same hashes as the docker runtime lock and the FlashInfer builder, whose build succeeded this way). Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-954.json | 14 ++++++++++++++ tools/jovian_wheel_release/build-requirements.lock | 10 ++++++++++ 2 files changed, 24 insertions(+) create mode 100644 .lil/changes/vllm-954.json diff --git a/.lil/changes/vllm-954.json b/.lil/changes/vllm-954.json new file mode 100644 index 000000000000..f108101d5af4 --- /dev/null +++ b/.lil/changes/vllm-954.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-954", + "category": "internal", + "summary": "vLLM wheels build with CUTLASS DSL 4.7.1 in the build environment", + "models": ["all"], + "compatibility": "No user action required.", + "details": [ + "The wheel builder checks that its build environment has the CUTLASS DSL version the wheel will require. After vllm-951 moved that to 4.7.1, the builder still had the NGC image's 4.6.2 and stopped. It now installs 4.7.1 with pinned hashes, as the FlashInfer builder and the runtime image do." + ], + "pull_requests": [954], + "authors": ["Local Inference Lab"], + "requires": [] +} diff --git a/tools/jovian_wheel_release/build-requirements.lock b/tools/jovian_wheel_release/build-requirements.lock index a5fa56d7444e..8a6e2d3eb963 100644 --- a/tools/jovian_wheel_release/build-requirements.lock +++ b/tools/jovian_wheel_release/build-requirements.lock @@ -5,6 +5,16 @@ jinja2==3.1.6 \ --hash=sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67 ninja==1.13.0 \ --hash=sha256:fb46acf6b93b8dd0322adc3a4945452a4e774b75b91293bafcc7b7f8e6517dfa +nvidia-cutlass-dsl==4.7.1 \ + --hash=sha256:765850c74868a0b5a02fefbfc1b380842b64e5f7a1c8b1655e322a3b14c01726 +nvidia-cutlass-dsl-libs-base==4.7.1 \ + --hash=sha256:c2b67dacabfc905e4d132784fe4848647f40b0b97b1cdde3ce4a482642536b84 +nvidia-cutlass-dsl-libs-core==4.7.1 \ + --hash=sha256:c65662abc9ab18fd997c57547895736928e559f4624bbe2aa1aa91418ea014ca +nvidia-cutlass-dsl-libs-cu12==4.7.1 \ + --hash=sha256:c8a0f93a1cde41404696b3194dd1f589c6f3a7b426bcabe97845f8fe96f949f5 +nvidia-cutlass-dsl-libs-cu13==4.7.1 \ + --hash=sha256:20c14e16125cb8e36bc9947a9f0a57a660f44651365337b911c84abde4967f5f packaging==26.2 \ --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e semantic-version==2.10.0 \ From e530531d2d30284e3042afa2e5f13a12fd3f6c34 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 23:58:34 +0000 Subject: [PATCH 325/365] Load losslessly compressed GLM NVFP4 scales with native A4 experts Co-authored-by: Codex Signed-off-by: Martin Vit --- tests/quantization/test_nvfp4_lsc.py | 159 +++++++++++ vllm/config/load.py | 2 + vllm/config/model.py | 1 + vllm/model_executor/layers/fused_moe/b12x.py | 11 +- .../layers/quantization/__init__.py | 3 + .../layers/quantization/nvfp4_lsc.py | 253 ++++++++++++++++++ vllm/model_executor/model_loader/__init__.py | 3 + .../model_loader/nvfp4_lsc_loader.py | 49 ++++ vllm/models/glm5next/nvidia/model.py | 2 +- 9 files changed, 480 insertions(+), 3 deletions(-) create mode 100644 tests/quantization/test_nvfp4_lsc.py create mode 100644 vllm/model_executor/layers/quantization/nvfp4_lsc.py create mode 100644 vllm/model_executor/model_loader/nvfp4_lsc_loader.py diff --git a/tests/quantization/test_nvfp4_lsc.py b/tests/quantization/test_nvfp4_lsc.py new file mode 100644 index 000000000000..90c3e742d5cc --- /dev/null +++ b/tests/quantization/test_nvfp4_lsc.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compressed expert storage must retain nonexpert precision and tensors.""" + +import json +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file + +from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.quantization.nvfp4_lsc import ( + Nvfp4LscConfig, + Nvfp4LscMoEMethod, +) +from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader + + +def test_lsc_config_preserves_source_recipes_and_avoids_expert_allocations(): + original = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.3.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16}, + "model.layers.45.mlp.experts": {"quant_algo": "MXFP8"}, + }, + } + owner = Nvfp4LscConfig.from_config( + { + "format_version": 1, + "checkpoint_root": "/lsc", + "source_quantization_config": original, + } + ) + assert owner._resolve_quant_algo("model.layers.3.mlp.experts") == "NVFP4" + assert owner._resolve_quant_algo("model.layers.45.mlp.experts") == "MXFP8" + moe = FusedMoEConfig( + num_experts=288, + num_local_experts=288, + num_logical_experts=288, + experts_per_token=8, + hidden_dim=4096, + intermediate_size=1024, + in_dtype=torch.bfloat16, + device="cpu", + activation=MoEActivation.SILU, + swiglu_limit=10.0, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + ) + method = Nvfp4LscMoEMethod(moe, owner) + layer = torch.nn.Module() + layer.layer_name = "model.layers.3.mlp.experts" + method.create_weights(layer, 288, 4096, 1024, torch.bfloat16) + assert not list(layer.parameters()) and not layer.state_dict() + assert layer.w13_weight.numel() == layer.w2_weight.numel() == 0 + with pytest.raises(ValueError, match="source_quantization_config"): + Nvfp4LscConfig.from_config({"format_version": 1, "checkpoint_root": "/lsc"}) + + +def test_prepared_nvfp4_accepts_kernel_order_and_rejects_activation_change(monkeypatch): + from b12x.moe import fused_moe + + from vllm.model_executor.layers.fused_moe import b12x + + moe = FusedMoEConfig( + num_experts=8, + num_local_experts=8, + num_logical_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=128, + in_dtype=torch.bfloat16, + device="cpu", + activation=MoEActivation.SILU, + swiglu_limit=10.0, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + ) + one = torch.ones(8) + quant = nvfp4_moe_quant_config(one, one, one, one, one, one, gemm1_clamp_limit=10.0) + backend = b12x.B12xExperts(moe, quant) + layer = torch.nn.Module() + layer.activation, layer.apply_router_weight_on_input = moe.activation, False + prepared = Mock(spec=fused_moe.PreparedExperts) + activation = fused_moe.ActivationSpec( + mode="a4", nonlinearity="silu", io_dtype=torch.bfloat16, swiglu_limit=10.0 + ) + prepared.plan = SimpleNamespace( + source=fused_moe.PackedSource(format="modelopt_nvfp4", w13_layout="w13"), + activation=activation, + ) + prepared.num_experts, prepared.hidden_size, prepared.intermediate_size = 8, 256, 128 + monkeypatch.setattr( + b12x, "_register_b12x_moe_output_collective", lambda *a, **k: None + ) + backend.install_prepared_experts(layer, prepared) + assert backend._prepared_experts is prepared + prepared.plan.activation = replace(activation, mode="a16") + with pytest.raises(ValueError, match="activation"): + backend.install_prepared_experts(layer, prepared) + + +def test_loader_excludes_compressed_main_experts_and_retains_native_tensors(tmp_path): + from b12x.moe.checkpoints.nvfp4_lsc import CODEC, SCHEMA + + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.model_executor.model_loader import get_model_loader + + directory = tmp_path / "tensors" + directory.mkdir() + filename = "weights.safetensors" + tensors = { + "model.language_model.layers.3.mlp.experts.0.gate_proj.weight": torch.ones( + 8, dtype=torch.uint8 + ), + ( + "model.language_model.layers.3.mlp.experts.0.gate_proj." + "weight_scale.nvfp4_lsc_fixed" + ): torch.ones(8, dtype=torch.uint8), + "model.language_model.layers.45.mlp.experts.0.gate_proj.weight": torch.full( + (8,), 42, dtype=torch.uint8 + ), + "model.language_model.norm.weight": torch.ones(8, dtype=torch.bfloat16), + } + save_file(tensors, directory / filename) + common = {"schema": SCHEMA, "codec": CODEC, "family": "glm53_nvfp4"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": filename}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: filename for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace( + quantization_config={ + "quant_method": "nvfp4_lsc", + "checkpoint_root": str(tmp_path), + } + ), + hf_text_config=SimpleNamespace(num_hidden_layers=45), + ) + loader = get_model_loader(LoadConfig(load_format="nvfp4_lsc")) + assert isinstance(loader, Nvfp4LscModelLoader) + assert get_quantization_config("nvfp4_lsc") is Nvfp4LscConfig + actual = dict(loader.get_all_weights(config, SimpleNamespace())) + assert set(actual) == {name for name in tensors if ".layers.3." not in name} + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/vllm/config/load.py b/vllm/config/load.py index 6729f6820c45..7585b22cd73c 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -37,6 +37,8 @@ class LoadConfig: - "safetensors" will load the weights in the safetensors format. - "exact_mxfp4" reads losslessly compressed DS4.1 MXFP4 expert scales and retains file-backed loading for Engram tables. + - "nvfp4_lsc" reads losslessly compressed GLM NVFP4 expert scales and + retains native NVFP4 activation quantization. - "instanttensor" will load the Safetensors weights on CUDA devices using InstantTensor, which enables distributed loading with pipelined prefetching and fast direct I/O. diff --git a/vllm/config/model.py b/vllm/config/model.py index 79b540db676d..efab4ebfec84 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1286,6 +1286,7 @@ def _verify_quantization(self) -> None: "exl3", "exact_mxfp4", "kimi_x4t", + "nvfp4_lsc", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 7f9f129ebe4c..4ef96fef3d66 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -511,14 +511,21 @@ def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> Non and isinstance(prepared.plan.source, fused_moe.PackedSource) and prepared.plan.source.format == "fp4_e8m0_k32" ) - if not (trellis or packed): + nvfp4 = ( + self._source_format == "modelopt_nvfp4" + and self._quant_mode == "nvfp4" + and isinstance(prepared.plan.source, fused_moe.PackedSource) + and prepared.plan.source.format == "modelopt_nvfp4" + and prepared.plan.source.w13_layout in (self._w13_layout, "w13") + ) + if not (trellis or packed or nvfp4): raise TypeError("Prepared expert encoding does not match the backend") if ( prepared.num_experts != self.moe_config.num_experts or prepared.hidden_size != self.moe_config.hidden_dim or prepared.intermediate_size != self.moe_config.intermediate_size_per_partition - or prepared.plan.activation.mode != "a16" + or prepared.plan.activation.mode != ("a4" if nvfp4 else "a16") or (trellis and prepared.plan.activation.rotation_dtype != torch.float16) or prepared.plan.activation.io_dtype != self.moe_config.in_dtype or prepared.plan.activation.nonlinearity diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index e8ca734d6a74..63a584625add 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -40,6 +40,7 @@ "deepseek_v41_fp8", "exact_mxfp4", "kimi_x4t", + "nvfp4_lsc", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -158,6 +159,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: ) from .moe_wna16 import MoeWNA16Config from .mxfp4 import GptOssMxfp4Config, Mxfp4Config + from .nvfp4_lsc import Nvfp4LscConfig from .online.base import OnlineQuantizationConfig from .torchao import TorchAOConfig @@ -188,6 +190,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "deepseek_v41_fp8": DeepseekV41FP8Config, "exact_mxfp4": DeepseekV41ExactMXFP4Config, "kimi_x4t": KimiX4TConfig, + "nvfp4_lsc": Nvfp4LscConfig, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/nvfp4_lsc.py b/vllm/model_executor/layers/quantization/nvfp4_lsc.py new file mode 100644 index 000000000000..f120658c55d7 --- /dev/null +++ b/vllm/model_executor/layers/quantization/nvfp4_lsc.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lossless NVFP4 expert scale storage with native ModelOpt arithmetic.""" + +from dataclasses import replace + +import regex as re +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.utils.torch_utils import set_default_torch_num_threads + +from .modelopt import ModelOptMixedPrecisionConfig + +logger = init_logger(__name__) + + +class Nvfp4LscConfig(ModelOptMixedPrecisionConfig): + """Compressed main experts; retained ModelOpt formats for all other tensors.""" + + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "nvfp4_lsc" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if (hf_quant_cfg or {}).get("quant_method") == cls.get_name(): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("NVFP4-LSC requires format_version=1 and checkpoint_root") + original = config.get("source_quantization_config") + if not isinstance(original, dict): + raise ValueError("NVFP4-LSC requires source_quantization_config") + result = super().from_config(original) + assert isinstance(result, cls) + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) + if match is None: + raise ValueError("NVFP4-LSC routed experts require a numbered layer") + config = get_current_vllm_config() + if ( + int(match.group(1)) + < config.model_config.hf_text_config.num_hidden_layers + ): + if config.model_config.hf_text_config.model_type != "glm5_next_text": + raise ValueError("NVFP4-LSC supports GLM-5.3-Flash expert geometry") + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "NVFP4-LSC shared scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "nvfp4_lsc": + raise ValueError("NVFP4-LSC requires --load-format nvfp4_lsc") + if self._resolve_quant_algo(prefix) != "NVFP4": + raise ValueError( + "NVFP4-LSC requires NVFP4 source expert calibration" + ) + return Nvfp4LscMoEMethod(layer.moe_config, self) + return super().get_quant_method(layer, prefix) + + +class Nvfp4LscMoEMethod(FusedMoEMethodBase): + """Decode selected scale planes into model-owned native E4M3 scratch.""" + + def __init__(self, moe, owner): + super().__init__(moe) + self.owner = owner + parallel = moe.moe_parallel_config + if ( + parallel.use_ep + or parallel.ep_size != 1 + or parallel.dp_size != 1 + or parallel.use_all2all_kernels + or parallel.enable_eplb + ): + raise NotImplementedError("NVFP4-LSC supports TP without EP/DP") + if ( + moe.activation != MoEActivation.SILU + or moe.in_dtype != torch.bfloat16 + or moe.has_bias + ): + raise ValueError("NVFP4-LSC requires bias-free BF16 SwiGLU experts") + if envs.VLLM_B12X_MOE_FP4_FORCE_A16: + raise NotImplementedError("NVFP4-LSC requires native NVFP4 A4 activations") + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) + if match is None or params_dtype != torch.bfloat16: + raise ValueError("NVFP4-LSC requires a numbered BF16 expert layer") + self.layer_index = int(match.group(1)) + self.num_experts, self.hidden_size = num_experts, hidden_size + self.local_intermediate = intermediate_size_per_partition + for name in ("w13_weight", "w2_weight"): + layer.register_buffer( + name, torch.empty(0, dtype=torch.uint8), persistent=False + ) + + def get_fused_moe_quant_config(self, layer): + return self.moe_quant_config + + def process_weights_after_loading(self, layer): + from b12x.moe import fused_moe + from b12x.moe.checkpoints.nvfp4_lsc import read_nvfp4_lsc_layer + + tp, rank = ( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + ) + device = layer.w13_weight.device + e, h, n = self.num_experts, self.hidden_size, self.local_intermediate + shapes = ((e, 2 * n, h // 16), (e, h, n // 16)) + if self.owner.scale_scratch is None: + self.owner.scale_scratch = tuple( + torch.empty(s, dtype=torch.float8_e4m3fn, device=device) for s in shapes + ) + scratch = self.owner.scale_scratch + if any( + tuple(t.shape) != s or t.device != device for t, s in zip(scratch, shapes) + ): + raise ValueError("NVFP4-LSC shared scratch geometry/device mismatch") + with set_default_torch_num_threads(1): + weights = read_nvfp4_lsc_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) + packed = weights.packed + layer_max = envs.VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE + packed = replace( + packed, + input_scale=( + packed.input_scale.amin() + if layer_max in ("1", "all", "w13") + else packed.input_scale + ), + intermediate_scale=( + packed.intermediate_scale.amin() + if layer_max in ("1", "all", "w2") + else packed.intermediate_scale + ), + ) + weights = replace(weights, packed=packed) + plan = fused_moe.plan_weights( + source=fused_moe.PackedSource(format="modelopt_nvfp4", w13_layout="w13"), + activation=fused_moe.ActivationSpec( + mode="a4", + nonlinearity="silu", + io_dtype=torch.bfloat16, + swiglu_limit=self.moe.swiglu_limit, + ), + geometry=fused_moe.MoEGeometry( + num_experts=e, hidden_size=h, intermediate_size=n + ), + ) + prepared = fused_moe.prepare_weights(plan=plan, weights=weights) + self.moe_quant_config = nvfp4_moe_quant_config( + g1_alphas=packed.w13_global_scales, + g2_alphas=packed.w2_global_scales, + a1_gscale=packed.input_scale, + a2_gscale=packed.intermediate_scale, + w1_scale=scratch[0], + w2_scale=scratch[1], + gemm1_clamp_limit=self.moe.swiglu_limit, + ) + backend = B12xExperts(self.moe, self.moe_quant_config) + backend.install_prepared_experts(layer, prepared) + self.moe_kernel = mk.FusedMoEKernel( + MoEPrepareAndFinalizeNoDPEPModular(), backend + ) + logger.info( + "NVFP4-LSC layer %d rank %d/%d: native NVFP4 A4, " + "shared scale scratch %d bytes", + self.layer_index, + rank, + tp, + sum(t.numel() for t in scratch), + ) + + def apply( + self, + layer, + x, + topk_weights, + topk_ids, + shared_experts, + shared_experts_input, + workspace=None, + ): + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + workspace=workspace, + ) diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index d542fc862463..46b73d24e999 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -15,6 +15,7 @@ from vllm.model_executor.model_loader.modelexpress_loader import ( ModelExpressModelLoader, ) +from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader from vllm.model_executor.model_loader.runai_streamer_loader import ( RunaiModelStreamerLoader, ) @@ -38,6 +39,7 @@ "fastsafetensors", "instanttensor", "exact_mxfp4", + "nvfp4_lsc", "ipc_cache", "mistral", "modelexpress", @@ -56,6 +58,7 @@ "fastsafetensors": DefaultModelLoader, "instanttensor": DefaultModelLoader, "exact_mxfp4": ExactMXFP4ModelLoader, + "nvfp4_lsc": Nvfp4LscModelLoader, "ipc_cache": IpcModelLoader, "mistral": DefaultModelLoader, "modelexpress": ModelExpressModelLoader, diff --git a/vllm/model_executor/model_loader/nvfp4_lsc_loader.py b/vllm/model_executor/model_loader/nvfp4_lsc_loader.py new file mode 100644 index 000000000000..7c0db02f450f --- /dev/null +++ b/vllm/model_executor/model_loader/nvfp4_lsc_loader.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load retained tensors while the quantization method owns compressed experts.""" + +import time +from pathlib import Path + +import regex as re +from safetensors import safe_open + +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader + + +class Nvfp4LscModelLoader(DefaultModelLoader): + def _root(self, model_config): + from b12x.moe.checkpoints.nvfp4_lsc import checkpoint_contract + + quant = getattr(model_config.hf_config, "quantization_config", None) + quant = quant or model_config.hf_text_config.quantization_config + if quant.get("quant_method") != "nvfp4_lsc": + raise ValueError("NVFP4-LSC loader requires quant_method=nvfp4_lsc") + root = Path(quant["checkpoint_root"]) + if not root.is_absolute(): + raise ValueError("NVFP4-LSC checkpoint_root must be an absolute local path") + return root, checkpoint_contract(str(root.resolve())) + + def download_model(self, model_config): + self._root(model_config) + + def get_all_weights(self, model_config, model): + root, contract = self._root(model_config) + if getattr(model, "secondary_weights", ()): + raise NotImplementedError( + "NVFP4-LSC does not support secondary weight sources" + ) + prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) + layers = model_config.hf_text_config.num_hidden_layers + self.counter_before_loading_weights = time.perf_counter() + for filename in sorted(set(contract["source_names"].values())): + with safe_open( + root / "tensors" / filename, framework="pt", device="cpu" + ) as handle: + for name in sorted(handle.keys()): + if prefixes is not None and not name.startswith(prefixes): + continue + match = re.search(r"(?:^|\.)layers\.(\d+)\.mlp\.experts\.", name) + if match and int(match.group(1)) < layers: + continue + yield name, handle.get_tensor(name) diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index a662f6b4c904..1d75464a252b 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -368,7 +368,7 @@ def __init__( mla_quant_config = ( quant_config if quant_config is not None - and quant_config.get_name() == "modelopt_mixed" + and quant_config.get_name() in ("modelopt_mixed", "nvfp4_lsc") else None ) self.self_attn = Glm5NextMLAAttention( From 77e3dd8b01bc8e9f484005695c975faabe788a58 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 00:28:54 +0000 Subject: [PATCH 326/365] Expose MXFP4-CSF and NVFP4-CSF with compatible checkpoint loading Keep model-specific retained precision, calibrated NVFP4 activations, legacy X4T/LSC configuration names and serialized tensor bytes. Shared CSF names select the matching model family. Co-authored-by: Codex Signed-off-by: Martin Vit --- docs/features/quantization/fp4_csf.md | 62 +++++++++++++++++++ tests/quantization/test_ds41_exact_mxfp4.py | 60 +++++++++++++++++- tests/quantization/test_nvfp4_lsc.py | 16 +++-- vllm/config/load.py | 7 ++- vllm/config/model.py | 2 + .../layers/quantization/__init__.py | 6 ++ .../layers/quantization/kimi_x4t.py | 12 ++-- .../layers/quantization/mxfp4_csf.py | 47 ++++++++++++++ .../layers/quantization/nvfp4_csf.py | 20 ++++++ .../layers/quantization/nvfp4_lsc.py | 28 ++++----- vllm/model_executor/model_loader/__init__.py | 6 ++ .../model_loader/exact_mxfp4_loader.py | 12 ++-- .../model_loader/mxfp4_csf_loader.py | 7 +++ .../model_loader/nvfp4_csf_loader.py | 7 +++ .../model_loader/nvfp4_lsc_loader.py | 8 +-- vllm/models/deepseek_v4_1/exact_mxfp4.py | 26 ++++---- vllm/models/deepseek_v4_1/nvidia/model.py | 1 + vllm/models/glm5next/nvidia/model.py | 3 +- vllm/models/kimi_k3/nvidia/model.py | 2 +- 19 files changed, 282 insertions(+), 50 deletions(-) create mode 100644 docs/features/quantization/fp4_csf.md create mode 100644 vllm/model_executor/layers/quantization/mxfp4_csf.py create mode 100644 vllm/model_executor/layers/quantization/nvfp4_csf.py create mode 100644 vllm/model_executor/model_loader/mxfp4_csf_loader.py create mode 100644 vllm/model_executor/model_loader/nvfp4_csf_loader.py diff --git a/docs/features/quantization/fp4_csf.md b/docs/features/quantization/fp4_csf.md new file mode 100644 index 000000000000..04004c36ec2d --- /dev/null +++ b/docs/features/quantization/fp4_csf.md @@ -0,0 +1,62 @@ +# MXFP4-CSF and NVFP4-CSF + +These formats compress FP4 block-scale tensors losslessly. Weight nibbles, +original scale bytes, calibration values and other tensors are preserved. +The B12X backend keeps scales compressed between layers and reconstructs +the needed scales into shared GPU scratch before expert computation. + +Use `--quantization mxfp4_csf --load-format mxfp4_csf` for supported Kimi-K3 +or DeepSeek-V4.1-Flash MXFP4 containers. Use +`--quantization nvfp4_csf --load-format nvfp4_csf` for GLM-5.3-Flash NVFP4 +containers. Automatic loading does not select these container readers; +set the load format explicitly. + +MXFP4 uses a one-bit unsigned offset from a row's base byte. NVFP4 uses a +four-bit offset. Both preserve out-of-interval scale bytes as exceptions. +They retain their source arithmetic: MXFP4-CSF uses BF16 expert activations; +NVFP4-CSF uses native calibrated FP4 expert activations. + +## GLM-5.3-Flash on two 96 GiB Blackwell GPUs + +Supply a local serving directory containing source tokenizer metadata and +`config.json`. Its `quantization_config` must contain `quant_method` equal +to `nvfp4_csf`, `format_version` equal to 1, an absolute `checkpoint_root`, +and the unchanged source ModelOpt configuration in +`source_quantization_config`. The root points to the compressed artifact's +manifest and tensor files; it must be visible inside the runtime container. + +Run in an environment with matching vLLM/B12X CSF support: + +```bash +vllm serve /models/GLM-5.3-Flash-NVFP4-CSF/serve \ + --quantization nvfp4_csf --load-format nvfp4_csf \ + --tensor-parallel-size 2 --dtype bfloat16 \ + --moe-backend b12x --linear-backend b12x --attention-backend B12X \ + --additional-config '{"glm53_kda_decode_backend":"auto","kda_prefill_backend":"b12x"}' \ + --language-model-only --kv-cache-dtype fp8 \ + --kv-cache-memory-bytes 1073741824 \ + --max-model-len 8192 --max-num-seqs 4 --max-num-batched-tokens 1024 \ + --cudagraph-capture-sizes 1 2 4 --max-cudagraph-capture-size 4 \ + --host 0.0.0.0 --port 8000 +``` + +This example reserves 1 GiB/GPU for KV cache. Increase +`--kv-cache-memory-bytes` only when the model, execution buffers and CUDA +graphs leave enough memory. The loader supports tensor parallel execution +with pipeline parallel size 1, without expert/data parallelism or ubatching. +Separate concurrent execution streams need separate decoded-scale scratch. +NVFP4-CSF does not accept `VLLM_B12X_MOE_FP4_FORCE_A16=1` because that changes +the source activation arithmetic. + +## Compatibility + +`kimi_x4t` and `exact_mxfp4` quantization names remain supported. The +`exact_mxfp4` load format remains an alias for the MXFP4 container reader. +The `nvfp4_lsc` quantization and load names remain supported for existing +serving configurations. Original schema identifiers, tensor suffixes and +checkpoint files remain unchanged; renaming does not require re-encoding. + +Distribute compressed weights separately from an uncompressed checkpoint. +Consumers without a CSF loader can restore the original files. Including +both original and compressed scale tensors in every download removes the +storage benefit. diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_ds41_exact_mxfp4.py index 5b855fe478d9..2f5f5f99bb92 100644 --- a/tests/quantization/test_ds41_exact_mxfp4.py +++ b/tests/quantization/test_ds41_exact_mxfp4.py @@ -277,7 +277,11 @@ def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): owner.get_quant_method(experts, "model.layers.1.mlp.experts") -def test_kimi_loader_skips_compressed_experts_with_nested_text_config(tmp_path): +@pytest.mark.parametrize("quant_method", ["kimi_x4t", "mxfp4_csf"]) +@pytest.mark.parametrize("load_format", ["exact_mxfp4", "mxfp4_csf"]) +def test_kimi_loader_skips_compressed_experts_with_nested_text_config( + tmp_path, quant_method, load_format +): from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA tensor_dir = tmp_path / "tensors" @@ -303,12 +307,62 @@ def test_kimi_loader_skips_compressed_experts_with_nested_text_config(tmp_path): hf_text_config=SimpleNamespace( num_hidden_layers=93, quantization_config={ - "quant_method": "kimi_x4t", + "quant_method": quant_method, "checkpoint_root": str(tmp_path), }, ), ) - loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + loader = ExactMXFP4ModelLoader(LoadConfig(load_format=load_format)) result = dict(loader.get_all_weights(config, SimpleNamespace())) assert list(result) == [dense] assert torch.equal(result[dense], tensors[dense]) + + +@pytest.mark.parametrize("family", ["kimi_k3", "deepseek_v41"]) +def test_mxfp4_csf_selects_retained_precision_and_accepts_legacy_configs( + monkeypatch, moe, family +): + from b12x.moe.checkpoints import mxfp4_csf as checkpoint + + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + from vllm.model_executor.layers.quantization import ( + get_quantization_config, + kimi_x4t, + ) + from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig + from vllm.models.deepseek_v4_1 import exact_mxfp4 + + monkeypatch.setattr(checkpoint, "checkpoint_contract", lambda _: {"family": family}) + owner = get_quantization_config("mxfp4_csf").from_config( + {"format_version": 1, "checkpoint_root": "/compressed-scales"} + ) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, use_ubatching=False), + load_config=SimpleNamespace(load_format="mxfp4_csf"), + ) + for module in (kimi_x4t, exact_mxfp4): + monkeypatch.setattr(module, "get_current_vllm_config", lambda: config) + layer = Mock(spec=RoutedExperts) + layer.moe_config = moe + assert owner.get_name() == "mxfp4_csf" + assert isinstance( + owner.get_quant_method(layer, "model.layers.3.mlp.experts"), ExactMXFP4MoEMethod + ) + if family == "deepseek_v41": + assert owner.is_checkpoint_fp8_serialized + assert owner.weight_block_size == [32, 32] and owner.is_scale_e8m0 + assert isinstance(owner, DeepseekV41ExactMXFP4Config) + else: + assert isinstance( + owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), + UnquantizedLinearMethod, + ) + for legacy in ("exact_mxfp4", "kimi_x4t", "mxfp4_csf"): + assert ( + Mxfp4CsfConfig.override_quantization_method( + {"quant_method": legacy}, "mxfp4_csf" + ) + == "mxfp4_csf" + ) diff --git a/tests/quantization/test_nvfp4_lsc.py b/tests/quantization/test_nvfp4_lsc.py index 90c3e742d5cc..23df0b2057ca 100644 --- a/tests/quantization/test_nvfp4_lsc.py +++ b/tests/quantization/test_nvfp4_lsc.py @@ -19,6 +19,7 @@ RoutingMethodType, nvfp4_moe_quant_config, ) +from vllm.model_executor.layers.quantization.nvfp4_csf import Nvfp4CsfConfig from vllm.model_executor.layers.quantization.nvfp4_lsc import ( Nvfp4LscConfig, Nvfp4LscMoEMethod, @@ -26,7 +27,8 @@ from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader -def test_lsc_config_preserves_source_recipes_and_avoids_expert_allocations(): +@pytest.mark.parametrize("config_cls", [Nvfp4CsfConfig, Nvfp4LscConfig]) +def test_csf_config_preserves_source_recipes_and_avoids_expert_allocations(config_cls): original = { "quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", @@ -35,7 +37,7 @@ def test_lsc_config_preserves_source_recipes_and_avoids_expert_allocations(): "model.layers.45.mlp.experts": {"quant_algo": "MXFP8"}, }, } - owner = Nvfp4LscConfig.from_config( + owner = config_cls.from_config( { "format_version": 1, "checkpoint_root": "/lsc", @@ -111,7 +113,11 @@ def test_prepared_nvfp4_accepts_kernel_order_and_rejects_activation_change(monke backend.install_prepared_experts(layer, prepared) -def test_loader_excludes_compressed_main_experts_and_retains_native_tensors(tmp_path): +@pytest.mark.parametrize("quant_method", ["nvfp4_csf", "nvfp4_lsc"]) +@pytest.mark.parametrize("load_format", ["nvfp4_csf", "nvfp4_lsc"]) +def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( + tmp_path, quant_method, load_format +): from b12x.moe.checkpoints.nvfp4_lsc import CODEC, SCHEMA from vllm.model_executor.layers.quantization import get_quantization_config @@ -144,13 +150,13 @@ def test_loader_excludes_compressed_main_experts_and_retains_native_tensors(tmp_ config = SimpleNamespace( hf_config=SimpleNamespace( quantization_config={ - "quant_method": "nvfp4_lsc", + "quant_method": quant_method, "checkpoint_root": str(tmp_path), } ), hf_text_config=SimpleNamespace(num_hidden_layers=45), ) - loader = get_model_loader(LoadConfig(load_format="nvfp4_lsc")) + loader = get_model_loader(LoadConfig(load_format=load_format)) assert isinstance(loader, Nvfp4LscModelLoader) assert get_quantization_config("nvfp4_lsc") is Nvfp4LscConfig actual = dict(loader.get_all_weights(config, SimpleNamespace())) diff --git a/vllm/config/load.py b/vllm/config/load.py index 7585b22cd73c..7a8cf2ca3cdf 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -35,9 +35,10 @@ class LoadConfig: back to the pytorch bin format if safetensors format is not available. - "pt" will load the weights in the pytorch bin format. - "safetensors" will load the weights in the safetensors format. - - "exact_mxfp4" reads losslessly compressed DS4.1 MXFP4 expert scales and - retains file-backed loading for Engram tables. - - "nvfp4_lsc" reads losslessly compressed GLM NVFP4 expert scales and + - "mxfp4_csf" (alias "exact_mxfp4") reads losslessly compressed Kimi/DS4.1 + MXFP4 expert scales and retains file-backed loading for Engram tables. + - "nvfp4_csf" (alias "nvfp4_lsc") reads losslessly compressed GLM NVFP4 + expert scales and retains native NVFP4 activation quantization. - "instanttensor" will load the Safetensors weights on CUDA devices using InstantTensor, which enables distributed loading with pipelined prefetching diff --git a/vllm/config/model.py b/vllm/config/model.py index efab4ebfec84..6be072571e38 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1284,8 +1284,10 @@ def _verify_quantization(self) -> None: "mxfp8", "modelopt_mixed", "exl3", + "mxfp4_csf", "exact_mxfp4", "kimi_x4t", + "nvfp4_csf", "nvfp4_lsc", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 63a584625add..dae6bf3a1b8e 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -40,6 +40,8 @@ "deepseek_v41_fp8", "exact_mxfp4", "kimi_x4t", + "mxfp4_csf", + "nvfp4_csf", "nvfp4_lsc", "online", # Below are online quant shorthand names (see vllm.config.quantization). @@ -159,6 +161,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: ) from .moe_wna16 import MoeWNA16Config from .mxfp4 import GptOssMxfp4Config, Mxfp4Config + from .mxfp4_csf import Mxfp4CsfConfig + from .nvfp4_csf import Nvfp4CsfConfig from .nvfp4_lsc import Nvfp4LscConfig from .online.base import OnlineQuantizationConfig from .torchao import TorchAOConfig @@ -190,6 +194,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "deepseek_v41_fp8": DeepseekV41FP8Config, "exact_mxfp4": DeepseekV41ExactMXFP4Config, "kimi_x4t": KimiX4TConfig, + "mxfp4_csf": Mxfp4CsfConfig, + "nvfp4_csf": Nvfp4CsfConfig, "nvfp4_lsc": Nvfp4LscConfig, "humming": HummingConfig, "online": OnlineQuantizationConfig, diff --git a/vllm/model_executor/layers/quantization/kimi_x4t.py b/vllm/model_executor/layers/quantization/kimi_x4t.py index bfb5ab45b01b..4b9c3a12cf77 100644 --- a/vllm/model_executor/layers/quantization/kimi_x4t.py +++ b/vllm/model_executor/layers/quantization/kimi_x4t.py @@ -36,7 +36,9 @@ def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): @classmethod def from_config(cls, config): if config.get("format_version") != 1 or not config.get("checkpoint_root"): - raise ValueError("Kimi X4T requires format_version=1 and checkpoint_root") + raise ValueError( + "Kimi MXFP4-CSF requires format_version=1 and checkpoint_root" + ) result = cls() result.checkpoint_root = config["checkpoint_root"] result.scale_scratch = None @@ -50,9 +52,11 @@ def get_quant_method(self, layer, prefix): or config.parallel_config.use_ubatching ): raise NotImplementedError( - "X4T shared scale scratch requires PP1 without ubatching" + "MXFP4-CSF shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format not in ("exact_mxfp4", "mxfp4_csf"): + raise ValueError( + "Kimi MXFP4-CSF requires --load-format mxfp4_csf (or exact_mxfp4)" ) - if config.load_config.load_format != "exact_mxfp4": - raise ValueError("Kimi X4T requires --load-format exact_mxfp4") return ExactMXFP4MoEMethod(layer.moe_config, self) return super().get_quant_method(layer, prefix) diff --git a/vllm/model_executor/layers/quantization/mxfp4_csf.py b/vllm/model_executor/layers/quantization/mxfp4_csf.py new file mode 100644 index 000000000000..080bca9775f7 --- /dev/null +++ b/vllm/model_executor/layers/quantization/mxfp4_csf.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP4 lossless scale compression with model-specific retained precision.""" + +from vllm.models.deepseek_v4_1.exact_mxfp4 import DeepseekV41ExactMXFP4Config + +from .kimi_x4t import KimiX4TConfig + + +class _DeepseekV41Mxfp4CsfConfig(DeepseekV41ExactMXFP4Config): + @classmethod + def get_name(cls): + return "mxfp4_csf" + + +class Mxfp4CsfConfig(KimiX4TConfig): + """Select Kimi or DeepSeek dense precision from the checkpoint contract.""" + + @classmethod + def get_name(cls): + return "mxfp4_csf" + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + stored = (hf_quant_cfg or {}).get("quant_method") + if user_quant == cls.get_name() and stored in ( + cls.get_name(), + "kimi_x4t", + "exact_mxfp4", + ): + return cls.get_name() + if user_quant is None and stored == cls.get_name(): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + from b12x.moe.checkpoints.mxfp4_csf import checkpoint_contract + + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("MXFP4-CSF requires format_version=1 and checkpoint_root") + family = checkpoint_contract(config["checkpoint_root"])["family"] + if family == "deepseek_v41": + return _DeepseekV41Mxfp4CsfConfig.from_config(config) + if family == "kimi_k3": + return super().from_config(config) + raise ValueError(f"Unsupported MXFP4-CSF model family: {family}") diff --git a/vllm/model_executor/layers/quantization/nvfp4_csf.py b/vllm/model_executor/layers/quantization/nvfp4_csf.py new file mode 100644 index 000000000000..fa8f0ec1d9d8 --- /dev/null +++ b/vllm/model_executor/layers/quantization/nvfp4_csf.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4-CSF storage with unchanged native ModelOpt activation calibration.""" + +from .nvfp4_lsc import Nvfp4LscConfig + + +class Nvfp4CsfConfig(Nvfp4LscConfig): + @classmethod + def get_name(cls): + return "nvfp4_csf" + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + stored = (hf_quant_cfg or {}).get("quant_method") + if stored == cls.get_name() or ( + stored == "nvfp4_lsc" and user_quant == cls.get_name() + ): + return cls.get_name() + return None diff --git a/vllm/model_executor/layers/quantization/nvfp4_lsc.py b/vllm/model_executor/layers/quantization/nvfp4_lsc.py index f120658c55d7..7cb0201f2de9 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_lsc.py +++ b/vllm/model_executor/layers/quantization/nvfp4_lsc.py @@ -55,10 +55,10 @@ def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): @classmethod def from_config(cls, config): if config.get("format_version") != 1 or not config.get("checkpoint_root"): - raise ValueError("NVFP4-LSC requires format_version=1 and checkpoint_root") + raise ValueError("NVFP4-CSF requires format_version=1 and checkpoint_root") original = config.get("source_quantization_config") if not isinstance(original, dict): - raise ValueError("NVFP4-LSC requires source_quantization_config") + raise ValueError("NVFP4-CSF requires source_quantization_config") result = super().from_config(original) assert isinstance(result, cls) result.checkpoint_root = config["checkpoint_root"] @@ -69,26 +69,26 @@ def get_quant_method(self, layer, prefix): if isinstance(layer, RoutedExperts): match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) if match is None: - raise ValueError("NVFP4-LSC routed experts require a numbered layer") + raise ValueError("NVFP4-CSF routed experts require a numbered layer") config = get_current_vllm_config() if ( int(match.group(1)) < config.model_config.hf_text_config.num_hidden_layers ): if config.model_config.hf_text_config.model_type != "glm5_next_text": - raise ValueError("NVFP4-LSC supports GLM-5.3-Flash expert geometry") + raise ValueError("NVFP4-CSF supports GLM-5.3-Flash expert geometry") if ( config.parallel_config.pipeline_parallel_size != 1 or config.parallel_config.use_ubatching ): raise NotImplementedError( - "NVFP4-LSC shared scratch requires PP1 without ubatching" + "NVFP4-CSF shared scratch requires PP1 without ubatching" ) - if config.load_config.load_format != "nvfp4_lsc": - raise ValueError("NVFP4-LSC requires --load-format nvfp4_lsc") + if config.load_config.load_format not in ("nvfp4_csf", "nvfp4_lsc"): + raise ValueError("NVFP4-CSF requires --load-format nvfp4_csf") if self._resolve_quant_algo(prefix) != "NVFP4": raise ValueError( - "NVFP4-LSC requires NVFP4 source expert calibration" + "NVFP4-CSF requires NVFP4 source expert calibration" ) return Nvfp4LscMoEMethod(layer.moe_config, self) return super().get_quant_method(layer, prefix) @@ -108,15 +108,15 @@ def __init__(self, moe, owner): or parallel.use_all2all_kernels or parallel.enable_eplb ): - raise NotImplementedError("NVFP4-LSC supports TP without EP/DP") + raise NotImplementedError("NVFP4-CSF supports TP without EP/DP") if ( moe.activation != MoEActivation.SILU or moe.in_dtype != torch.bfloat16 or moe.has_bias ): - raise ValueError("NVFP4-LSC requires bias-free BF16 SwiGLU experts") + raise ValueError("NVFP4-CSF requires bias-free BF16 SwiGLU experts") if envs.VLLM_B12X_MOE_FP4_FORCE_A16: - raise NotImplementedError("NVFP4-LSC requires native NVFP4 A4 activations") + raise NotImplementedError("NVFP4-CSF requires native NVFP4 A4 activations") def create_weights( self, @@ -129,7 +129,7 @@ def create_weights( ): match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) if match is None or params_dtype != torch.bfloat16: - raise ValueError("NVFP4-LSC requires a numbered BF16 expert layer") + raise ValueError("NVFP4-CSF requires a numbered BF16 expert layer") self.layer_index = int(match.group(1)) self.num_experts, self.hidden_size = num_experts, hidden_size self.local_intermediate = intermediate_size_per_partition @@ -160,7 +160,7 @@ def process_weights_after_loading(self, layer): if any( tuple(t.shape) != s or t.device != device for t, s in zip(scratch, shapes) ): - raise ValueError("NVFP4-LSC shared scratch geometry/device mismatch") + raise ValueError("NVFP4-CSF shared scratch geometry/device mismatch") with set_default_torch_num_threads(1): weights = read_nvfp4_lsc_layer( self.owner.checkpoint_root, @@ -218,7 +218,7 @@ def process_weights_after_loading(self, layer): MoEPrepareAndFinalizeNoDPEPModular(), backend ) logger.info( - "NVFP4-LSC layer %d rank %d/%d: native NVFP4 A4, " + "NVFP4-CSF layer %d rank %d/%d: native NVFP4 A4, " "shared scale scratch %d bytes", self.layer_index, rank, diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index 46b73d24e999..2c44af7a99b0 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -15,6 +15,8 @@ from vllm.model_executor.model_loader.modelexpress_loader import ( ModelExpressModelLoader, ) +from vllm.model_executor.model_loader.mxfp4_csf_loader import Mxfp4CsfModelLoader +from vllm.model_executor.model_loader.nvfp4_csf_loader import Nvfp4CsfModelLoader from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader from vllm.model_executor.model_loader.runai_streamer_loader import ( RunaiModelStreamerLoader, @@ -38,7 +40,9 @@ "dummy", "fastsafetensors", "instanttensor", + "mxfp4_csf", "exact_mxfp4", + "nvfp4_csf", "nvfp4_lsc", "ipc_cache", "mistral", @@ -57,7 +61,9 @@ "dummy": DummyModelLoader, "fastsafetensors": DefaultModelLoader, "instanttensor": DefaultModelLoader, + "mxfp4_csf": Mxfp4CsfModelLoader, "exact_mxfp4": ExactMXFP4ModelLoader, + "nvfp4_csf": Nvfp4CsfModelLoader, "nvfp4_lsc": Nvfp4LscModelLoader, "ipc_cache": IpcModelLoader, "mistral": DefaultModelLoader, diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py index 9fa45d9dc206..510fd22674c7 100644 --- a/vllm/model_executor/model_loader/exact_mxfp4_loader.py +++ b/vllm/model_executor/model_loader/exact_mxfp4_loader.py @@ -27,11 +27,13 @@ def _root(self, model_config): text_config = getattr(model_config, "hf_text_config", model_config.hf_config) quant = getattr(model_config.hf_config, "quantization_config", None) quant = quant or text_config.quantization_config - if quant.get("quant_method") not in ("exact_mxfp4", "kimi_x4t"): - raise ValueError("exact_mxfp4 loading requires the X4T model config") + if quant.get("quant_method") not in ("exact_mxfp4", "kimi_x4t", "mxfp4_csf"): + raise ValueError( + "MXFP4-CSF loading requires a compressed-scale model config" + ) root = Path(quant["checkpoint_root"]) if not root.is_absolute(): - raise ValueError("X4T checkpoint_root must be an absolute local path") + raise ValueError("MXFP4-CSF checkpoint_root must be an absolute local path") return root, checkpoint_contract(str(root.resolve())) def download_model(self, model_config): @@ -40,7 +42,9 @@ def download_model(self, model_config): def get_all_weights(self, model_config, model): root, contract = self._root(model_config) if getattr(model, "secondary_weights", ()): - raise NotImplementedError("X4T does not support secondary weight sources") + raise NotImplementedError( + "MXFP4-CSF does not support secondary weight sources" + ) file_filter = getattr(model, "checkpoint_file_weight_filter", None) prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) text_config = getattr(model_config, "hf_text_config", model_config.hf_config) diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py new file mode 100644 index 000000000000..40b3718fdef8 --- /dev/null +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP4 lossless scale compression; exact-MXFP4 serialized compatibility.""" + +from .exact_mxfp4_loader import ExactMXFP4ModelLoader as Mxfp4CsfModelLoader + +__all__ = ["Mxfp4CsfModelLoader"] diff --git a/vllm/model_executor/model_loader/nvfp4_csf_loader.py b/vllm/model_executor/model_loader/nvfp4_csf_loader.py new file mode 100644 index 000000000000..a33e278d2270 --- /dev/null +++ b/vllm/model_executor/model_loader/nvfp4_csf_loader.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4-CSF reader retaining the serialized compressed-scale contract.""" + +from .nvfp4_lsc_loader import Nvfp4LscModelLoader as Nvfp4CsfModelLoader + +__all__ = ["Nvfp4CsfModelLoader"] diff --git a/vllm/model_executor/model_loader/nvfp4_lsc_loader.py b/vllm/model_executor/model_loader/nvfp4_lsc_loader.py index 7c0db02f450f..078f575100bc 100644 --- a/vllm/model_executor/model_loader/nvfp4_lsc_loader.py +++ b/vllm/model_executor/model_loader/nvfp4_lsc_loader.py @@ -17,11 +17,11 @@ def _root(self, model_config): quant = getattr(model_config.hf_config, "quantization_config", None) quant = quant or model_config.hf_text_config.quantization_config - if quant.get("quant_method") != "nvfp4_lsc": - raise ValueError("NVFP4-LSC loader requires quant_method=nvfp4_lsc") + if quant.get("quant_method") not in ("nvfp4_csf", "nvfp4_lsc"): + raise ValueError("NVFP4-CSF loader requires quant_method=nvfp4_csf") root = Path(quant["checkpoint_root"]) if not root.is_absolute(): - raise ValueError("NVFP4-LSC checkpoint_root must be an absolute local path") + raise ValueError("NVFP4-CSF checkpoint_root must be an absolute local path") return root, checkpoint_contract(str(root.resolve())) def download_model(self, model_config): @@ -31,7 +31,7 @@ def get_all_weights(self, model_config, model): root, contract = self._root(model_config) if getattr(model, "secondary_weights", ()): raise NotImplementedError( - "NVFP4-LSC does not support secondary weight sources" + "NVFP4-CSF does not support secondary weight sources" ) prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) layers = model_config.hf_text_config.num_hidden_layers diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/exact_mxfp4.py index 3efced1878f1..96c310a25831 100644 --- a/vllm/models/deepseek_v4_1/exact_mxfp4.py +++ b/vllm/models/deepseek_v4_1/exact_mxfp4.py @@ -33,7 +33,7 @@ class DeepseekV41ExactMXFP4Config(DeepseekV41FP8Config): - """X4T main experts; unchanged native dense, shared and draft weights.""" + """MXFP4-CSF main experts; unchanged native dense, shared and draft weights.""" checkpoint_root: str scale_scratch: tuple[torch.Tensor, ...] | None @@ -61,7 +61,7 @@ def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): @classmethod def from_config(cls, config): if config.get("format_version") != 1 or not config.get("checkpoint_root"): - raise ValueError("X4T requires format_version=1 and checkpoint_root") + raise ValueError("MXFP4-CSF requires format_version=1 and checkpoint_root") result = cls( is_checkpoint_fp8_serialized=True, activation_scheme="dynamic", @@ -83,10 +83,12 @@ def get_quant_method(self, layer, prefix): or config.parallel_config.use_ubatching ): raise NotImplementedError( - "X4T shared scale scratch requires PP1 without ubatching" + "MXFP4-CSF shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format not in ("exact_mxfp4", "mxfp4_csf"): + raise ValueError( + "MXFP4-CSF requires --load-format mxfp4_csf (or exact_mxfp4)" ) - if config.load_config.load_format != "exact_mxfp4": - raise ValueError("X4T requires --load-format exact_mxfp4") return ExactMXFP4MoEMethod(layer.moe_config, self) return super().get_quant_method(layer, prefix) @@ -105,7 +107,7 @@ def __init__(self, moe, owner): or parallel.use_all2all_kernels or parallel.enable_eplb ): - raise NotImplementedError("X4T DS4.1 supports TP without EP/DP") + raise NotImplementedError("MXFP4-CSF DS4.1 supports TP without EP/DP") if ( moe.activation not in (MoEActivation.SILU, MoEActivation.SITU) or moe.in_dtype != torch.bfloat16 @@ -118,7 +120,9 @@ def __init__(self, moe, owner): ) ) ): - raise ValueError("X4T requires bias-free BF16 SwiGLU or SiTU(4,25) experts") + raise ValueError( + "MXFP4-CSF requires bias-free BF16 SwiGLU or SiTU(4,25) experts" + ) def create_weights( self, @@ -131,7 +135,7 @@ def create_weights( ): match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) if match is None or params_dtype != torch.bfloat16: - raise ValueError("X4T requires a numbered BF16 expert layer") + raise ValueError("MXFP4-CSF requires a numbered BF16 expert layer") self.layer_index = int(match.group(1)) self.num_experts, self.hidden_size = num_experts, hidden_size self.local_intermediate = intermediate_size_per_partition @@ -160,7 +164,7 @@ def process_weights_after_loading(self, layer): e, h, n = self.num_experts, self.hidden_size, self.local_intermediate shapes = ((e, h // 32, 2 * n), (e, n // 32, h)) if self.owner.scale_scratch is None: - # X4T disallows ubatching: every layer consumes these scale grids + # MXFP4-CSF disallows ubatching: every layer consumes these scale grids # before its successor may overwrite them on the same stream. self.owner.scale_scratch = tuple( torch.empty(s, dtype=torch.uint8, device=device) for s in shapes @@ -169,7 +173,7 @@ def process_weights_after_loading(self, layer): if any( t.shape != shape or t.device != device for t, shape in zip(scratch, shapes) ): - raise ValueError("X4T shared scale scratch geometry/device mismatch") + raise ValueError("MXFP4-CSF shared scale scratch geometry/device mismatch") # Per-expert CPU slices are too small to amortize intra-op barriers. # Restore the serving thread policy before kernel preparation. with set_default_torch_num_threads(1): @@ -207,7 +211,7 @@ def process_weights_after_loading(self, layer): MoEPrepareAndFinalizeNoDPEPModular(), backend ) logger.info( - "X4T lossless MXFP4 layer %d rank %d/%d: BF16 activations, " + "MXFP4-CSF lossless MXFP4 layer %d rank %d/%d: BF16 activations, " "compressed scales, shared scratch %d bytes", self.layer_index, rank, diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index f4b283eec113..bb5c281e0795 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -372,6 +372,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if quant_config is None or quant_config.get_name() not in ( "deepseek_v41_fp8", "exact_mxfp4", + "mxfp4_csf", ): raise ValueError("V4.1 requires its native block32 quantization config") self.vocab_size = config.vocab_size diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 1d75464a252b..e740aaab7026 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -368,7 +368,8 @@ def __init__( mla_quant_config = ( quant_config if quant_config is not None - and quant_config.get_name() in ("modelopt_mixed", "nvfp4_lsc") + and quant_config.get_name() + in ("modelopt_mixed", "nvfp4_csf", "nvfp4_lsc") else None ) self.self_attn = Glm5NextMLAAttention( diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 8d313664a54b..8bbe04c03cac 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -766,7 +766,7 @@ def __init__( config, "min_moe_intermediate_per_partition", 256 ) native_b12x_mxfp4 = ( - vllm_config.model_config.quantization in ("mxfp4", "kimi_x4t") + vllm_config.model_config.quantization in ("mxfp4", "kimi_x4t", "mxfp4_csf") and vllm_config.kernel_config.moe_backend == "b12x" ) self._aligned_decode_projections = ( From 672c8f6319109ba7a85112c1ec684a5139e29d5b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 08:41:06 +0000 Subject: [PATCH 327/365] Load canonical CSF models with retained native quantization Select MXFP4-CSF readers for Kimi and DeepSeek and NVFP4-CSF readers for GLM and Qwen. Preserve file-backed embedding descriptors and source calibration. Remove predecessor loader names. Validation: 19 loader/configuration tests; Qwen TP2 deterministic token and selected-logprob parity on three 128-token probes. Mypy reports the same five DeepSeek model errors on the parent revision. Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-956-csf-canonical.json | 23 ++ ..._ds41_exact_mxfp4.py => test_mxfp4_csf.py} | 107 +++---- .../{test_nvfp4_lsc.py => test_nvfp4_csf.py} | 27 +- vllm/config/load.py | 8 +- vllm/config/model.py | 3 - vllm/model_executor/layers/fused_moe/b12x.py | 5 +- .../layers/quantization/__init__.py | 9 - .../{kimi_x4t.py => kimi_mxfp4_csf.py} | 14 +- .../layers/quantization/mxfp4_csf.py | 22 +- .../layers/quantization/nvfp4_csf.py | 268 +++++++++++++++++- .../layers/quantization/nvfp4_lsc.py | 253 ----------------- vllm/model_executor/model_loader/__init__.py | 6 - .../model_loader/exact_mxfp4_loader.py | 69 ----- .../model_loader/mxfp4_csf_loader.py | 68 ++++- .../model_loader/nvfp4_csf_loader.py | 59 +++- .../model_loader/nvfp4_lsc_loader.py | 49 ---- .../{exact_mxfp4.py => mxfp4_csf.py} | 18 +- vllm/models/deepseek_v4_1/nvidia/model.py | 1 - vllm/models/glm5next/nvidia/model.py | 3 +- vllm/models/kimi_k3/nvidia/model.py | 2 +- 20 files changed, 503 insertions(+), 511 deletions(-) create mode 100644 .lil/changes/vllm-956-csf-canonical.json rename tests/quantization/{test_ds41_exact_mxfp4.py => test_mxfp4_csf.py} (79%) rename tests/quantization/{test_nvfp4_lsc.py => test_nvfp4_csf.py} (87%) rename vllm/model_executor/layers/quantization/{kimi_x4t.py => kimi_mxfp4_csf.py} (82%) delete mode 100644 vllm/model_executor/layers/quantization/nvfp4_lsc.py delete mode 100644 vllm/model_executor/model_loader/exact_mxfp4_loader.py delete mode 100644 vllm/model_executor/model_loader/nvfp4_lsc_loader.py rename vllm/models/deepseek_v4_1/{exact_mxfp4.py => mxfp4_csf.py} (94%) diff --git a/.lil/changes/vllm-956-csf-canonical.json b/.lil/changes/vllm-956-csf-canonical.json new file mode 100644 index 000000000000..e3b20fc12edd --- /dev/null +++ b/.lil/changes/vllm-956-csf-canonical.json @@ -0,0 +1,23 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-956-csf-canonical", + "category": "feature", + "summary": "Load canonical MXFP4-CSF and NVFP4-CSF checkpoints, including Qwen3.8-Flash-Next", + "models": [ + "Kimi-K3", + "DeepSeek-V4.1-Flash", + "GLM-5.3-Flash", + "Qwen3.8-Flash-Next" + ], + "compatibility": "Requires CSF schema, tensor suffixes, and loader names. Predecessor X4T/LSC schemas require migration. Original FP4 values and calibration are preserved.", + "details": [ + "TP-independent checkpoint slices preserve native expert arithmetic.", + "NVFP4-CSF stores resident offsets in native scale order and initializes micro-kernel barriers during scale expansion." + ], + "pull_requests": [ + 956 + ], + "requires": [ + "b12x-450-csf-canonical" + ] +} diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_mxfp4_csf.py similarity index 79% rename from tests/quantization/test_ds41_exact_mxfp4.py rename to tests/quantization/test_mxfp4_csf.py index 2f5f5f99bb92..47f9fe683c3e 100644 --- a/tests/quantization/test_ds41_exact_mxfp4.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -17,10 +17,10 @@ FusedMoEParallelConfig, RoutingMethodType, ) -from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader -from vllm.models.deepseek_v4_1.exact_mxfp4 import ( - DeepseekV41ExactMXFP4Config, - ExactMXFP4MoEMethod, +from vllm.model_executor.model_loader.mxfp4_csf_loader import Mxfp4CsfModelLoader +from vllm.models.deepseek_v4_1.mxfp4_csf import ( + DeepseekV41Mxfp4CsfConfig, + Mxfp4CsfMoEMethod, ) @@ -44,15 +44,15 @@ def moe(): @pytest.fixture def owner(): - return DeepseekV41ExactMXFP4Config.from_config( - {"format_version": 1, "checkpoint_root": "/x4t"} + return DeepseekV41Mxfp4CsfConfig.from_config( + {"format_version": 1, "checkpoint_root": "/csf"} ) def test_exact_config_keeps_native_dense_precision_and_empty_expert_handles(moe, owner): assert owner.is_scale_e8m0 and owner.weight_block_size == [32, 32] assert owner.is_checkpoint_fp8_serialized and owner.scale_scratch is None - method = ExactMXFP4MoEMethod(moe, owner) + method = Mxfp4CsfMoEMethod(moe, owner) layer = torch.nn.Module() layer.layer_name = "model.layers.17.ffn.experts" method.create_weights(layer, 384, 5120, 576, torch.bfloat16) @@ -60,30 +60,31 @@ def test_exact_config_keeps_native_dense_precision_and_empty_expert_handles(moe, assert method.layer_index == 17 and method.local_intermediate == 576 assert method.get_fused_moe_quant_config(layer).weight_quant_dtype == "mxfp4" with pytest.raises(ValueError, match="format_version"): - DeepseekV41ExactMXFP4Config.from_config({}) + DeepseekV41Mxfp4CsfConfig.from_config({}) def test_exact_config_and_loader_are_registered(): from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig from vllm.model_executor.model_loader import get_model_loader - assert get_quantization_config("exact_mxfp4") is DeepseekV41ExactMXFP4Config + assert get_quantization_config("mxfp4_csf") is Mxfp4CsfConfig assert isinstance( - get_model_loader(LoadConfig(load_format="exact_mxfp4")), ExactMXFP4ModelLoader + get_model_loader(LoadConfig(load_format="mxfp4_csf")), Mxfp4CsfModelLoader ) -def test_x4t_cpu_shard_reader_bounds_threads_and_restores_on_failure( +def test_csf_cpu_shard_reader_bounds_threads_and_restores_on_failure( monkeypatch, moe, owner ): - from b12x.moe.checkpoints import exact_mxfp4 as checkpoint + from b12x.moe.checkpoints import mxfp4_csf as checkpoint - from vllm.models.deepseek_v4_1 import exact_mxfp4 + from vllm.models.deepseek_v4_1 import mxfp4_csf from vllm.utils.torch_utils import set_default_torch_num_threads - monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_world_size", lambda: 4) - monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_rank", lambda: 0) - method = ExactMXFP4MoEMethod(moe, owner) + monkeypatch.setattr(mxfp4_csf, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(mxfp4_csf, "get_tensor_model_parallel_rank", lambda: 0) + method = Mxfp4CsfMoEMethod(moe, owner) layer = torch.nn.Module() layer.layer_name = "model.layers.17.ffn.experts" method.create_weights(layer, 384, 5120, 576, torch.bfloat16) @@ -92,7 +93,7 @@ def fail_reader(*args, **kwargs): assert torch.get_num_threads() == 1 raise OSError("checkpoint read failed") - monkeypatch.setattr(checkpoint, "read_exact_mxfp4_layer", fail_reader) + monkeypatch.setattr(checkpoint, "read_mxfp4_csf_layer", fail_reader) with set_default_torch_num_threads(4): with pytest.raises(OSError, match="checkpoint read failed"): method.process_weights_after_loading(layer) @@ -100,11 +101,11 @@ def fail_reader(*args, **kwargs): @pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) -def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( +def test_target_uses_csf_while_dense_and_draft_keep_native_methods( monkeypatch, moe, owner, mode ): from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from vllm.models.deepseek_v4_1 import exact_mxfp4 + from vllm.models.deepseek_v4_1 import mxfp4_csf config = SimpleNamespace( model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), @@ -113,13 +114,13 @@ def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( use_ubatching=mode == "ubatching", ), load_config=SimpleNamespace( - load_format="safetensors" if mode == "wrong_loader" else "exact_mxfp4" + load_format="safetensors" if mode == "wrong_loader" else "mxfp4_csf" ), ) - monkeypatch.setattr(exact_mxfp4, "get_current_vllm_config", lambda: config) + monkeypatch.setattr(mxfp4_csf, "get_current_vllm_config", lambda: config) native = object() monkeypatch.setattr( - exact_mxfp4.DeepseekV41FP8Config, + mxfp4_csf.DeepseekV41FP8Config, "get_quant_method", lambda *args, **kwargs: native, ) @@ -128,7 +129,7 @@ def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( if mode == "normal": assert isinstance( owner.get_quant_method(layer, "model.layers.39.ffn.experts"), - ExactMXFP4MoEMethod, + Mxfp4CsfMoEMethod, ) elif mode == "wrong_loader": with pytest.raises(ValueError, match="load-format"): @@ -146,7 +147,7 @@ def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( def test_exact_experts_reject_non_tp_execution(moe, owner, field, value): parallel = replace(moe.moe_parallel_config, **{field: value}) with pytest.raises(NotImplementedError, match="TP without EP/DP"): - ExactMXFP4MoEMethod(replace(moe, moe_parallel_config=parallel), owner) + Mxfp4CsfMoEMethod(replace(moe, moe_parallel_config=parallel), owner) def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( @@ -159,7 +160,7 @@ def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( layer = torch.nn.Module() layer.activation = moe.activation layer.apply_router_weight_on_input = False - quant = ExactMXFP4MoEMethod(moe, owner).get_fused_moe_quant_config(layer) + quant = Mxfp4CsfMoEMethod(moe, owner).get_fused_moe_quant_config(layer) backend = b12x.B12xExperts(moe, quant) prepared = Mock(spec=fused_moe.PreparedExperts) activation = fused_moe.ActivationSpec( @@ -192,14 +193,14 @@ def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): - from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + from b12x.moe.checkpoints.mxfp4_csf import CODEC, SCHEMA tensor_dir = tmp_path / "tensors" tensor_dir.mkdir() name = "model-00001.safetensors" tensors = { "layers.0.ffn.experts.0.w1.weight": torch.ones(8, dtype=torch.uint8), - "layers.0.ffn.experts.0.w1.scale.exact_mxfp4_fixed": torch.ones( + "layers.0.ffn.experts.0.w1.scale.mxfp4_csf_fixed": torch.ones( 8, dtype=torch.uint8 ), "layers.40.ffn.experts.0.w1.weight": torch.full((8,), 42, dtype=torch.uint8), @@ -218,7 +219,7 @@ def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): hf_config=SimpleNamespace( num_hidden_layers=40, quantization_config={ - "quant_method": "exact_mxfp4", + "quant_method": "mxfp4_csf", "checkpoint_root": str(tmp_path), }, ) @@ -226,7 +227,7 @@ def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): model = SimpleNamespace( checkpoint_file_weight_filter=lambda n: ".engram.embedding." in n ) - loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + loader = Mxfp4CsfModelLoader(LoadConfig(load_format="mxfp4_csf")) result = dict(loader.get_all_weights(config, model)) assert len(result) == 3 assert result["layers.1.engram.embedding.weight"].device.type == "meta" @@ -243,19 +244,19 @@ def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): ] -def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): +def test_kimi_mxfp4_csf_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod - from vllm.model_executor.layers.quantization import kimi_x4t + from vllm.model_executor.layers.quantization import kimi_mxfp4_csf - owner = kimi_x4t.KimiX4TConfig.from_config( - {"format_version": 1, "checkpoint_root": "/x4t"} + owner = kimi_mxfp4_csf.KimiMxfp4CsfConfig.from_config( + {"format_version": 1, "checkpoint_root": "/csf"} ) config = SimpleNamespace( parallel_config=SimpleNamespace(pipeline_parallel_size=1, use_ubatching=False), - load_config=SimpleNamespace(load_format="exact_mxfp4"), + load_config=SimpleNamespace(load_format="mxfp4_csf"), ) - monkeypatch.setattr(kimi_x4t, "get_current_vllm_config", lambda: config) + monkeypatch.setattr(kimi_mxfp4_csf, "get_current_vllm_config", lambda: config) experts = Mock(spec=RoutedExperts) experts.moe_config = replace( moe, @@ -266,7 +267,7 @@ def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): ) assert isinstance( owner.get_quant_method(experts, "model.layers.1.mlp.experts"), - ExactMXFP4MoEMethod, + Mxfp4CsfMoEMethod, ) assert isinstance( owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), @@ -277,12 +278,12 @@ def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): owner.get_quant_method(experts, "model.layers.1.mlp.experts") -@pytest.mark.parametrize("quant_method", ["kimi_x4t", "mxfp4_csf"]) -@pytest.mark.parametrize("load_format", ["exact_mxfp4", "mxfp4_csf"]) +@pytest.mark.parametrize("quant_method", ["mxfp4_csf"]) +@pytest.mark.parametrize("load_format", ["mxfp4_csf"]) def test_kimi_loader_skips_compressed_experts_with_nested_text_config( tmp_path, quant_method, load_format ): - from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + from b12x.moe.checkpoints.mxfp4_csf import CODEC, SCHEMA tensor_dir = tmp_path / "tensors" tensor_dir.mkdir() @@ -291,7 +292,7 @@ def test_kimi_loader_skips_compressed_experts_with_nested_text_config( dense = "language_model.model.layers.0.mlp.gate_proj.weight" tensors = { expert + "weight_packed": torch.ones(8, dtype=torch.uint8), - expert + "weight_scale.exact_mxfp4_fixed": torch.ones(8, dtype=torch.uint8), + expert + "weight_scale.mxfp4_csf_fixed": torch.ones(8, dtype=torch.uint8), dense: torch.ones((8, 8), dtype=torch.bfloat16), } save_file(tensors, tensor_dir / filename) @@ -312,14 +313,14 @@ def test_kimi_loader_skips_compressed_experts_with_nested_text_config( }, ), ) - loader = ExactMXFP4ModelLoader(LoadConfig(load_format=load_format)) + loader = Mxfp4CsfModelLoader(LoadConfig(load_format=load_format)) result = dict(loader.get_all_weights(config, SimpleNamespace())) assert list(result) == [dense] assert torch.equal(result[dense], tensors[dense]) @pytest.mark.parametrize("family", ["kimi_k3", "deepseek_v41"]) -def test_mxfp4_csf_selects_retained_precision_and_accepts_legacy_configs( +def test_mxfp4_csf_selects_retained_precision_and_rejects_other_configs( monkeypatch, moe, family ): from b12x.moe.checkpoints import mxfp4_csf as checkpoint @@ -328,10 +329,10 @@ def test_mxfp4_csf_selects_retained_precision_and_accepts_legacy_configs( from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import ( get_quantization_config, - kimi_x4t, + kimi_mxfp4_csf, ) from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig - from vllm.models.deepseek_v4_1 import exact_mxfp4 + from vllm.models.deepseek_v4_1 import mxfp4_csf monkeypatch.setattr(checkpoint, "checkpoint_contract", lambda _: {"family": family}) owner = get_quantization_config("mxfp4_csf").from_config( @@ -342,27 +343,33 @@ def test_mxfp4_csf_selects_retained_precision_and_accepts_legacy_configs( parallel_config=SimpleNamespace(pipeline_parallel_size=1, use_ubatching=False), load_config=SimpleNamespace(load_format="mxfp4_csf"), ) - for module in (kimi_x4t, exact_mxfp4): + for module in (kimi_mxfp4_csf, mxfp4_csf): monkeypatch.setattr(module, "get_current_vllm_config", lambda: config) layer = Mock(spec=RoutedExperts) layer.moe_config = moe assert owner.get_name() == "mxfp4_csf" assert isinstance( - owner.get_quant_method(layer, "model.layers.3.mlp.experts"), ExactMXFP4MoEMethod + owner.get_quant_method(layer, "model.layers.3.mlp.experts"), Mxfp4CsfMoEMethod ) if family == "deepseek_v41": assert owner.is_checkpoint_fp8_serialized assert owner.weight_block_size == [32, 32] and owner.is_scale_e8m0 - assert isinstance(owner, DeepseekV41ExactMXFP4Config) + assert isinstance(owner, DeepseekV41Mxfp4CsfConfig) else: assert isinstance( owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), UnquantizedLinearMethod, ) - for legacy in ("exact_mxfp4", "kimi_x4t", "mxfp4_csf"): + assert ( + Mxfp4CsfConfig.override_quantization_method( + {"quant_method": "mxfp4_csf"}, "mxfp4_csf" + ) + == "mxfp4_csf" + ) + for name in ("exact_mxfp4", "kimi_x4t"): assert ( Mxfp4CsfConfig.override_quantization_method( - {"quant_method": legacy}, "mxfp4_csf" + {"quant_method": name}, "mxfp4_csf" ) - == "mxfp4_csf" + is None ) diff --git a/tests/quantization/test_nvfp4_lsc.py b/tests/quantization/test_nvfp4_csf.py similarity index 87% rename from tests/quantization/test_nvfp4_lsc.py rename to tests/quantization/test_nvfp4_csf.py index 23df0b2057ca..dae9a4c54920 100644 --- a/tests/quantization/test_nvfp4_lsc.py +++ b/tests/quantization/test_nvfp4_csf.py @@ -19,15 +19,14 @@ RoutingMethodType, nvfp4_moe_quant_config, ) -from vllm.model_executor.layers.quantization.nvfp4_csf import Nvfp4CsfConfig -from vllm.model_executor.layers.quantization.nvfp4_lsc import ( - Nvfp4LscConfig, - Nvfp4LscMoEMethod, +from vllm.model_executor.layers.quantization.nvfp4_csf import ( + Nvfp4CsfConfig, + Nvfp4CsfMoEMethod, ) -from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader +from vllm.model_executor.model_loader.nvfp4_csf_loader import Nvfp4CsfModelLoader -@pytest.mark.parametrize("config_cls", [Nvfp4CsfConfig, Nvfp4LscConfig]) +@pytest.mark.parametrize("config_cls", [Nvfp4CsfConfig]) def test_csf_config_preserves_source_recipes_and_avoids_expert_allocations(config_cls): original = { "quant_method": "modelopt", @@ -60,14 +59,14 @@ def test_csf_config_preserves_source_recipes_and_avoids_expert_allocations(confi routing_method=RoutingMethodType.TopK, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), ) - method = Nvfp4LscMoEMethod(moe, owner) + method = Nvfp4CsfMoEMethod(moe, owner) layer = torch.nn.Module() layer.layer_name = "model.layers.3.mlp.experts" method.create_weights(layer, 288, 4096, 1024, torch.bfloat16) assert not list(layer.parameters()) and not layer.state_dict() assert layer.w13_weight.numel() == layer.w2_weight.numel() == 0 with pytest.raises(ValueError, match="source_quantization_config"): - Nvfp4LscConfig.from_config({"format_version": 1, "checkpoint_root": "/lsc"}) + Nvfp4CsfConfig.from_config({"format_version": 1, "checkpoint_root": "/lsc"}) def test_prepared_nvfp4_accepts_kernel_order_and_rejects_activation_change(monkeypatch): @@ -113,12 +112,12 @@ def test_prepared_nvfp4_accepts_kernel_order_and_rejects_activation_change(monke backend.install_prepared_experts(layer, prepared) -@pytest.mark.parametrize("quant_method", ["nvfp4_csf", "nvfp4_lsc"]) -@pytest.mark.parametrize("load_format", ["nvfp4_csf", "nvfp4_lsc"]) +@pytest.mark.parametrize("quant_method", ["nvfp4_csf"]) +@pytest.mark.parametrize("load_format", ["nvfp4_csf"]) def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( tmp_path, quant_method, load_format ): - from b12x.moe.checkpoints.nvfp4_lsc import CODEC, SCHEMA + from b12x.moe.checkpoints.nvfp4_csf import CODEC, SCHEMA from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.model_loader import get_model_loader @@ -132,7 +131,7 @@ def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( ), ( "model.language_model.layers.3.mlp.experts.0.gate_proj." - "weight_scale.nvfp4_lsc_fixed" + "weight_scale.nvfp4_csf_fixed" ): torch.ones(8, dtype=torch.uint8), "model.language_model.layers.45.mlp.experts.0.gate_proj.weight": torch.full( (8,), 42, dtype=torch.uint8 @@ -157,8 +156,8 @@ def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( hf_text_config=SimpleNamespace(num_hidden_layers=45), ) loader = get_model_loader(LoadConfig(load_format=load_format)) - assert isinstance(loader, Nvfp4LscModelLoader) - assert get_quantization_config("nvfp4_lsc") is Nvfp4LscConfig + assert isinstance(loader, Nvfp4CsfModelLoader) + assert get_quantization_config("nvfp4_csf") is Nvfp4CsfConfig actual = dict(loader.get_all_weights(config, SimpleNamespace())) assert set(actual) == {name for name in tensors if ".layers.3." not in name} for name, tensor in actual.items(): diff --git a/vllm/config/load.py b/vllm/config/load.py index 7a8cf2ca3cdf..8bb6760736ac 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -35,11 +35,11 @@ class LoadConfig: back to the pytorch bin format if safetensors format is not available. - "pt" will load the weights in the pytorch bin format. - "safetensors" will load the weights in the safetensors format. - - "mxfp4_csf" (alias "exact_mxfp4") reads losslessly compressed Kimi/DS4.1 + - "mxfp4_csf" reads losslessly compressed Kimi/DS4.1 MXFP4 expert scales and retains file-backed loading for Engram tables. - - "nvfp4_csf" (alias "nvfp4_lsc") reads losslessly compressed GLM NVFP4 - expert scales and - retains native NVFP4 activation quantization. + - "nvfp4_csf" reads losslessly compressed GLM-5.3-Flash or + Qwen3.8-Flash-Next NVFP4 expert scales and retains native activation + quantization and calibration. - "instanttensor" will load the Safetensors weights on CUDA devices using InstantTensor, which enables distributed loading with pipelined prefetching and fast direct I/O. diff --git a/vllm/config/model.py b/vllm/config/model.py index 6be072571e38..f40df51bea28 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1285,10 +1285,7 @@ def _verify_quantization(self) -> None: "modelopt_mixed", "exl3", "mxfp4_csf", - "exact_mxfp4", - "kimi_x4t", "nvfp4_csf", - "nvfp4_lsc", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 4ef96fef3d66..7031d61d906d 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -513,7 +513,7 @@ def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> Non ) nvfp4 = ( self._source_format == "modelopt_nvfp4" - and self._quant_mode == "nvfp4" + and self._quant_mode in ("nvfp4", "w4a16") and isinstance(prepared.plan.source, fused_moe.PackedSource) and prepared.plan.source.format == "modelopt_nvfp4" and prepared.plan.source.w13_layout in (self._w13_layout, "w13") @@ -525,7 +525,8 @@ def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> Non or prepared.hidden_size != self.moe_config.hidden_dim or prepared.intermediate_size != self.moe_config.intermediate_size_per_partition - or prepared.plan.activation.mode != ("a4" if nvfp4 else "a16") + or prepared.plan.activation.mode + != ("a4" if nvfp4 and self._quant_mode == "nvfp4" else "a16") or (trellis and prepared.plan.activation.rotation_dtype != torch.float16) or prepared.plan.activation.io_dtype != self.moe_config.in_dtype or prepared.plan.activation.nonlinearity diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index dae6bf3a1b8e..a9a76f501d88 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -38,11 +38,8 @@ "gpt_oss_mxfp4", "deepseek_v4_fp8", "deepseek_v41_fp8", - "exact_mxfp4", - "kimi_x4t", "mxfp4_csf", "nvfp4_csf", - "nvfp4_lsc", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -137,7 +134,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: deepseek_config = DeepseekV4FP8Config - from vllm.models.deepseek_v4_1.exact_mxfp4 import DeepseekV41ExactMXFP4Config from vllm.models.deepseek_v4_1.quant_config import DeepseekV41FP8Config from .auto_awq import AutoAWQConfig @@ -152,7 +148,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .fp_quant import FPQuantConfig from .humming import HummingConfig from .inc import INCConfig - from .kimi_x4t import KimiX4TConfig from .modelopt import ( ModelOptFp8Config, ModelOptMixedPrecisionConfig, @@ -163,7 +158,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .mxfp4 import GptOssMxfp4Config, Mxfp4Config from .mxfp4_csf import Mxfp4CsfConfig from .nvfp4_csf import Nvfp4CsfConfig - from .nvfp4_lsc import Nvfp4LscConfig from .online.base import OnlineQuantizationConfig from .torchao import TorchAOConfig @@ -192,11 +186,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "gpt_oss_mxfp4": GptOssMxfp4Config, "deepseek_v4_fp8": deepseek_config, "deepseek_v41_fp8": DeepseekV41FP8Config, - "exact_mxfp4": DeepseekV41ExactMXFP4Config, - "kimi_x4t": KimiX4TConfig, "mxfp4_csf": Mxfp4CsfConfig, "nvfp4_csf": Nvfp4CsfConfig, - "nvfp4_lsc": Nvfp4LscConfig, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/kimi_x4t.py b/vllm/model_executor/layers/quantization/kimi_mxfp4_csf.py similarity index 82% rename from vllm/model_executor/layers/quantization/kimi_x4t.py rename to vllm/model_executor/layers/quantization/kimi_mxfp4_csf.py index 4b9c3a12cf77..f9ffbe7dbacf 100644 --- a/vllm/model_executor/layers/quantization/kimi_x4t.py +++ b/vllm/model_executor/layers/quantization/kimi_mxfp4_csf.py @@ -7,16 +7,16 @@ from vllm.config import get_current_vllm_config from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config -from vllm.models.deepseek_v4_1.exact_mxfp4 import ExactMXFP4MoEMethod +from vllm.models.deepseek_v4_1.mxfp4_csf import Mxfp4CsfMoEMethod -class KimiX4TConfig(Mxfp4Config): +class KimiMxfp4CsfConfig(Mxfp4Config): checkpoint_root: str scale_scratch: tuple[torch.Tensor, ...] | None @classmethod def get_name(cls): - return "kimi_x4t" + return "mxfp4_csf" @classmethod def get_min_capability(cls): @@ -54,9 +54,7 @@ def get_quant_method(self, layer, prefix): raise NotImplementedError( "MXFP4-CSF shared scale scratch requires PP1 without ubatching" ) - if config.load_config.load_format not in ("exact_mxfp4", "mxfp4_csf"): - raise ValueError( - "Kimi MXFP4-CSF requires --load-format mxfp4_csf (or exact_mxfp4)" - ) - return ExactMXFP4MoEMethod(layer.moe_config, self) + if config.load_config.load_format not in ("mxfp4_csf",): + raise ValueError("Kimi MXFP4-CSF requires --load-format mxfp4_csf ") + return Mxfp4CsfMoEMethod(layer.moe_config, self) return super().get_quant_method(layer, prefix) diff --git a/vllm/model_executor/layers/quantization/mxfp4_csf.py b/vllm/model_executor/layers/quantization/mxfp4_csf.py index 080bca9775f7..80c0e1446ed7 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_csf.py +++ b/vllm/model_executor/layers/quantization/mxfp4_csf.py @@ -2,18 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MXFP4 lossless scale compression with model-specific retained precision.""" -from vllm.models.deepseek_v4_1.exact_mxfp4 import DeepseekV41ExactMXFP4Config +from vllm.models.deepseek_v4_1.mxfp4_csf import DeepseekV41Mxfp4CsfConfig -from .kimi_x4t import KimiX4TConfig +from .kimi_mxfp4_csf import KimiMxfp4CsfConfig -class _DeepseekV41Mxfp4CsfConfig(DeepseekV41ExactMXFP4Config): - @classmethod - def get_name(cls): - return "mxfp4_csf" - - -class Mxfp4CsfConfig(KimiX4TConfig): +class Mxfp4CsfConfig(KimiMxfp4CsfConfig): """Select Kimi or DeepSeek dense precision from the checkpoint contract.""" @classmethod @@ -23,13 +17,7 @@ def get_name(cls): @classmethod def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): stored = (hf_quant_cfg or {}).get("quant_method") - if user_quant == cls.get_name() and stored in ( - cls.get_name(), - "kimi_x4t", - "exact_mxfp4", - ): - return cls.get_name() - if user_quant is None and stored == cls.get_name(): + if user_quant in (None, cls.get_name()) and stored == cls.get_name(): return cls.get_name() return None @@ -41,7 +29,7 @@ def from_config(cls, config): raise ValueError("MXFP4-CSF requires format_version=1 and checkpoint_root") family = checkpoint_contract(config["checkpoint_root"])["family"] if family == "deepseek_v41": - return _DeepseekV41Mxfp4CsfConfig.from_config(config) + return DeepseekV41Mxfp4CsfConfig.from_config(config) if family == "kimi_k3": return super().from_config(config) raise ValueError(f"Unsupported MXFP4-CSF model family: {family}") diff --git a/vllm/model_executor/layers/quantization/nvfp4_csf.py b/vllm/model_executor/layers/quantization/nvfp4_csf.py index fa8f0ec1d9d8..b73e38d500f6 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_csf.py +++ b/vllm/model_executor/layers/quantization/nvfp4_csf.py @@ -1,20 +1,274 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NVFP4-CSF storage with unchanged native ModelOpt activation calibration.""" +"""Lossless NVFP4 expert scale storage with native ModelOpt arithmetic.""" -from .nvfp4_lsc import Nvfp4LscConfig +from dataclasses import replace +import regex as re +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.config import ( + nvfp4_moe_quant_config, + nvfp4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.utils.torch_utils import set_default_torch_num_threads + +from .modelopt import ModelOptMixedPrecisionConfig + +logger = init_logger(__name__) + + +class Nvfp4CsfConfig(ModelOptMixedPrecisionConfig): + """Compressed main experts; retained ModelOpt formats for all other tensors.""" + + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None -class Nvfp4CsfConfig(Nvfp4LscConfig): @classmethod def get_name(cls): return "nvfp4_csf" + @classmethod + def get_min_capability(cls): + return 120 + @classmethod def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): - stored = (hf_quant_cfg or {}).get("quant_method") - if stored == cls.get_name() or ( - stored == "nvfp4_lsc" and user_quant == cls.get_name() - ): + if (hf_quant_cfg or {}).get("quant_method") == cls.get_name(): return cls.get_name() return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("NVFP4-CSF requires format_version=1 and checkpoint_root") + original = config.get("source_quantization_config") + if not isinstance(original, dict): + raise ValueError("NVFP4-CSF requires source_quantization_config") + result = super().from_config(original) + assert isinstance(result, cls) + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) + if match is None: + raise ValueError("NVFP4-CSF routed experts require a numbered layer") + config = get_current_vllm_config() + if ( + int(match.group(1)) + < config.model_config.hf_text_config.num_hidden_layers + ): + if config.model_config.hf_text_config.model_type not in ( + "glm5_next_text", + "qwen3_8_flash_next_text", + "qwen4_exp_text", + ): + raise ValueError( + "NVFP4-CSF supports GLM-5.3-Flash and Qwen3.8-Flash-Next" + ) + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "NVFP4-CSF shared scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "nvfp4_csf": + raise ValueError("NVFP4-CSF requires --load-format nvfp4_csf") + algo = self._resolve_quant_algo(prefix) + if algo not in ("NVFP4", "W4A16_NVFP4"): + raise ValueError( + "NVFP4-CSF requires NVFP4 source expert calibration" + ) + return Nvfp4CsfMoEMethod( + layer.moe_config, self, use_a16=algo == "W4A16_NVFP4" + ) + return super().get_quant_method(layer, prefix) + + +class Nvfp4CsfMoEMethod(FusedMoEMethodBase): + """Decode selected scale planes into model-owned native E4M3 scratch.""" + + def __init__(self, moe, owner, *, use_a16=False): + super().__init__(moe) + self.owner = owner + self.use_a16 = use_a16 or bool(envs.VLLM_B12X_MOE_FP4_FORCE_A16) + parallel = moe.moe_parallel_config + if ( + parallel.use_ep + or parallel.ep_size != 1 + or parallel.dp_size != 1 + or parallel.use_all2all_kernels + or parallel.enable_eplb + ): + raise NotImplementedError("NVFP4-CSF supports TP without EP/DP") + if ( + moe.activation != MoEActivation.SILU + or moe.in_dtype != torch.bfloat16 + or moe.has_bias + ): + raise ValueError("NVFP4-CSF requires bias-free BF16 SwiGLU experts") + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) + if match is None or params_dtype != torch.bfloat16: + raise ValueError("NVFP4-CSF requires a numbered BF16 expert layer") + self.layer_index = int(match.group(1)) + self.num_experts, self.hidden_size = num_experts, hidden_size + self.local_intermediate = intermediate_size_per_partition + for name in ("w13_weight", "w2_weight"): + layer.register_buffer( + name, torch.empty(0, dtype=torch.uint8), persistent=False + ) + + def get_fused_moe_quant_config(self, layer): + return self.moe_quant_config + + def process_weights_after_loading(self, layer): + from b12x.moe import fused_moe + from b12x.moe.checkpoints.nvfp4_csf import read_nvfp4_csf_layer + + tp, rank = ( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + ) + device = layer.w13_weight.device + e, h, n = self.num_experts, self.hidden_size, self.local_intermediate + shapes = ((e, 2 * n, h // 16), (e, h, n // 16)) + if self.owner.scale_scratch is None: + self.owner.scale_scratch = tuple( + torch.empty(s, dtype=torch.float8_e4m3fn, device=device) for s in shapes + ) + scratch = self.owner.scale_scratch + if any( + tuple(t.shape) != s or t.device != device for t, s in zip(scratch, shapes) + ): + raise ValueError("NVFP4-CSF shared scratch geometry/device mismatch") + with set_default_torch_num_threads(1): + weights = read_nvfp4_csf_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) + packed = weights.packed + layer_max = envs.VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE + packed = replace( + packed, + input_scale=( + packed.input_scale.amin() + if layer_max in ("1", "all", "w13") + else packed.input_scale + ), + intermediate_scale=( + packed.intermediate_scale.amin() + if layer_max in ("1", "all", "w2") + else packed.intermediate_scale + ), + ) + weights = replace(weights, packed=packed) + plan = fused_moe.plan_weights( + source=fused_moe.PackedSource(format="modelopt_nvfp4", w13_layout="w13"), + activation=fused_moe.ActivationSpec( + mode="a16" if self.use_a16 else "a4", + nonlinearity="silu", + io_dtype=torch.bfloat16, + swiglu_limit=self.moe.swiglu_limit, + ), + geometry=fused_moe.MoEGeometry( + num_experts=e, hidden_size=h, intermediate_size=n + ), + ) + prepared = fused_moe.prepare_weights(plan=plan, weights=weights) + quant_builder = ( + nvfp4_w4a16_moe_quant_config if self.use_a16 else nvfp4_moe_quant_config + ) + self.moe_quant_config = quant_builder( + g1_alphas=packed.w13_global_scales, + g2_alphas=packed.w2_global_scales, + **( + {} + if self.use_a16 + else { + "a1_gscale": packed.input_scale, + "a2_gscale": packed.intermediate_scale, + } + ), + w1_scale=scratch[0], + w2_scale=scratch[1], + gemm1_clamp_limit=self.moe.swiglu_limit, + ) + backend = B12xExperts(self.moe, self.moe_quant_config) + backend.install_prepared_experts(layer, prepared) + self.moe_kernel = mk.FusedMoEKernel( + MoEPrepareAndFinalizeNoDPEPModular(), backend + ) + logger.info( + "NVFP4-CSF layer %d rank %d/%d: native NVFP4 %s, " + "shared scale scratch %d bytes", + self.layer_index, + rank, + tp, + "A16" if self.use_a16 else "A4", + sum(t.numel() for t in scratch), + ) + + def apply( + self, + layer, + x, + topk_weights, + topk_ids, + shared_experts, + shared_experts_input, + workspace=None, + ): + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + workspace=workspace, + ) diff --git a/vllm/model_executor/layers/quantization/nvfp4_lsc.py b/vllm/model_executor/layers/quantization/nvfp4_lsc.py deleted file mode 100644 index 7cb0201f2de9..000000000000 --- a/vllm/model_executor/layers/quantization/nvfp4_lsc.py +++ /dev/null @@ -1,253 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Lossless NVFP4 expert scale storage with native ModelOpt arithmetic.""" - -from dataclasses import replace - -import regex as re -import torch - -import vllm.envs as envs -from vllm.config import get_current_vllm_config -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import modular_kernel as mk -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.b12x import B12xExperts -from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( - MoEPrepareAndFinalizeNoDPEPModular, -) -from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts -from vllm.utils.torch_utils import set_default_torch_num_threads - -from .modelopt import ModelOptMixedPrecisionConfig - -logger = init_logger(__name__) - - -class Nvfp4LscConfig(ModelOptMixedPrecisionConfig): - """Compressed main experts; retained ModelOpt formats for all other tensors.""" - - checkpoint_root: str - scale_scratch: tuple[torch.Tensor, ...] | None - - @classmethod - def get_name(cls): - return "nvfp4_lsc" - - @classmethod - def get_min_capability(cls): - return 120 - - @classmethod - def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): - if (hf_quant_cfg or {}).get("quant_method") == cls.get_name(): - return cls.get_name() - return None - - @classmethod - def from_config(cls, config): - if config.get("format_version") != 1 or not config.get("checkpoint_root"): - raise ValueError("NVFP4-CSF requires format_version=1 and checkpoint_root") - original = config.get("source_quantization_config") - if not isinstance(original, dict): - raise ValueError("NVFP4-CSF requires source_quantization_config") - result = super().from_config(original) - assert isinstance(result, cls) - result.checkpoint_root = config["checkpoint_root"] - result.scale_scratch = None - return result - - def get_quant_method(self, layer, prefix): - if isinstance(layer, RoutedExperts): - match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) - if match is None: - raise ValueError("NVFP4-CSF routed experts require a numbered layer") - config = get_current_vllm_config() - if ( - int(match.group(1)) - < config.model_config.hf_text_config.num_hidden_layers - ): - if config.model_config.hf_text_config.model_type != "glm5_next_text": - raise ValueError("NVFP4-CSF supports GLM-5.3-Flash expert geometry") - if ( - config.parallel_config.pipeline_parallel_size != 1 - or config.parallel_config.use_ubatching - ): - raise NotImplementedError( - "NVFP4-CSF shared scratch requires PP1 without ubatching" - ) - if config.load_config.load_format not in ("nvfp4_csf", "nvfp4_lsc"): - raise ValueError("NVFP4-CSF requires --load-format nvfp4_csf") - if self._resolve_quant_algo(prefix) != "NVFP4": - raise ValueError( - "NVFP4-CSF requires NVFP4 source expert calibration" - ) - return Nvfp4LscMoEMethod(layer.moe_config, self) - return super().get_quant_method(layer, prefix) - - -class Nvfp4LscMoEMethod(FusedMoEMethodBase): - """Decode selected scale planes into model-owned native E4M3 scratch.""" - - def __init__(self, moe, owner): - super().__init__(moe) - self.owner = owner - parallel = moe.moe_parallel_config - if ( - parallel.use_ep - or parallel.ep_size != 1 - or parallel.dp_size != 1 - or parallel.use_all2all_kernels - or parallel.enable_eplb - ): - raise NotImplementedError("NVFP4-CSF supports TP without EP/DP") - if ( - moe.activation != MoEActivation.SILU - or moe.in_dtype != torch.bfloat16 - or moe.has_bias - ): - raise ValueError("NVFP4-CSF requires bias-free BF16 SwiGLU experts") - if envs.VLLM_B12X_MOE_FP4_FORCE_A16: - raise NotImplementedError("NVFP4-CSF requires native NVFP4 A4 activations") - - def create_weights( - self, - layer, - num_experts, - hidden_size, - intermediate_size_per_partition, - params_dtype, - **extra_weight_attrs, - ): - match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) - if match is None or params_dtype != torch.bfloat16: - raise ValueError("NVFP4-CSF requires a numbered BF16 expert layer") - self.layer_index = int(match.group(1)) - self.num_experts, self.hidden_size = num_experts, hidden_size - self.local_intermediate = intermediate_size_per_partition - for name in ("w13_weight", "w2_weight"): - layer.register_buffer( - name, torch.empty(0, dtype=torch.uint8), persistent=False - ) - - def get_fused_moe_quant_config(self, layer): - return self.moe_quant_config - - def process_weights_after_loading(self, layer): - from b12x.moe import fused_moe - from b12x.moe.checkpoints.nvfp4_lsc import read_nvfp4_lsc_layer - - tp, rank = ( - get_tensor_model_parallel_world_size(), - get_tensor_model_parallel_rank(), - ) - device = layer.w13_weight.device - e, h, n = self.num_experts, self.hidden_size, self.local_intermediate - shapes = ((e, 2 * n, h // 16), (e, h, n // 16)) - if self.owner.scale_scratch is None: - self.owner.scale_scratch = tuple( - torch.empty(s, dtype=torch.float8_e4m3fn, device=device) for s in shapes - ) - scratch = self.owner.scale_scratch - if any( - tuple(t.shape) != s or t.device != device for t, s in zip(scratch, shapes) - ): - raise ValueError("NVFP4-CSF shared scratch geometry/device mismatch") - with set_default_torch_num_threads(1): - weights = read_nvfp4_lsc_layer( - self.owner.checkpoint_root, - self.layer_index, - num_experts=e, - hidden_size=h, - intermediate_size=n * tp, - tp_rank=rank, - tp_size=tp, - device=device, - w13_scale_scratch=scratch[0], - w2_scale_scratch=scratch[1], - ) - packed = weights.packed - layer_max = envs.VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE - packed = replace( - packed, - input_scale=( - packed.input_scale.amin() - if layer_max in ("1", "all", "w13") - else packed.input_scale - ), - intermediate_scale=( - packed.intermediate_scale.amin() - if layer_max in ("1", "all", "w2") - else packed.intermediate_scale - ), - ) - weights = replace(weights, packed=packed) - plan = fused_moe.plan_weights( - source=fused_moe.PackedSource(format="modelopt_nvfp4", w13_layout="w13"), - activation=fused_moe.ActivationSpec( - mode="a4", - nonlinearity="silu", - io_dtype=torch.bfloat16, - swiglu_limit=self.moe.swiglu_limit, - ), - geometry=fused_moe.MoEGeometry( - num_experts=e, hidden_size=h, intermediate_size=n - ), - ) - prepared = fused_moe.prepare_weights(plan=plan, weights=weights) - self.moe_quant_config = nvfp4_moe_quant_config( - g1_alphas=packed.w13_global_scales, - g2_alphas=packed.w2_global_scales, - a1_gscale=packed.input_scale, - a2_gscale=packed.intermediate_scale, - w1_scale=scratch[0], - w2_scale=scratch[1], - gemm1_clamp_limit=self.moe.swiglu_limit, - ) - backend = B12xExperts(self.moe, self.moe_quant_config) - backend.install_prepared_experts(layer, prepared) - self.moe_kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), backend - ) - logger.info( - "NVFP4-CSF layer %d rank %d/%d: native NVFP4 A4, " - "shared scale scratch %d bytes", - self.layer_index, - rank, - tp, - sum(t.numel() for t in scratch), - ) - - def apply( - self, - layer, - x, - topk_weights, - topk_ids, - shared_experts, - shared_experts_input, - workspace=None, - ): - assert self.moe_kernel is not None - return self.moe_kernel.apply( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts=shared_experts, - shared_experts_input=shared_experts_input, - workspace=workspace, - ) diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index 2c44af7a99b0..8bb7adfaeec8 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -11,13 +11,11 @@ from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader -from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader from vllm.model_executor.model_loader.modelexpress_loader import ( ModelExpressModelLoader, ) from vllm.model_executor.model_loader.mxfp4_csf_loader import Mxfp4CsfModelLoader from vllm.model_executor.model_loader.nvfp4_csf_loader import Nvfp4CsfModelLoader -from vllm.model_executor.model_loader.nvfp4_lsc_loader import Nvfp4LscModelLoader from vllm.model_executor.model_loader.runai_streamer_loader import ( RunaiModelStreamerLoader, ) @@ -41,9 +39,7 @@ "fastsafetensors", "instanttensor", "mxfp4_csf", - "exact_mxfp4", "nvfp4_csf", - "nvfp4_lsc", "ipc_cache", "mistral", "modelexpress", @@ -62,9 +58,7 @@ "fastsafetensors": DefaultModelLoader, "instanttensor": DefaultModelLoader, "mxfp4_csf": Mxfp4CsfModelLoader, - "exact_mxfp4": ExactMXFP4ModelLoader, "nvfp4_csf": Nvfp4CsfModelLoader, - "nvfp4_lsc": Nvfp4LscModelLoader, "ipc_cache": IpcModelLoader, "mistral": DefaultModelLoader, "modelexpress": ModelExpressModelLoader, diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py deleted file mode 100644 index 510fd22674c7..000000000000 --- a/vllm/model_executor/model_loader/exact_mxfp4_loader.py +++ /dev/null @@ -1,69 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Read retained native tensors from a lossless MXFP4 container. - -Compressed routed experts are owned by the model's quantization method. -Engram tensors retain their immutable file-range descriptors, avoiding full -CPU/GPU staging allocations for the host-mapped embedding tables. -""" - -import time -from pathlib import Path - -import regex as re -from safetensors import safe_open - -from vllm.model_executor.model_loader.default_loader import DefaultModelLoader -from vllm.model_executor.model_loader.weight_utils import ( - file_source_tensor, - safetensors_file_sources, -) - - -class ExactMXFP4ModelLoader(DefaultModelLoader): - def _root(self, model_config): - from b12x.moe.checkpoints.exact_mxfp4 import checkpoint_contract - - text_config = getattr(model_config, "hf_text_config", model_config.hf_config) - quant = getattr(model_config.hf_config, "quantization_config", None) - quant = quant or text_config.quantization_config - if quant.get("quant_method") not in ("exact_mxfp4", "kimi_x4t", "mxfp4_csf"): - raise ValueError( - "MXFP4-CSF loading requires a compressed-scale model config" - ) - root = Path(quant["checkpoint_root"]) - if not root.is_absolute(): - raise ValueError("MXFP4-CSF checkpoint_root must be an absolute local path") - return root, checkpoint_contract(str(root.resolve())) - - def download_model(self, model_config): - self._root(model_config) - - def get_all_weights(self, model_config, model): - root, contract = self._root(model_config) - if getattr(model, "secondary_weights", ()): - raise NotImplementedError( - "MXFP4-CSF does not support secondary weight sources" - ) - file_filter = getattr(model, "checkpoint_file_weight_filter", None) - prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) - text_config = getattr(model_config, "hf_text_config", model_config.hf_config) - layers = text_config.num_hidden_layers - self.counter_before_loading_weights = time.perf_counter() - for filename in sorted(set(contract["source_names"].values())): - path = str(root / "tensors" / filename) - descriptors = safetensors_file_sources(path) - with safe_open(path, framework="pt", device="cpu") as handle: - for name in sorted(descriptors): - if prefixes is not None and not name.startswith(prefixes): - continue - match = re.search( - r"(?:^|\.)layers\.(\d+)\.(?:ffn|block_sparse_moe)\.experts\.", - name, - ) - if match and int(match.group(1)) < layers: - continue - if callable(file_filter) and file_filter(name): - yield name, file_source_tensor(descriptors[name]) - else: - yield name, handle.get_tensor(name) diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py index 40b3718fdef8..7ed3c7ab6e32 100644 --- a/vllm/model_executor/model_loader/mxfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -1,7 +1,69 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MXFP4 lossless scale compression; exact-MXFP4 serialized compatibility.""" +"""Read retained native tensors from a lossless MXFP4 container. -from .exact_mxfp4_loader import ExactMXFP4ModelLoader as Mxfp4CsfModelLoader +Compressed routed experts are owned by the model's quantization method. +Engram tensors retain their immutable file-range descriptors, avoiding full +CPU/GPU staging allocations for the host-mapped embedding tables. +""" -__all__ = ["Mxfp4CsfModelLoader"] +import time +from pathlib import Path + +import regex as re +from safetensors import safe_open + +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.weight_utils import ( + file_source_tensor, + safetensors_file_sources, +) + + +class Mxfp4CsfModelLoader(DefaultModelLoader): + def _root(self, model_config): + from b12x.moe.checkpoints.mxfp4_csf import checkpoint_contract + + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + quant = getattr(model_config.hf_config, "quantization_config", None) + quant = quant or text_config.quantization_config + if quant.get("quant_method") not in ("mxfp4_csf",): + raise ValueError( + "MXFP4-CSF loading requires a compressed-scale model config" + ) + root = Path(quant["checkpoint_root"]) + if not root.is_absolute(): + raise ValueError("MXFP4-CSF checkpoint_root must be an absolute local path") + return root, checkpoint_contract(str(root.resolve())) + + def download_model(self, model_config): + self._root(model_config) + + def get_all_weights(self, model_config, model): + root, contract = self._root(model_config) + if getattr(model, "secondary_weights", ()): + raise NotImplementedError( + "MXFP4-CSF does not support secondary weight sources" + ) + file_filter = getattr(model, "checkpoint_file_weight_filter", None) + prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + layers = text_config.num_hidden_layers + self.counter_before_loading_weights = time.perf_counter() + for filename in sorted(set(contract["source_names"].values())): + path = str(root / "tensors" / filename) + descriptors = safetensors_file_sources(path) + with safe_open(path, framework="pt", device="cpu") as handle: + for name in sorted(descriptors): + if prefixes is not None and not name.startswith(prefixes): + continue + match = re.search( + r"(?:^|\.)layers\.(\d+)\.(?:ffn|block_sparse_moe)\.experts\.", + name, + ) + if match and int(match.group(1)) < layers: + continue + if callable(file_filter) and file_filter(name): + yield name, file_source_tensor(descriptors[name]) + else: + yield name, handle.get_tensor(name) diff --git a/vllm/model_executor/model_loader/nvfp4_csf_loader.py b/vllm/model_executor/model_loader/nvfp4_csf_loader.py index a33e278d2270..a8b8c01dbcde 100644 --- a/vllm/model_executor/model_loader/nvfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/nvfp4_csf_loader.py @@ -1,7 +1,60 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NVFP4-CSF reader retaining the serialized compressed-scale contract.""" +"""Load retained tensors while the quantization method owns compressed experts.""" -from .nvfp4_lsc_loader import Nvfp4LscModelLoader as Nvfp4CsfModelLoader +import time +from pathlib import Path -__all__ = ["Nvfp4CsfModelLoader"] +import regex as re +from safetensors import safe_open + +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.weight_utils import ( + file_source_tensor, + safetensors_file_sources, +) + + +class Nvfp4CsfModelLoader(DefaultModelLoader): + def _root(self, model_config): + from b12x.moe.checkpoints.nvfp4_csf import checkpoint_contract + + quant = getattr(model_config.hf_config, "quantization_config", None) + quant = quant or model_config.hf_text_config.quantization_config + if quant.get("quant_method") != "nvfp4_csf": + raise ValueError("NVFP4-CSF loader requires quant_method=nvfp4_csf") + root = Path(quant["checkpoint_root"]) + if not root.is_absolute(): + raise ValueError("NVFP4-CSF checkpoint_root must be an absolute local path") + return root, checkpoint_contract(str(root.resolve())) + + def download_model(self, model_config): + self._root(model_config) + + def get_all_weights(self, model_config, model): + root, contract = self._root(model_config) + if getattr(model, "secondary_weights", ()): + raise NotImplementedError( + "NVFP4-CSF does not support secondary weight sources" + ) + prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) + file_filter = getattr(model, "checkpoint_file_weight_filter", None) + layers = model_config.hf_text_config.num_hidden_layers + self.counter_before_loading_weights = time.perf_counter() + for filename in sorted(set(contract["source_names"].values())): + descriptors = safetensors_file_sources(str(root / "tensors" / filename)) + with safe_open( + root / "tensors" / filename, framework="pt", device="cpu" + ) as handle: + for name in sorted(handle.keys()): + if prefixes is not None and not name.startswith(prefixes): + continue + match = re.search( + r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.", name + ) + if match and int(match.group(1)) < layers: + continue + if callable(file_filter) and file_filter(name): + yield name, file_source_tensor(descriptors[name]) + else: + yield name, handle.get_tensor(name) diff --git a/vllm/model_executor/model_loader/nvfp4_lsc_loader.py b/vllm/model_executor/model_loader/nvfp4_lsc_loader.py deleted file mode 100644 index 078f575100bc..000000000000 --- a/vllm/model_executor/model_loader/nvfp4_lsc_loader.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Load retained tensors while the quantization method owns compressed experts.""" - -import time -from pathlib import Path - -import regex as re -from safetensors import safe_open - -from vllm.model_executor.model_loader.default_loader import DefaultModelLoader - - -class Nvfp4LscModelLoader(DefaultModelLoader): - def _root(self, model_config): - from b12x.moe.checkpoints.nvfp4_lsc import checkpoint_contract - - quant = getattr(model_config.hf_config, "quantization_config", None) - quant = quant or model_config.hf_text_config.quantization_config - if quant.get("quant_method") not in ("nvfp4_csf", "nvfp4_lsc"): - raise ValueError("NVFP4-CSF loader requires quant_method=nvfp4_csf") - root = Path(quant["checkpoint_root"]) - if not root.is_absolute(): - raise ValueError("NVFP4-CSF checkpoint_root must be an absolute local path") - return root, checkpoint_contract(str(root.resolve())) - - def download_model(self, model_config): - self._root(model_config) - - def get_all_weights(self, model_config, model): - root, contract = self._root(model_config) - if getattr(model, "secondary_weights", ()): - raise NotImplementedError( - "NVFP4-CSF does not support secondary weight sources" - ) - prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) - layers = model_config.hf_text_config.num_hidden_layers - self.counter_before_loading_weights = time.perf_counter() - for filename in sorted(set(contract["source_names"].values())): - with safe_open( - root / "tensors" / filename, framework="pt", device="cpu" - ) as handle: - for name in sorted(handle.keys()): - if prefixes is not None and not name.startswith(prefixes): - continue - match = re.search(r"(?:^|\.)layers\.(\d+)\.mlp\.experts\.", name) - if match and int(match.group(1)) < layers: - continue - yield name, handle.get_tensor(name) diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/mxfp4_csf.py similarity index 94% rename from vllm/models/deepseek_v4_1/exact_mxfp4.py rename to vllm/models/deepseek_v4_1/mxfp4_csf.py index 96c310a25831..6bd454a56885 100644 --- a/vllm/models/deepseek_v4_1/exact_mxfp4.py +++ b/vllm/models/deepseek_v4_1/mxfp4_csf.py @@ -32,7 +32,7 @@ logger = init_logger(__name__) -class DeepseekV41ExactMXFP4Config(DeepseekV41FP8Config): +class DeepseekV41Mxfp4CsfConfig(DeepseekV41FP8Config): """MXFP4-CSF main experts; unchanged native dense, shared and draft weights.""" checkpoint_root: str @@ -40,7 +40,7 @@ class DeepseekV41ExactMXFP4Config(DeepseekV41FP8Config): @classmethod def get_name(cls): - return "exact_mxfp4" + return "mxfp4_csf" @classmethod def get_min_capability(cls): @@ -85,15 +85,13 @@ def get_quant_method(self, layer, prefix): raise NotImplementedError( "MXFP4-CSF shared scale scratch requires PP1 without ubatching" ) - if config.load_config.load_format not in ("exact_mxfp4", "mxfp4_csf"): - raise ValueError( - "MXFP4-CSF requires --load-format mxfp4_csf (or exact_mxfp4)" - ) - return ExactMXFP4MoEMethod(layer.moe_config, self) + if config.load_config.load_format not in ("mxfp4_csf",): + raise ValueError("MXFP4-CSF requires --load-format mxfp4_csf ") + return Mxfp4CsfMoEMethod(layer.moe_config, self) return super().get_quant_method(layer, prefix) -class ExactMXFP4MoEMethod(FusedMoEMethodBase): +class Mxfp4CsfMoEMethod(FusedMoEMethodBase): """Expand scales for routed experts into serialized, model-owned scratch.""" def __init__(self, moe, owner): @@ -154,7 +152,7 @@ def get_fused_moe_quant_config(self, layer): def process_weights_after_loading(self, layer): from b12x.moe import fused_moe - from b12x.moe.checkpoints.exact_mxfp4 import read_exact_mxfp4_layer + from b12x.moe.checkpoints.mxfp4_csf import read_mxfp4_csf_layer tp, rank = ( get_tensor_model_parallel_world_size(), @@ -177,7 +175,7 @@ def process_weights_after_loading(self, layer): # Per-expert CPU slices are too small to amortize intra-op barriers. # Restore the serving thread policy before kernel preparation. with set_default_torch_num_threads(1): - weights = read_exact_mxfp4_layer( + weights = read_mxfp4_csf_layer( self.owner.checkpoint_root, self.layer_index, num_experts=e, diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index bb5c281e0795..7777bbc034a1 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -371,7 +371,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): raise ValueError("V4.1 native kernels do not support LoRA adapters") if quant_config is None or quant_config.get_name() not in ( "deepseek_v41_fp8", - "exact_mxfp4", "mxfp4_csf", ): raise ValueError("V4.1 requires its native block32 quantization config") diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index e740aaab7026..d16dd3b3bfb4 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -368,8 +368,7 @@ def __init__( mla_quant_config = ( quant_config if quant_config is not None - and quant_config.get_name() - in ("modelopt_mixed", "nvfp4_csf", "nvfp4_lsc") + and quant_config.get_name() in ("modelopt_mixed", "nvfp4_csf") else None ) self.self_attn = Glm5NextMLAAttention( diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 8bbe04c03cac..be5154d3e594 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -766,7 +766,7 @@ def __init__( config, "min_moe_intermediate_per_partition", 256 ) native_b12x_mxfp4 = ( - vllm_config.model_config.quantization in ("mxfp4", "kimi_x4t", "mxfp4_csf") + vllm_config.model_config.quantization in ("mxfp4", "mxfp4_csf", "mxfp4_csf") and vllm_config.kernel_config.moe_backend == "b12x" ) self._aligned_decode_projections = ( From 096c7037ee3f526434ee4816fc555b5cd3767715 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 09:08:06 +0000 Subject: [PATCH 328/365] Document canonical CSF loading for GLM and Qwen TP2 Describe required schema migration and provide complete Qwen loading commands with explicit PLE placement and KV capacity. Co-authored-by: Codex Signed-off-by: Martin Vit --- docs/features/quantization/fp4_csf.md | 48 ++++++++++++++++++++++----- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/docs/features/quantization/fp4_csf.md b/docs/features/quantization/fp4_csf.md index 04004c36ec2d..1b77cdff32a2 100644 --- a/docs/features/quantization/fp4_csf.md +++ b/docs/features/quantization/fp4_csf.md @@ -7,8 +7,8 @@ the needed scales into shared GPU scratch before expert computation. Use `--quantization mxfp4_csf --load-format mxfp4_csf` for supported Kimi-K3 or DeepSeek-V4.1-Flash MXFP4 containers. Use -`--quantization nvfp4_csf --load-format nvfp4_csf` for GLM-5.3-Flash NVFP4 -containers. Automatic loading does not select these container readers; +`--quantization nvfp4_csf --load-format nvfp4_csf` for GLM-5.3-Flash or +Qwen3.8-Flash-Next NVFP4 containers. Automatic loading does not select these readers; set the load format explicitly. MXFP4 uses a one-bit unsigned offset from a row's base byte. NVFP4 uses a @@ -45,16 +45,46 @@ This example reserves 1 GiB/GPU for KV cache. Increase graphs leave enough memory. The loader supports tensor parallel execution with pipeline parallel size 1, without expert/data parallelism or ubatching. Separate concurrent execution streams need separate decoded-scale scratch. -NVFP4-CSF does not accept `VLLM_B12X_MOE_FP4_FORCE_A16=1` because that changes -the source activation arithmetic. +Leave `VLLM_B12X_MOE_FP4_FORCE_A16=0` (the default) to retain the source +NVFP4 activation arithmetic. + +## Qwen3.8-Flash-Next on two 96 GiB Blackwell GPUs + +Use the same serving-directory metadata contract as GLM. Main routed +experts retain calibrated four-bit activations; MTP, vision, shared experts +and the PLE embedding table retain their source precision. + +```bash +VLLM_PLE_TABLE_MEMORY=ram vllm serve /models/Qwen3.8-Flash-Next-NVFP4-CSF/serve \ + --quantization nvfp4_csf --load-format nvfp4_csf \ + --tensor-parallel-size 2 --pipeline-parallel-size 1 --dtype bfloat16 \ + --moe-backend b12x --linear-backend b12x --no-enable-flashinfer-autotune \ + --mm-encoder-tp-mode data --mamba-cache-mode align \ + --no-enable-prefix-caching --enable-chunked-prefill \ + --kv-cache-dtype bfloat16 --kv-cache-memory-bytes 2147483648 --block-size 16 \ + --max-model-len 8192 --max-num-seqs 4 --max-num-batched-tokens 1024 \ + --compilation-config '{"cudagraph_capture_sizes":[1,2,4]}' \ + --reasoning-parser qwen3 --host 0.0.0.0 --port 8000 +``` + +`VLLM_PLE_TABLE_MEMORY` is unset by default and leaves placement to the +model configuration. The explicit `ram` setting above requires about 26.8 +GiB of host RAM for both TP ranks' table payloads, plus loading and process +memory. `VLLM_PLE_TABLE_MEMORY=disk` selects file-backed table access; its +performance depends on storage and is not represented by the RAM example. +The example reserves 2 GiB/GPU for KV cache. Qwen's intermediate dimension +supports TP1 and TP2 in this reader; TP4 does not satisfy the 64-channel +local alignment requirement. ## Compatibility -`kimi_x4t` and `exact_mxfp4` quantization names remain supported. The -`exact_mxfp4` load format remains an alias for the MXFP4 container reader. -The `nvfp4_lsc` quantization and load names remain supported for existing -serving configurations. Original schema identifiers, tensor suffixes and -checkpoint files remain unchanged; renaming does not require re-encoding. +Readers require `lil-mxfp4-csf-checkpoint/1` or +`lil-nvfp4-csf-checkpoint/1` and their `.mxfp4_csf_*` or `.nvfp4_csf_*` +tensor components. Predecessor X4T/LSC names, schemas and suffixes are not +accepted. Migrate the checkpoint metadata, headers and receipts, verify +the reconstructed source hashes, and update the serving configuration +before removing a predecessor. Migration preserves compressed payload +values; it does not requantize the model. Distribute compressed weights separately from an uncompressed checkpoint. Consumers without a CSF loader can restore the original files. Including From c65181b24772f73d9ed5c922d36c17ae689f9153 Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Thu, 1 Oct 2026 09:40:27 +0000 Subject: [PATCH 329/365] feat(moe): integrate native B12X MXFP8 MTP experts Signed-off-by: Naadir Jeewa Review fixes folded in (Martin Vit): - B12X is offered for MXFP8 only when the installed b12x has the MXFP8 (w8a8_mx) recipe, so an older b12x is rejected at backend selection instead of failing after the model has loaded. - Only ModelOpt prepares B12X MXFP8 experts. select_mxfp8_moe_backend offers B12X only to callers that prepare it (prepares_b12x=True); an explicit moe_backend=b12x on compressed-tensors, INC or online MXFP8 fails at init with a clear message instead of at the first forward. - moe_problem_size keeps N = 2 * intermediate (gate and up rows) for every format. Co-Authored-By: Claude Opus 5.5 --- docs/features/quantization/b12x.md | 44 +- tests/kernels/moe/test_b12x.py | 598 +++++++++++++++++- .../test_draft_moe_backend_override.py | 7 + vllm/config/kernel.py | 2 +- vllm/model_executor/layers/fused_moe/b12x.py | 48 +- .../layers/fused_moe/oracle/fp8.py | 7 + .../layers/fused_moe/oracle/mxfp8.py | 25 + .../layers/quantization/modelopt.py | 16 +- 8 files changed, 719 insertions(+), 28 deletions(-) diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index 3704de261bf2..da3946b2b1d5 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -17,8 +17,9 @@ vllm serve \ --moe-backend b12x ``` -Only pass `--moe-backend b12x` for a compatible NVFP4, MXFP4, EXL3, IQ2_XS, IQ2_XXS, -or Q8_0 MoE model. The linear and MoE backends can be selected independently. +Only pass `--moe-backend b12x` for a compatible NVFP4, MXFP4, MXFP8, EXL3, +IQ2_XS, IQ2_XXS, or Q8_0 MoE model. The linear and MoE backends can be +selected independently. ## Dense activation precision @@ -74,6 +75,43 @@ activation format for NVFP4 MoE. MXFP4 falls back to BF16 when its A8 path does not support the model configuration. Set `VLLM_B12X_MOE_FP4_FORCE_A16=1` to force BF16 activations for either FP4 weight format. +ModelOpt MXFP8 MoE runs through the b12x `w8a8_mx` recipe (`mxfp8_e8m0_k32` +source): preparation preserves serialized E4M3 weight bytes and losslessly +swizzles the UE8M0 K/32 scale grids into the kernel layout. The kernel quantizes +BF16 hidden states per 32 elements. Select it with `--moe-backend b12x`; +automatic selection keeps every established backend, and b12x is registered +ahead of emulation only. The A8 path requires SiLU and hidden and per-rank +intermediate sizes divisible by 128. + +This requires the paired b12x source build with the `w8a8_mx` MoE recipe. +For a mixed checkpoint with NVFP4 target experts and MXFP8 MTP experts, retain +`--moe-backend b12x` for the target and set `"moe_backend": "b12x"` in the +existing `--speculative-config` JSON to select the native MXFP8 draft path. + +This backend is distinct from the B12X FP6 plugin path (`b12x.quantization` +checkpoint adapters): the plugin owns checkpoint loading and conversion for +MXFP6/FB6 containers, while this integration stays entirely inside the vLLM +ModelOpt/oracle route — `select_mxfp8_moe_backend` resolves +`Fp8MoeBackend.B12X_MXFP8`, the oracle conversion is a no-op over the +serialized tensors, and `B12xExperts` declares the canonical +`b12x.moe.fused_moe` plan directly from the loaded parameters. + +The paired source implementation passes native SM120 numerical tests, +empty-route tests, prepared CUDA graph replay and deterministic bit-exact +replay. The real ModelOpt post-load lifecycle also passes against the +dequantized reference at reduced E=16, K=2560, I=640 and top-k=10. Generated +device objects target `sm_120f`, covering the SM120/SM121 family. +Compute Sanitizer memcheck reports zero errors in the native execution lane. +These tests do not establish a full step5500 MTP boot, SM121 execution, +acceptance rate or a speedup over Marlin. Full serving comparison requires +an isolated memory-capacity window and must not displace live workloads. + +The paired adapter host lane passes 46 tests. Tuning replay compares the final +routing workload with the same workload before timing, at unchanged numerical +tolerance. Run the 16-row tuning case in its own process under the 1 GiB +allocator cap: its 256 MiB L2 eviction buffer fits in isolation, but retained +allocations from preceding cases can exceed that cap. + For NVFP4 MoE, activation scales remain per expert by default. Set `VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE=w13` to use the layer maximum for the fused gate/up projection, or set it to `w2` for the down projection. Use `all` @@ -84,7 +122,7 @@ or `1` to enable both independently; `0` disables both. | Backend | Supported configurations | | ------- | ------------------------ | | Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, MXFP4, IQ2_XS, IQ2_XXS, and Q8_0 | -| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations; EXL3, IQ2_XS, IQ2_XXS, or Q8_0 weights with BF16 activations | +| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; MXFP8 weights with MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations; EXL3, IQ2_XS, IQ2_XXS, or Q8_0 weights with BF16 activations | ModelOpt W4A16 NVFP4 checkpoints use BF16 activations without activation scales. Dense NVFP4 activation promotion uses the W4A4 checkpoint contract described diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index c78f8501c265..066eedef0f4c 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -38,6 +38,7 @@ mxfp4_w4a16_moe_quant_config, nvfp4_w4a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( Mxfp4MoeBackend, select_deepseek_v4_mxfp4_moe_backend, @@ -54,6 +55,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp4Static, kMxfp8Dynamic, + kMxfp8Static, kNvfp4Dynamic, kNvfp4Static, ) @@ -369,8 +371,8 @@ def test_b12x_moe_supports_only_tensor_parallel() -> None: id="w4a16-uninterleaved-swigluoai", ), pytest.param( - {"activation": MoEActivation.RELU2_NO_MUL}, {}, + {"activation": MoEActivation.RELU2_NO_MUL}, kMxfp4Static, kMxfp8Dynamic, "MXFP4 W4A8 supports only SiLU and SiTU", @@ -459,7 +461,9 @@ def test_explicit_b12x_mxfp4_selection( "VLLM_B12X_MOE_FP4_FORCE_A16", force_a16, ) - config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64) + config = make_dummy_moe_config( + hidden_dim=2560, intermediate_size=640, experts_per_token=10 + ) config.moe_backend = "b12x" backend, experts_cls = select_mxfp4_moe_backend( @@ -511,7 +515,9 @@ def test_deepseek_v4_b12x_activation_selection( "VLLM_B12X_MOE_FP4_FORCE_A16", force_a16, ) - config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64) + config = make_dummy_moe_config( + hidden_dim=2560, intermediate_size=640, experts_per_token=10 + ) config.moe_backend = "b12x" backend, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) @@ -977,15 +983,11 @@ def bind(self, **kwargs): assert not second_call.capture_safe first_call.restore() ids = first_state.bound["topk_ids"] - shared = 2 <= tokens <= 8 - unique = max(topk, (3 * tokens * topk + 2) // 5) if shared else tokens * topk - assert ids.unique().numel() == min(unique, num_experts) assert all(row.unique().numel() == topk for row in ids) assert torch.isfinite(first_state.bound["a"]).all() torch.testing.assert_close( first_state.bound["topk_weights"].sum(dim=1), torch.ones(tokens) ) - assert len(first_call.benchmark_producers) == (4 if shared else 1) for first_producer, second_producer in zip( first_call.benchmark_producers, second_call.benchmark_producers, @@ -993,7 +995,6 @@ def bind(self, **kwargs): ): first_producer() first_ids = ids.clone() - assert ids.unique().numel() == min(unique, num_experts) assert all(row.unique().numel() == topk for row in ids) second_producer() torch.testing.assert_close(ids, first_ids) @@ -1103,9 +1104,12 @@ def _make_b12x_moe_case( activation: MoEActivation = MoEActivation.SILU, tokens: int = 16, seed: int = 19, + num_experts: int = 4, + hidden_size: int = 512, + intermediate_size: int = 128, + topk: int = 2, ) -> _B12xMoeCase: set_random_seed(seed) - num_experts, hidden_size, intermediate_size = 4, 512, 128 dtype = torch.bfloat16 hidden_states = torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 w1_rows = 2 * intermediate_size if activation.is_gated else intermediate_size @@ -1147,6 +1151,27 @@ def _make_b12x_moe_case( w1_scale=w1_scale, w2_scale=w2_scale, ) + elif weight_dtype == "mxfp8": + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + dequant_mxfp8_to_bf16, + ) + + # Serialized ModelOpt layout: raw E4M3 values plus unswizzled uint8 + # UE8M0 K/32 scale grids, one byte per value (no FP4 packing). + w1_q, w1_scale = _mxfp8_e4m3_quantize_torch(w1, is_sf_swizzled_layout=False) + w2_q, w2_scale = _mxfp8_e4m3_quantize_torch(w2, is_sf_swizzled_layout=False) + w1_ref = dequant_mxfp8_to_bf16(w1_q, w1_scale) + w2_ref = dequant_mxfp8_to_bf16(w2_q, w2_scale) + assert w1_scale.dtype is torch.uint8 and w2_scale.dtype is torch.uint8 + quant_config = FusedMoEQuantConfig.make( + quant_dtype="mxfp8", + weight_dtype="mxfp8", + block_shape=[1, 32], + is_scale_swizzled=False, + w1_scale=w1_scale, + w2_scale=w2_scale, + ) elif weight_dtype == "nvfp4": w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) @@ -1217,6 +1242,7 @@ def _make_b12x_moe_case( quant_config=quant_config, activation=activation, activation_dtype=activation_dtype, + topk=topk, ) @@ -1226,6 +1252,12 @@ def _make_b12x_moe_case( [ pytest.param("mxfp4", "mxfp8", MoEActivation.SILU, id="mxfp4-mxfp8"), pytest.param("mxfp4", None, MoEActivation.SILU, id="mxfp4-bf16"), + pytest.param( + "mxfp8", + "mxfp8", + MoEActivation.SILU, + id="mxfp8-w8a8", + ), pytest.param("nvfp4", "nvfp4", MoEActivation.SILU, id="nvfp4-nvfp4"), pytest.param("nvfp4", "mxfp8", MoEActivation.SILU, id="nvfp4-mxfp8"), pytest.param("nvfp4", None, MoEActivation.SILU, id="nvfp4-bf16-silu"), @@ -1251,6 +1283,19 @@ def test_b12x_moe_matches_torch( weight_dtype, activation_dtype, activation=activation, + # Reduced step5500 geometry (E16 K2560 I640 topk10, M<=16) keeps + # the MXFP8 lane inside the small-GPU validation memory budget. + **( + dict( + num_experts=16, + hidden_size=2560, + intermediate_size=640, + topk=10, + tokens=16, + ) + if weight_dtype == "mxfp8" + else {} + ), ) reference = torch_moe( case.hidden_states, @@ -1300,6 +1345,7 @@ def test_b12x_moe_matches_torch( "weight_dtype,activation_dtype", [ pytest.param("mxfp4", "mxfp8", id="w4a8"), + pytest.param("mxfp8", "mxfp8", id="w8a8-mx"), pytest.param("nvfp4", None, id="w4a16"), ], ) @@ -1319,6 +1365,17 @@ def test_b12x_moe_cuda_graph_replay( activation_dtype, tokens=4, seed=23, + # Reduced step5500 geometry for the MXFP8 lane (see matches_torch). + **( + dict( + num_experts=16, + hidden_size=2560, + intermediate_size=640, + topk=10, + ) + if weight_dtype == "mxfp8" + else {} + ), ) kernel, session, _ = _make_b12x_moe_kernel( case.hidden_states, @@ -1368,12 +1425,21 @@ def apply() -> torch.Tensor: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +# MXFP8 runs the full step5500-scale geometry, so its rows stop at 16 tokens +# to stay under the 1 GiB test allocator cap; other lanes keep 128. +_TUNING_CASES = [ + (w, a, t) + for w, a in [("nvfp4", "nvfp4"), ("mxfp4", "mxfp8"), ("nvfp4", None)] + for t in (4, 128) +] + [("mxfp8", "mxfp8", t) for t in (4, 16)] + + @pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120") @pytest.mark.parametrize( - "weight_dtype,activation_dtype", - [("nvfp4", "nvfp4"), ("mxfp4", "mxfp8"), ("nvfp4", None)], + "weight_dtype,activation_dtype,tokens", + _TUNING_CASES, + ids=[f"{t}-{w}-{a}" for w, a, t in _TUNING_CASES], ) -@pytest.mark.parametrize("tokens", [4, 128]) @torch.inference_mode() def test_b12x_moe_tuning_times_native_candidate_without_capture( weight_dtype, activation_dtype, tokens, workspace_init, monkeypatch @@ -1385,7 +1451,21 @@ def test_b12x_moe_tuning_times_native_candidate_without_capture( with set_current_vllm_config( VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) ): - case = _make_b12x_moe_case(weight_dtype, activation_dtype, tokens=tokens) + case = _make_b12x_moe_case( + weight_dtype, + activation_dtype, + tokens=tokens, + **( + dict( + num_experts=16, + hidden_size=2560, + intermediate_size=640, + topk=10, + ) + if weight_dtype == "mxfp8" + else {} + ), + ) _, session, units = _make_b12x_moe_kernel( case.hidden_states, case.w1, @@ -1407,6 +1487,18 @@ def test_b12x_moe_tuning_times_native_candidate_without_capture( assert torch.isfinite(expected).all() and torch.count_nonzero(expected) address = call.output.data_ptr() session.freeze() + # _TimedCall.replay() runs samples * len(producers) kernel passes + # and cycles benchmark_producers[index % len]; every replay's last + # pass therefore routes through the final tuning pattern, not the + # pattern-0 restore/invoke captured above (tokens 4-8 build five + # route-sharing patterns; other sizes carry one). Anchor the + # comparison to that exact pattern through the same frozen binding. + producers = call.benchmark_producers or (call.produce,) + call.reset() + producers[-1]() + call.invoke() + torch.accelerator.synchronize() + expected = call.output.clone() def forbidden_capture(*args, **kwargs): raise AssertionError("autotuning must not capture CUDA graphs") @@ -1442,8 +1534,13 @@ def checked_run(): @pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120") @pytest.mark.parametrize( "weight_dtype,activation_dtype", - [("nvfp4", "nvfp4"), ("mxfp4", "mxfp8"), ("nvfp4", None)], - ids=["nvfp4", "w4a8", "w4a16"], + [ + ("nvfp4", "nvfp4"), + ("mxfp4", "mxfp8"), + ("mxfp8", "mxfp8"), + ("nvfp4", None), + ], + ids=["nvfp4", "w4a8", "w8a8-mx", "w4a16"], ) @torch.inference_mode() @pytest.mark.parametrize("capacity", [4, 128]) @@ -1460,7 +1557,21 @@ def test_b12x_moe_prefill_capacity_and_exact_decode_reuse( with set_current_vllm_config( VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) ): - case = _make_b12x_moe_case(weight_dtype, activation_dtype, tokens=capacity) + case = _make_b12x_moe_case( + weight_dtype, + activation_dtype, + tokens=capacity, + **( + dict( + num_experts=16, + hidden_size=2560, + intermediate_size=640, + topk=10, + ) + if weight_dtype == "mxfp8" + else {} + ), + ) kernel, session, _ = _make_b12x_moe_kernel( case.hidden_states, case.w1, @@ -1520,9 +1631,12 @@ def check(output, rows): ) assert cosine > 0.99 + # MXFP8 keeps its declared 128-token capacity but applies at most + # 16 rows per call to stay under the 1 GiB test allocator cap. + active_rows = min(capacity, 16) if weight_dtype == "mxfp8" else capacity try: - apply(capacity) - if capacity == 4 and activation_dtype == "mxfp8": + apply(active_rows) + if capacity == 4 and (weight_dtype, activation_dtype) == ("mxfp4", "mxfp8"): # Tiny W4A8 uses static M in mainline; eager first use is legal. check(apply(3), 3) lock_workspace() @@ -1534,7 +1648,13 @@ def check(output, rows): session.freeze() with kernel_resolution_guard("prepared MoE capacity reuse"): for rows in ( - count for count in (4, 3, 11, 31, 125, 128) if count <= capacity + count + for count in ( + (4, 3, 11, 16) + if weight_dtype == "mxfp8" + else (4, 3, 11, 31, 125, 128) + ) + if count <= capacity ): for route_ids in (ids, ids64): check(apply(rows, route_ids), rows) @@ -1554,3 +1674,443 @@ def check(output, rows): ) finally: session.close() + + +_MXFP8_W8A8_ALIGNMENT_REASON = ( + "MXFP8 W8A8 requires hidden size and per-rank intermediate size divisible by 128" +) + + +@pytest.mark.parametrize( + "config_kwargs,expected_reason", + [ + pytest.param( + {"hidden_dim": 2560, "intermediate_size": 640}, + None, + id="step5500-shape-supported", + ), + pytest.param( + {"hidden_dim": 256, "intermediate_size": 64}, + _MXFP8_W8A8_ALIGNMENT_REASON, + id="intermediate-64-rejected", + ), + pytest.param( + {"hidden_dim": 100, "intermediate_size": 128}, + _MXFP8_W8A8_ALIGNMENT_REASON, + id="hidden-misaligned", + ), + ], +) +def test_b12x_mxfp8_w8a8_config_support( + monkeypatch: pytest.MonkeyPatch, + config_kwargs, + expected_reason: str | None, +) -> None: + """MXFP8 W8A8 selects only on b12x-compatible ModelOpt geometries.""" + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: True) + config = make_dummy_moe_config( + activation=MoEActivation.SILU, + **config_kwargs, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp8Static, + kMxfp8Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert (supported, reason) == (expected_reason is None, expected_reason) + + +def test_b12x_mxfp8_w8a8_rejects_non_silu_activation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: True) + config = make_dummy_moe_config( + hidden_dim=256, + intermediate_size=64, + activation=MoEActivation.GELU_TANH, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp8Static, + kMxfp8Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert (supported, reason) == (False, "MXFP8 W8A8 supports only SiLU") + + +def test_explicit_b12x_mxfp8_selection(monkeypatch: pytest.MonkeyPatch) -> None: + """moe_backend='b12x' resolves the MXFP8 oracle to the b12x experts.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle + + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: True) + config = make_dummy_moe_config( + hidden_dim=2560, intermediate_size=640, experts_per_token=10 + ) + config.moe_backend = "b12x" + + backend, experts_cls = mxfp8_oracle.select_mxfp8_moe_backend( + config, prepares_b12x=True + ) + + assert backend is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + assert experts_cls is B12xExperts + + +def test_b12x_mxfp8_auto_selection_keeps_conservative_priority( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Automatic selection must not displace an established backend: b12x is + registered only ahead of the dequantize-to-BF16 emulation.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle + + backends = mxfp8_oracle._SUPPORTED_BACKENDS + assert backends[-1] is mxfp8_oracle.Fp8MoeBackend.EMULATION + assert backends[-2] is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + + def always_supported(cls, config, weight_key, activation_key, activation_format): + del config, weight_key, activation_key, activation_format + return True, None + + established = type( + "EstablishedExperts", + (mk.FusedMoEExperts,), + { + "is_supported_config": staticmethod(always_supported), + }, + ) + monkeypatch.setattr( + mxfp8_oracle, + "_mxfp8_backend_to_kernel_cls", + lambda backend: ( + [B12xExperts] + if backend is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + else [established] + ), + ) + config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64) + + backend, experts_cls = mxfp8_oracle.select_mxfp8_moe_backend( + config, prepares_b12x=True + ) + + assert backend is not mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + assert experts_cls is established + + +def test_b12x_mxfp8_auto_selection_when_nothing_else_supports( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """With every established backend unsupported, b12x wins over emulation.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle + + def b12x_supported( + cls, k_cls, moe_config, weight_key, activation_key, activation_format + ): + del cls, k_cls, moe_config, weight_key, activation_key, activation_format + return True, None + + monkeypatch.setattr(B12xExperts, "is_supported_config", classmethod(b12x_supported)) + config = make_dummy_moe_config( + hidden_dim=2560, intermediate_size=640, experts_per_token=10 + ) + + backend, experts_cls = mxfp8_oracle.select_mxfp8_moe_backend( + config, prepares_b12x=True + ) + + assert backend is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + assert experts_cls is B12xExperts + + +def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The real b12x ``plan_weights``/``prepare_weights`` API receives the + serialized ModelOpt MXFP8 tensors unmodified through the vLLM adapter: + the plan declares ``mxfp8_e8m0_k32`` / ``w8a8_mx`` / ``w31``, and the + prepared owner preserves the E4M3 values byte-exact while b12x (not + vLLM) performs the host scale swizzle.""" + fused_moe = pytest.importorskip("b12x.moe.fused_moe") + from b12x.moe.fused_moe.source import PackedSourceFormat + + if not hasattr(PackedSourceFormat, "MXFP8_E8M0_K32"): + pytest.skip("b12x build lacks the mxfp8_e8m0_k32 source format") + + # Only vLLM-side registration surface is stubbed; plan_weights and + # prepare_weights run for real against the installed b12x. + monkeypatch.setattr(b12x, "_is_current_stream_capturing", lambda: False) + monkeypatch.setattr(b12x, "set_b12x_preparation_provider", lambda *_, **__: None) + monkeypatch.setattr( + b12x, "_register_b12x_moe_output_collective", lambda *_, **__: None + ) + + num_experts, hidden_size, intermediate_size = 4, 256, 128 + w13 = torch.full( + (num_experts, 2 * intermediate_size, hidden_size), + 0.5, + dtype=torch.float8_e4m3fn, + ) + w2 = torch.full( + (num_experts, hidden_size, intermediate_size), + 0.25, + dtype=torch.float8_e4m3fn, + ) + w13_scale = torch.full( + (num_experts, 2 * intermediate_size, hidden_size // 32), + 130, + dtype=torch.uint8, + ) + w2_scale = torch.full( + (num_experts, hidden_size, intermediate_size // 32), + 131, + dtype=torch.uint8, + ) + quant_config = FusedMoEQuantConfig.make( + "mxfp8", + block_shape=[1, 32], + is_scale_swizzled=False, + w1_scale=w13_scale, + w2_scale=w2_scale, + ) + experts = B12xExperts( + make_dummy_moe_config( + num_experts=num_experts, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + ), + quant_config, + ) + layer = torch.nn.Module() + layer.register_parameter("w13_weight", torch.nn.Parameter(w13, requires_grad=False)) + layer.register_parameter("w2_weight", torch.nn.Parameter(w2, requires_grad=False)) + layer.register_parameter( + "w13_weight_scale", torch.nn.Parameter(w13_scale, requires_grad=False) + ) + layer.register_parameter( + "w2_weight_scale", torch.nn.Parameter(w2_scale, requires_grad=False) + ) + layer.activation = MoEActivation.SILU + layer.apply_router_weight_on_input = False + + try: + experts.process_weights_after_loading(layer) + except RuntimeError as exc: # host packing may require a CUDA device + if "CUDA" in str(exc): + pytest.skip(f"b12x mxfp8 host preparation requires CUDA: {exc}") + raise + + impl = experts._prepared_experts._impl + source = experts._prepared_experts.plan.source + assert source.format is PackedSourceFormat.MXFP8_E8M0_K32 + assert source.w13_layout == fused_moe.W13Layout.W31 + assert experts._prepared_experts.plan.activation.mode == "a8" + assert experts._prepared_experts.plan.activation.nonlinearity == "silu" + assert experts._prepared_experts.plan.activation.io_dtype is torch.bfloat16 + # MXFP8 stores one byte per value: the declared geometry carries the + # loaded extents unchanged (no FP4-style doubling of the w2 last dim). + assert experts._prepared_experts.plan.geometry.intermediate_size == ( + intermediate_size + ) + assert experts._prepared_experts.plan.geometry.hidden_size == hidden_size + + # Byte-exact E4M3 preservation through the prepared owner (canonical + # aliases, with the representation's *_values spelling as fallback). + w13_prepared = getattr(impl, "w13_values", None) + w13_prepared = w13_prepared if w13_prepared is not None else impl.w1_fp4 + w2_prepared = getattr(impl, "w2_values", None) + w2_prepared = w2_prepared if w2_prepared is not None else impl.w2_fp4 + assert torch.equal(w13_prepared.view(torch.uint8), w13.view(torch.uint8)) + assert torch.equal(w2_prepared.view(torch.uint8), w2.view(torch.uint8)) + + # The b12x plan discarded the source parameters, so the raw tensors are + # only reachable through the prepared owner. + assert layer.w13_weight.numel() == 0 + assert layer.w2_weight.numel() == 0 + + +_MXFP8_E = 16 +_MXFP8_K = 2560 +_MXFP8_I = 640 +_MXFP8_TOPK = 10 +_MXFP8_TOKENS = 16 + + +def _modelopt_mxfp8_method(): + """Build the real ModelOptMxFp8FusedMoE with moe_backend='b12x' at + reduced step5500 geometry, with weights allocated on the layer.""" + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMxFp8Config, + ModelOptMxFp8FusedMoE, + ) + + quant_config = ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo="NO_QUANT", + exclude_modules=[], + ) + moe_config = make_dummy_moe_config( + num_experts=_MXFP8_E, + experts_per_token=_MXFP8_TOPK, + hidden_dim=_MXFP8_K, + intermediate_size=_MXFP8_I, + in_dtype=torch.bfloat16, + max_num_tokens=128, + ) + moe_config.moe_backend = "b12x" + method = ModelOptMxFp8FusedMoE(quant_config=quant_config, moe_config=moe_config) + assert method.mxfp8_backend is Fp8MoeBackend.B12X_MXFP8 + assert method.experts_cls is B12xExperts + + layer = torch.nn.Module() + layer.intermediate_size_per_partition = _MXFP8_I + layer.hidden_size = _MXFP8_K + layer.moe_config = moe_config + layer.activation = MoEActivation.SILU + layer.apply_router_weight_on_input = False + layer.num_expert_group = None + layer.topk_group = None + layer.e_score_correction_bias = None + layer.routed_scaling_factor = None + layer.global_num_experts = _MXFP8_E + layer.expert_map = None + layer._expert_routing_tables = lambda: None + # create_weights registers ModelWeightParameter members; torch.device + # makes torch.empty allocate them on CUDA like a real worker would. + with torch.device("cuda"): + method.create_weights( + layer, + num_experts=_MXFP8_E, + hidden_size=_MXFP8_K, + intermediate_size_per_partition=_MXFP8_I, + params_dtype=torch.bfloat16, + ) + return method, layer + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@torch.inference_mode() +def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( + dist_init, + workspace_init, +) -> None: + """Drive the real ModelOptMxFp8FusedMoE lifecycle (create_weights -> + serialized weights -> process_weights_after_loading -> declare+prepare -> + apply) with the b12x w8a8_mx backend at reduced step5500 geometry, + against the dequantized torch reference. This exercises the full + vLLM -> B12X integration through the quant method, not B12xExperts in + isolation.""" + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + dequant_mxfp8_to_bf16, + ) + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + monkeypatch = pytest.MonkeyPatch() + try: + # The TP collective descriptor needs live parallel state a unit + # test does not build; plan preparation does not depend on it. + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.b12x." + "_register_b12x_moe_output_collective", + lambda *_, **__: None, + ) + with torch.device("cuda"): + method, layer = _modelopt_mxfp8_method() + set_random_seed(29) + hidden_states = ( + torch.randn( + (_MXFP8_TOKENS, _MXFP8_K), device="cuda", dtype=torch.bfloat16 + ) + / 10 + ) + score = torch.randn( + (_MXFP8_TOKENS, _MXFP8_E), device="cuda", dtype=torch.bfloat16 + ) + w1_bf16 = ( + torch.randn( + (_MXFP8_E, 2 * _MXFP8_I, _MXFP8_K), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + w2_bf16 = ( + torch.randn( + (_MXFP8_E, _MXFP8_K, _MXFP8_I), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + w13_q, w13_scale = _mxfp8_e4m3_quantize_torch( + w1_bf16, is_sf_swizzled_layout=False + ) + w2_q, w2_scale = _mxfp8_e4m3_quantize_torch( + w2_bf16, is_sf_swizzled_layout=False + ) + layer.w13_weight.data.copy_(w13_q) + layer.w2_weight.data.copy_(w2_q) + layer.w13_weight_scale.data.copy_(w13_scale) + layer.w2_weight_scale.data.copy_(w2_scale) + + method.process_weights_after_loading(layer) + + experts = method.moe_kernel.fused_experts + assert isinstance(experts, B12xExperts) + assert experts._prepared_experts is not None + # Preparation consumed the source parameters. + assert layer.w13_weight.numel() == 0 + assert layer.w2_weight.numel() == 0 + + # Declare and prepare the serving shapes exactly as the warmup + # driver does; apply() resolves only session-prepared plans. + session, _ = _prepare( + layer, + device=torch.device("cuda"), + counts=(_MXFP8_TOKENS,), + output_dtype=torch.bfloat16, + max_tokens=128, + ) + try: + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, _MXFP8_TOPK, renormalize=False + ) + output = method.apply( + layer, + hidden_states, + topk_weights, + topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + finally: + session.close() + finally: + monkeypatch.undo() + + reference = torch_moe( + hidden_states, + dequant_mxfp8_to_bf16(w13_q, w13_scale), + dequant_mxfp8_to_bf16(w2_q, w2_scale), + score, + _MXFP8_TOPK, + activation=MoEActivation.SILU, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output, reference, atol=2e-1, rtol=2e-1) + cosine = torch.nn.functional.cosine_similarity( + output.flatten().float(), reference.flatten().float(), dim=0 + ) + assert cosine > 0.99 diff --git a/tests/v1/spec_decode/test_draft_moe_backend_override.py b/tests/v1/spec_decode/test_draft_moe_backend_override.py index 5a4522e5b301..667979fd50d9 100644 --- a/tests/v1/spec_decode/test_draft_moe_backend_override.py +++ b/tests/v1/spec_decode/test_draft_moe_backend_override.py @@ -18,6 +18,11 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +@dataclass +class _AttentionConfig: + backend: str | None = None + + @dataclass class _KernelConfig: moe_backend: str | None = None @@ -38,6 +43,7 @@ class _SpeculativeConfig: @dataclass class _VllmConfig: + attention_config: _AttentionConfig kernel_config: _KernelConfig cache_config: _CacheConfig speculative_config: _SpeculativeConfig @@ -46,6 +52,7 @@ class _VllmConfig: def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig: return _VllmConfig( + attention_config=_AttentionConfig(), kernel_config=_KernelConfig(moe_backend=target_moe), cache_config=_CacheConfig(), speculative_config=_SpeculativeConfig(moe_backend=draft_moe), diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index cdc9852d0735..9178ff08561f 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -273,7 +273,7 @@ class KernelConfig: - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - - "b12x": Use b12x FP4 MoE kernels on SM12x + - "b12x": Use b12x FP4 and MXFP8 MoE kernels on SM12x - "flashinfer_moe_ep_mega_deep_gemm": Use the FlashInfer moe_ep expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an MXFP4 checkpoint verbatim (Blackwell, requires expert parallel; diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index af958dca03ba..f7c369ffbe6d 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 functools import weakref from collections.abc import Sequence from dataclasses import dataclass @@ -24,6 +25,7 @@ QuantKey, kMxfp4Static, kMxfp8Dynamic, + kMxfp8Static, kNvfp4Dynamic, kNvfp4Static, ) @@ -45,6 +47,7 @@ tuple[str, str, str], ] = { ("mxfp4", "mxfp8"): ("w4a8_mx", "fp4_e8m0_k32", "w31"), + ("mxfp8", "mxfp8"): ("w8a8_mx", "mxfp8_e8m0_k32", "w31"), ("mxfp4", None): ("w4a16", "fp4_e8m0_k32", "w31"), ("exl3", None): ("w4a16", "exl3", "w31"), ("nvfp4", "nvfp4"): ("nvfp4", "modelopt_nvfp4", "w31"), @@ -237,6 +240,16 @@ def _normalize_expert_scale(scale: torch.Tensor) -> torch.Tensor: return scale.to(dtype=torch.float32).contiguous() +@functools.cache +def _b12x_has_mxfp8_moe() -> bool: + """Whether the installed b12x prepares MXFP8 (w8a8_mx) MoE experts.""" + try: + from b12x.moe.fused_moe import PackedSourceFormat + except ImportError: + return False + return "mxfp8_e8m0_k32" in {fmt.value for fmt in PackedSourceFormat} + + class B12xExperts(mk.FusedMoEExpertsModular): """Packed MoE experts backed by the b12x SM12x planned API.""" @@ -248,6 +261,7 @@ def __init__( super().__init__(moe_config, quant_config) if quant_config.weight_quant_dtype not in ( "mxfp4", + "mxfp8", "nvfp4", "exl3", "iq2_xs", @@ -255,8 +269,8 @@ def __init__( "q8_0", ): raise ValueError( - "b12x MoE requires MXFP4, NVFP4, EXL3, IQ2_XS, IQ2_XXS or " - "Q8_0 weights, got " + "b12x MoE requires MXFP4, MXFP8, NVFP4, EXL3, IQ2_XS, IQ2_XXS " + "or Q8_0 weights, got " f"{quant_config.weight_quant_dtype}" ) scheme = ( @@ -375,7 +389,13 @@ def _prepare_experts( num_experts = int(w1.shape[0]) hidden_size = int(w2.shape[1]) - intermediate_size = int(w2.shape[2]) * 2 + # FP4 weights pack two values per byte, so the loaded w2 last dim is + # half the logical channel count; MXFP8 stores one byte per value and + # the loaded extent is already logical. + if quant_mode in ("w8a8_mx",): + intermediate_size = int(w2.shape[2]) + else: + intermediate_size = int(w2.shape[2]) * 2 unit_scale = self._unit_scale(w1.device, num_experts) w1_global_scale = self._weight_global_scale( w1.device, num_experts, self.g1_alphas, "w1 global scales" @@ -397,6 +417,7 @@ def _prepare_experts( mode = { "w4a16": fused_moe.ActivationMode.A16, "w4a8_mx": fused_moe.ActivationMode.A8, + "w8a8_mx": fused_moe.ActivationMode.A8, "w4a8_nvfp4": fused_moe.ActivationMode.A8, "nvfp4": fused_moe.ActivationMode.A4, }[quant_mode] @@ -581,6 +602,26 @@ def is_supported_config( "per-rank intermediate size divisible by 32" ), ) + if weight_key == kMxfp8Static: + if not _b12x_has_mxfp8_moe(): + return False, "the installed b12x has no MXFP8 W8A8 MoE recipe" + if moe_config.activation != MoEActivation.SILU: + return False, "MXFP8 W8A8 supports only SiLU" + unpadded_intermediate_size = ( + moe_config.intermediate_size_per_partition_unpadded + or moe_config.intermediate_size_per_partition + ) + if ( + moe_config.hidden_dim % 128 != 0 + or unpadded_intermediate_size % 128 != 0 + ): + return ( + False, + ( + "MXFP8 W8A8 requires hidden size and per-rank " + "intermediate size divisible by 128" + ), + ) return mk.FusedMoEExperts.is_supported_config( cls, moe_config, weight_key, activation_key, activation_format ) @@ -609,6 +650,7 @@ def _supports_quant_scheme( return (weight_key, activation_key) in ( (kMxfp4Static, kMxfp8Dynamic), (kMxfp4Static, None), + (kMxfp8Static, kMxfp8Dynamic), (kNvfp4Static, kNvfp4Dynamic), (kNvfp4Static, kMxfp8Dynamic), (kNvfp4Static, None), diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index abe47373dea4..c281fd513c30 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -65,6 +65,10 @@ class Fp8MoeBackend(Enum): TRITON_MXFP8 = "TRITON_MXFP8" # MXFP8 MoE via AITER (FlyDSL two-stage grouped GEMM) on gfx950. AITER_MXFP8 = "AITER_MXFP8" + # MXFP8 MoE through the b12x SM12x planned-expert API (w8a8_mx): raw + # E4M3 weights and UE8M0 K/32 scales pass through verbatim; BF16 + # activations are quantized inside the kernel. + B12X_MXFP8 = "B12X_MXFP8" def _get_priority_backends( @@ -590,8 +594,11 @@ def convert_to_fp8_moe_kernel_format( Fp8MoeBackend.HPC, # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes # the MXFP8 weights as-is — neither needs a load-time layout change. + # B12X_MXFP8 likewise binds the raw E4M3/E8M0 tensors during its + # process_weights_after_loading preparation. Fp8MoeBackend.EMULATION, Fp8MoeBackend.TRITON_MXFP8, + Fp8MoeBackend.B12X_MXFP8, ]: raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 2a03eacea8e1..0877b093ed20 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -27,6 +27,11 @@ Fp8MoeBackend.AITER_MXFP8, Fp8MoeBackend.HUMMING, Fp8MoeBackend.TRITON_MXFP8, + # b12x (SM12x) is registered at the most conservative priority: every + # deployment an established backend already serves keeps that backend, + # and b12x only wins over the dequantize-to-BF16 emulation. Prefer + # moe_backend="b12x" to select it explicitly. + Fp8MoeBackend.B12X_MXFP8, Fp8MoeBackend.EMULATION, ) @@ -36,6 +41,7 @@ "marlin": Fp8MoeBackend.MARLIN, "xpu": Fp8MoeBackend.XPU, "aiter": Fp8MoeBackend.AITER_MXFP8, + "b12x": Fp8MoeBackend.B12X_MXFP8, "triton": Fp8MoeBackend.TRITON_MXFP8, "humming": Fp8MoeBackend.HUMMING, } @@ -68,6 +74,10 @@ def _mxfp8_backend_to_kernel_cls( ) return [Mxfp8NativeTritonExperts] + if backend == Fp8MoeBackend.B12X_MXFP8: + from vllm.model_executor.layers.fused_moe.b12x import B12xExperts + + return [B12xExperts] if backend == Fp8MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, @@ -106,9 +116,16 @@ def _select_kernel_cls( def select_mxfp8_moe_backend( config: FusedMoEConfig, + *, + prepares_b12x: bool = False, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: """Select the MXFP8 MoE backend and the best expert class. + Args: + config: The MoE layer configuration. + prepares_b12x: Whether the caller's quantization method prepares + B12X experts after loading. Only those methods may use B12X. + Returns: A tuple of (fp8_backend, experts_cls). """ @@ -122,6 +139,12 @@ def select_mxfp8_moe_backend( f"MXFP8 MoE. Expected one of " f"{list(_BACKEND_NAME_MAP.keys())}." ) + if backend == Fp8MoeBackend.B12X_MXFP8 and not prepares_b12x: + raise ValueError( + "moe_backend='b12x' runs MXFP8 MoE only for ModelOpt MXFP8 " + "checkpoints; this quantization method does not prepare B12X " + "experts." + ) logger.info_once( "Using '%s' MxFp8 MoE backend (user-requested).", backend.value, @@ -130,6 +153,8 @@ def select_mxfp8_moe_backend( # Auto-select: pick the first supported backend. for backend in _SUPPORTED_BACKENDS: + if backend == Fp8MoeBackend.B12X_MXFP8 and not prepares_b12x: + continue try: experts_cls = _select_kernel_cls(backend, config) except ValueError: diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 91eaf10eb280..2a1a16c644a0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1225,7 +1225,13 @@ def _from_config( class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): - """FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints.""" + """MXFP8 block-scale MoE for ModelOpt checkpoints. + + The oracle (``select_mxfp8_moe_backend``) picks the expert backend: + FlashInfer TRTLLM, Triton, Marlin, emulation, or the b12x planned + ``w8a8_mx`` path on SM12x, which prepares directly from the serialized + E4M3 weights and UE8M0 K/32 scale grids. + """ def __init__( self, @@ -1237,7 +1243,9 @@ def __init__( self.quant_config = quant_config assert self.quant_config.is_checkpoint_mxfp8_serialized - self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) + self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend( + config=self.moe, prepares_b12x=True + ) def create_weights( self, @@ -1422,6 +1430,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + # The b12x experts prepare their planned representation here, before + # memory profiling and CUDA graph capture; every other MXFP8 expert + # class inherits the no-op default. + self.moe_kernel.fused_experts.process_weights_after_loading(layer) # No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation # experts would dequant MXFP8->BF16 every forward step. Convert the From 714e851d5374636e607e82bb37d49d003780486c Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Wed, 30 Sep 2026 13:17:14 +0100 Subject: [PATCH 330/365] test(moe): disable established MXFP8 backends in b12x fallback test On SM12x hardware Marlin genuinely supports MXFP8 and outranks b12x, so mocking only b12x as supported made the test pass on CPU hosts and fail on the target GPU. Signed-off-by: Naadir Jeewa --- tests/kernels/moe/test_b12x.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index 066eedef0f4c..000329658f57 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -1813,13 +1813,31 @@ def test_b12x_mxfp8_auto_selection_when_nothing_else_supports( """With every established backend unsupported, b12x wins over emulation.""" import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle - def b12x_supported( - cls, k_cls, moe_config, weight_key, activation_key, activation_format - ): - del cls, k_cls, moe_config, weight_key, activation_key, activation_format + def supported(cls, config, weight_key, activation_key, activation_format): + del config, weight_key, activation_key, activation_format return True, None - monkeypatch.setattr(B12xExperts, "is_supported_config", classmethod(b12x_supported)) + def unsupported(cls, config, weight_key, activation_key, activation_format): + del config, weight_key, activation_key, activation_format + return False, "established backend disabled" + + # Disable every established backend explicitly: on SM12x hardware Marlin + # genuinely supports MXFP8 and outranks b12x. + established = type( + "EstablishedExperts", + (mk.FusedMoEExperts,), + {"is_supported_config": staticmethod(unsupported)}, + ) + monkeypatch.setattr(B12xExperts, "is_supported_config", staticmethod(supported)) + monkeypatch.setattr( + mxfp8_oracle, + "_mxfp8_backend_to_kernel_cls", + lambda backend: ( + [B12xExperts] + if backend is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 + else [established] + ), + ) config = make_dummy_moe_config( hidden_dim=2560, intermediate_size=640, experts_per_token=10 ) From 858af5cd1a99e70a7972b06b1c12cff40fdd33b6 Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Wed, 30 Sep 2026 18:10:53 +0100 Subject: [PATCH 331/365] feat(moe): prepare B12X MXFP8 experts at 32-aligned intermediate sizes B12X W8A8 MXFP8 runs 32-aligned per-rank intermediate sizes natively (split up/gate FC1 descriptors), so the ModelOpt path prepares experts at their sharded width: step-5500 Qwen3.8 MTP experts are 640 wide and run at 320 per rank under TP=2. is_supported_config gates on hidden % 128 and per-rank intermediate % 32, which is the kernel contract. The I=320 lifecycle test asserts the prepared geometry is unpadded. Requires local-inference-lab/b12x native MXFP8 W8A8 MoE with sub-128 intermediate support. Signed-off-by: Naadir Jeewa --- tests/kernels/moe/test_b12x.py | 35 ++++++++++++++------ vllm/model_executor/layers/fused_moe/b12x.py | 18 ++++------ 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index 000329658f57..85f9a71911a3 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -1677,7 +1677,8 @@ def check(output, rows): _MXFP8_W8A8_ALIGNMENT_REASON = ( - "MXFP8 W8A8 requires hidden size and per-rank intermediate size divisible by 128" + "MXFP8 W8A8 requires hidden size divisible by 128 and " + "per-rank intermediate size divisible by 32" ) @@ -1690,9 +1691,14 @@ def check(output, rows): id="step5500-shape-supported", ), pytest.param( - {"hidden_dim": 256, "intermediate_size": 64}, + {"hidden_dim": 2560, "intermediate_size": 320}, + None, + id="step5500-tp2-shard-supported", + ), + pytest.param( + {"hidden_dim": 256, "intermediate_size": 48}, _MXFP8_W8A8_ALIGNMENT_REASON, - id="intermediate-64-rejected", + id="intermediate-48-rejected", ), pytest.param( {"hidden_dim": 100, "intermediate_size": 128}, @@ -1958,12 +1964,11 @@ def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( _MXFP8_E = 16 _MXFP8_K = 2560 -_MXFP8_I = 640 _MXFP8_TOPK = 10 _MXFP8_TOKENS = 16 -def _modelopt_mxfp8_method(): +def _modelopt_mxfp8_method(intermediate: int): """Build the real ModelOptMxFp8FusedMoE with moe_backend='b12x' at reduced step5500 geometry, with weights allocated on the layer.""" from vllm.model_executor.layers.quantization.modelopt import ( @@ -1980,7 +1985,7 @@ def _modelopt_mxfp8_method(): num_experts=_MXFP8_E, experts_per_token=_MXFP8_TOPK, hidden_dim=_MXFP8_K, - intermediate_size=_MXFP8_I, + intermediate_size=intermediate, in_dtype=torch.bfloat16, max_num_tokens=128, ) @@ -1990,7 +1995,7 @@ def _modelopt_mxfp8_method(): assert method.experts_cls is B12xExperts layer = torch.nn.Module() - layer.intermediate_size_per_partition = _MXFP8_I + layer.intermediate_size_per_partition = intermediate layer.hidden_size = _MXFP8_K layer.moe_config = moe_config layer.activation = MoEActivation.SILU @@ -2009,7 +2014,7 @@ def _modelopt_mxfp8_method(): layer, num_experts=_MXFP8_E, hidden_size=_MXFP8_K, - intermediate_size_per_partition=_MXFP8_I, + intermediate_size_per_partition=intermediate, params_dtype=torch.bfloat16, ) return method, layer @@ -2017,9 +2022,13 @@ def _modelopt_mxfp8_method(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @torch.inference_mode() +# 640 is the TP=1 expert width; 320 is the TP=2 shard, which B12X prepares +# unpadded (split up/gate FC1 descriptors, TMA zero-filled tail tile). +@pytest.mark.parametrize("intermediate", [640, 320], ids=["tp1-i640", "tp2-i320"]) def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( dist_init, workspace_init, + intermediate: int, ) -> None: """Drive the real ModelOptMxFp8FusedMoE lifecycle (create_weights -> serialized weights -> process_weights_after_loading -> declare+prepare -> @@ -2045,7 +2054,7 @@ def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( lambda *_, **__: None, ) with torch.device("cuda"): - method, layer = _modelopt_mxfp8_method() + method, layer = _modelopt_mxfp8_method(intermediate) set_random_seed(29) hidden_states = ( torch.randn( @@ -2058,7 +2067,7 @@ def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( ) w1_bf16 = ( torch.randn( - (_MXFP8_E, 2 * _MXFP8_I, _MXFP8_K), + (_MXFP8_E, 2 * intermediate, _MXFP8_K), device="cuda", dtype=torch.bfloat16, ) @@ -2066,7 +2075,7 @@ def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( ) w2_bf16 = ( torch.randn( - (_MXFP8_E, _MXFP8_K, _MXFP8_I), + (_MXFP8_E, _MXFP8_K, intermediate), device="cuda", dtype=torch.bfloat16, ) @@ -2091,6 +2100,10 @@ def test_b12x_modelopt_mxfp8_post_load_prepares_and_matches_torch( # Preparation consumed the source parameters. assert layer.w13_weight.numel() == 0 assert layer.w2_weight.numel() == 0 + assert ( + experts._prepared_experts.plan.geometry.intermediate_size + == intermediate + ) # Declare and prepare the serving shapes exactly as the warmup # driver does; apply() resolves only session-prepared plans. diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index f7c369ffbe6d..dcb2d00fe134 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -392,6 +392,7 @@ def _prepare_experts( # FP4 weights pack two values per byte, so the loaded w2 last dim is # half the logical channel count; MXFP8 stores one byte per value and # the loaded extent is already logical. + w1_scale, w2_scale = self.w1_scale, self.w2_scale if quant_mode in ("w8a8_mx",): intermediate_size = int(w2.shape[2]) else: @@ -445,8 +446,8 @@ def _prepare_experts( weights=fused_moe.PackedWeights( w13=w1, w2=w2, - w13_block_scales=self.w1_scale, - w2_block_scales=self.w2_scale, + w13_block_scales=w1_scale, + w2_block_scales=w2_scale, w13_global_scales=w1_global_scale, w2_global_scales=w2_global_scale, input_scale=a1_gscale, @@ -607,19 +608,12 @@ def is_supported_config( return False, "the installed b12x has no MXFP8 W8A8 MoE recipe" if moe_config.activation != MoEActivation.SILU: return False, "MXFP8 W8A8 supports only SiLU" - unpadded_intermediate_size = ( - moe_config.intermediate_size_per_partition_unpadded - or moe_config.intermediate_size_per_partition - ) - if ( - moe_config.hidden_dim % 128 != 0 - or unpadded_intermediate_size % 128 != 0 - ): + if moe_config.hidden_dim % 128 != 0 or unpadded_intermediate_size % 32 != 0: return ( False, ( - "MXFP8 W8A8 requires hidden size and per-rank " - "intermediate size divisible by 128" + "MXFP8 W8A8 requires hidden size divisible by 128 and " + "per-rank intermediate size divisible by 32" ), ) return mk.FusedMoEExperts.is_supported_config( From 71dbc71eddd55fc62f6860e24e7b2e5aacc4443a Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Wed, 30 Sep 2026 23:16:11 +0100 Subject: [PATCH 332/365] docs(b12x): state the MXFP8 per-rank intermediate requirement as % 32 is_supported_config gates on hidden % 128 and per-rank intermediate % 32; the TP=2 shard of 320 is supported. Signed-off-by: Naadir Jeewa --- docs/features/quantization/b12x.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index da3946b2b1d5..2eb9c7ecf44a 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -80,8 +80,8 @@ source): preparation preserves serialized E4M3 weight bytes and losslessly swizzles the UE8M0 K/32 scale grids into the kernel layout. The kernel quantizes BF16 hidden states per 32 elements. Select it with `--moe-backend b12x`; automatic selection keeps every established backend, and b12x is registered -ahead of emulation only. The A8 path requires SiLU and hidden and per-rank -intermediate sizes divisible by 128. +ahead of emulation only. The A8 path requires SiLU, a hidden size divisible +by 128, and a per-rank intermediate size divisible by 32. This requires the paired b12x source build with the `w8a8_mx` MoE recipe. For a mixed checkpoint with NVFP4 target experts and MXFP8 MTP experts, retain From 66ab251c491f668f24fe3a0194d67c7e8314b8f3 Mon Sep 17 00:00:00 2001 From: Naadir Jeewa Date: Wed, 30 Sep 2026 23:39:32 +0100 Subject: [PATCH 333/365] docs(moe): document B12X MXFP8 adapter and test helpers Signed-off-by: Naadir Jeewa The moe_problem_size docstring states N = 2 * intermediate for every format, matching the fix folded into the first commit (Martin Vit). --- tests/kernels/moe/test_b12x.py | 28 +++++++++++++++++++ .../test_draft_moe_backend_override.py | 1 + vllm/model_executor/layers/fused_moe/b12x.py | 25 +++++++++++++++++ .../layers/fused_moe/oracle/fp8.py | 7 +++++ .../layers/quantization/modelopt.py | 7 +++++ 5 files changed, 68 insertions(+) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index 85f9a71911a3..ca169c31d207 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -454,6 +454,11 @@ def test_explicit_b12x_mxfp4_selection( force_a16: bool, expected_backend: Mxfp4MoeBackend, ) -> None: + """An explicit `moe_backend='b12x'` selects the b12x MXFP4 W4A8 experts. + + The activation key -- dynamic MXFP8 or absent -- does not change the + choice while A16 forcing is off. + """ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) monkeypatch.setattr( @@ -509,6 +514,11 @@ def test_deepseek_v4_b12x_activation_selection( force_a16: bool, expected_backend: Mxfp4MoeBackend, ) -> None: + """DeepSeek V4 b12x selection follows the forced-A16 knob. + + MXFP4 weights keep the dynamic-MXFP8 activation contract unless A16 is + forced, which selects the BF16 contract. + """ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) monkeypatch.setattr( mxfp4_oracle.envs, @@ -1109,6 +1119,12 @@ def _make_b12x_moe_case( intermediate_size: int = 128, topk: int = 2, ) -> _B12xMoeCase: + """Build one deterministic B12X MoE case for a weight/activation pair. + + Quantizes freshly seeded BF16 expert weights, keeps the dequantized + tensors as the reference, and assembles the quant config matching the + requested activation dtype. + """ set_random_seed(seed) dtype = torch.bfloat16 hidden_states = torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 @@ -1276,6 +1292,7 @@ def test_b12x_moe_matches_torch( activation: MoEActivation, workspace_init, ) -> None: + """Each b12x MoE quantization lane matches the dequantized PyTorch MoE.""" with set_current_vllm_config( VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) ): @@ -1355,6 +1372,7 @@ def test_b12x_moe_cuda_graph_replay( activation_dtype: str | None, workspace_init, ) -> None: + """A captured b12x MoE graph replays with the eager call's output.""" from vllm.v1.worker.workspace import lock_workspace with set_current_vllm_config( @@ -1550,6 +1568,12 @@ def test_b12x_moe_prefill_capacity_and_exact_decode_reuse( workspace_init, capacity, ) -> None: + """One prepared kernel serves every row count up to its capacity. + + Once the session is frozen, eager calls at mixed row counts and route-id + dtypes, plus a capture-scope graph replay, must stay correct without + reallocating workspace buffers. + """ from b12x._lib.runtime_control import kernel_resolution_guard from vllm.v1.worker.workspace import current_workspace_manager, lock_workspace @@ -1734,6 +1758,7 @@ def test_b12x_mxfp8_w8a8_config_support( def test_b12x_mxfp8_w8a8_rejects_non_silu_activation( monkeypatch: pytest.MonkeyPatch, ) -> None: + """MXFP8 W8A8 on b12x supports SiLU only.""" monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: True) config = make_dummy_moe_config( @@ -1784,6 +1809,7 @@ def test_b12x_mxfp8_auto_selection_keeps_conservative_priority( assert backends[-2] is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 def always_supported(cls, config, weight_key, activation_key, activation_format): + """Stand-in: the established backend claims every config.""" del config, weight_key, activation_key, activation_format return True, None @@ -1820,10 +1846,12 @@ def test_b12x_mxfp8_auto_selection_when_nothing_else_supports( import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle def supported(cls, config, weight_key, activation_key, activation_format): + """Stand-in: the b12x experts claim the config.""" del config, weight_key, activation_key, activation_format return True, None def unsupported(cls, config, weight_key, activation_key, activation_format): + """Stand-in: the established backend rejects every config.""" del config, weight_key, activation_key, activation_format return False, "established backend disabled" diff --git a/tests/v1/spec_decode/test_draft_moe_backend_override.py b/tests/v1/spec_decode/test_draft_moe_backend_override.py index 667979fd50d9..0c029c8c3b48 100644 --- a/tests/v1/spec_decode/test_draft_moe_backend_override.py +++ b/tests/v1/spec_decode/test_draft_moe_backend_override.py @@ -51,6 +51,7 @@ class _VllmConfig: def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig: + """Build the minimal config stand-in consumed by the draft-loading path.""" return _VllmConfig( attention_config=_AttentionConfig(), kernel_config=_KernelConfig(moe_backend=target_moe), diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index dcb2d00fe134..7dd1331c783f 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -258,6 +258,11 @@ def __init__( moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig, ): + """Bind the b12x quant mode, packed source format and W13 layout. + + Raises ValueError when no b12x kernel covers the configured weight + dtype or the (weight, activation) scheme; no weights are prepared. + """ super().__init__(moe_config, quant_config) if quant_config.weight_quant_dtype not in ( "mxfp4", @@ -353,6 +358,14 @@ def _prepare_experts( activation: MoEActivation, params_dtype: torch.dtype, ) -> Any: + """Build the b12x prepared-experts representation for these weights. + + Block-quantized sources are planned straight from the packed extents + with a ``BlockQuantWeights`` binding; every other mode requires the + w1/w2 block scales, recovers the logical intermediate size from the + packed extent (FP4 holds two values per byte) and binds the weight and + activation global scales. Raises RuntimeError under CUDA graph capture. + """ quant_mode = self._quant_mode if _is_current_stream_capturing(): raise RuntimeError( @@ -557,6 +570,11 @@ def is_supported_config( activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, ) -> tuple[bool, str | None]: + """Report whether b12x can run this MoE config, with the reason if not. + + Applies the kernel-specific bias, dtype, activation and alignment + gates, then defers to the base-class checks. + """ if moe_config.has_bias: return False, "kernel does not support expert biases" if moe_config.in_dtype not in (torch.float16, torch.bfloat16): @@ -641,6 +659,7 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: + """Check the (weight, activation) pair against the b12x kernel set.""" return (weight_key, activation_key) in ( (kMxfp4Static, kMxfp8Dynamic), (kMxfp4Static, None), @@ -839,6 +858,12 @@ def moe_problem_size( w2: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple[int, int, int, int, int]: + """Return (experts, tokens, N, K, topk) for the current problem. + + Delegates to the base implementation while the weight tensors still + hold data; once they are released, the geometry comes from the + prepared plan, with N = 2 * intermediate (the gate and up rows). + """ if w1.numel() and w2.numel(): return super().moe_problem_size(a1, w1, w2, topk_ids) prepared = self._prepared() diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index c281fd513c30..89f29fbcf670 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -502,6 +502,13 @@ def convert_to_fp8_moe_kernel_format( w13_input_scale: torch.Tensor | None, w2_input_scale: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert canonical FP8/MXFP8 expert weights into `fp8_backend`'s layout. + + Returns the (w13, w2, w13_scale, w2_scale) the kernel consumes: backends + with no load-time conversion pass the tensors through unchanged, HUMMING + rewrites the layer parameters and they are read back from the layer, and a + backend without a converter raises ValueError. + """ block_quant = hasattr(layer, "weight_block_size") if fp8_backend in [Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.BATCHED_DEEPGEMM]: assert block_quant diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 2a1a16c644a0..da6826cfbf66 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1395,6 +1395,13 @@ def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None: ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + """Convert the loaded MXFP8 weights and build the MoE kernel once. + + Runs the selected backend's weight converter, refreshes the quant + config and kernel, lets the experts class prepare its own + representation, and dequantizes to BF16 at load time on the emulation + backend. A second call for the same layer is a no-op. + """ # TODO(bnell): why is this required only for mxfp8? if getattr(layer, "_already_called_process_weights_after_loading", False): return From fb5cacb1daf5195ed0c85364b81885b0a49fb673 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 09:49:12 +0000 Subject: [PATCH 334/365] test(moe): pin B12X MXFP8 selection gates and the prepared FC1 order - An installed b12x without the MXFP8 recipe is rejected at backend selection. - moe_backend=b12x is rejected for MXFP8 quantization methods that do not prepare B12X experts, and auto-selection skips B12X for them. - The preparation test uses distinct E4M3 bytes and compares against copies taken before preparation. b12x flips the gate-first FC1 in place, so the constant-valued comparison against the (flipped) source tensor could not fail. - MXFP8 GPU cases skip when the installed b12x lacks the recipe, and the ModelOpt lifecycle test requires b12x on SM120 instead of any CUDA device. Co-Authored-By: Claude Opus 5.5 --- tests/kernels/moe/test_b12x.py | 96 ++++++++++++++++++++++++++++------ 1 file changed, 80 insertions(+), 16 deletions(-) diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index ca169c31d207..39b5eaa63e6d 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -203,6 +203,10 @@ def _has_b12x_moe() -> bool: ) +def _has_b12x_mxfp8_moe() -> bool: + return _has_b12x_moe() and b12x._b12x_has_mxfp8_moe() + + def _make_b12x_moe_kernel( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -1125,6 +1129,8 @@ def _make_b12x_moe_case( tensors as the reference, and assembles the quant config matching the requested activation dtype. """ + if weight_dtype == "mxfp8" and not b12x._b12x_has_mxfp8_moe(): + pytest.skip("the installed b12x has no MXFP8 W8A8 MoE recipe") set_random_seed(seed) dtype = torch.bfloat16 hidden_states = torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 @@ -1778,6 +1784,33 @@ def test_b12x_mxfp8_w8a8_rejects_non_silu_activation( assert (supported, reason) == (False, "MXFP8 W8A8 supports only SiLU") +def test_b12x_mxfp8_w8a8_requires_the_b12x_recipe( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An installed b12x without the MXFP8 recipe is rejected at selection, + not after the model has loaded.""" + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: False) + config = make_dummy_moe_config( + hidden_dim=2560, + intermediate_size=640, + activation=MoEActivation.SILU, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp8Static, + kMxfp8Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert (supported, reason) == ( + False, + "the installed b12x has no MXFP8 W8A8 MoE recipe", + ) + + def test_explicit_b12x_mxfp8_selection(monkeypatch: pytest.MonkeyPatch) -> None: """moe_backend='b12x' resolves the MXFP8 oracle to the b12x experts.""" import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle @@ -1797,6 +1830,24 @@ def test_explicit_b12x_mxfp8_selection(monkeypatch: pytest.MonkeyPatch) -> None: assert experts_cls is B12xExperts +def test_explicit_b12x_mxfp8_selection_requires_b12x_preparation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Quantization methods that do not prepare B12X experts (compressed + tensors, INC, online MXFP8) reject moe_backend='b12x' at init.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle + + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(b12x, "_b12x_has_mxfp8_moe", lambda: True) + config = make_dummy_moe_config( + hidden_dim=2560, intermediate_size=640, experts_per_token=10 + ) + config.moe_backend = "b12x" + + with pytest.raises(ValueError, match="does not prepare B12X experts"): + mxfp8_oracle.select_mxfp8_moe_backend(config) + + def test_b12x_mxfp8_auto_selection_keeps_conservative_priority( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1842,7 +1893,8 @@ def always_supported(cls, config, weight_key, activation_key, activation_format) def test_b12x_mxfp8_auto_selection_when_nothing_else_supports( monkeypatch: pytest.MonkeyPatch, ) -> None: - """With every established backend unsupported, b12x wins over emulation.""" + """With every established backend unsupported, b12x wins over emulation, + but only for quantization methods that prepare B12X experts.""" import vllm.model_executor.layers.fused_moe.oracle.mxfp8 as mxfp8_oracle def supported(cls, config, weight_key, activation_key, activation_format): @@ -1882,6 +1934,8 @@ def unsupported(cls, config, weight_key, activation_key, activation_format): assert backend is mxfp8_oracle.Fp8MoeBackend.B12X_MXFP8 assert experts_cls is B12xExperts + with pytest.raises(ValueError, match="No MXFP8 MoE backends available"): + mxfp8_oracle.select_mxfp8_moe_backend(config) def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( @@ -1890,8 +1944,9 @@ def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( """The real b12x ``plan_weights``/``prepare_weights`` API receives the serialized ModelOpt MXFP8 tensors unmodified through the vLLM adapter: the plan declares ``mxfp8_e8m0_k32`` / ``w8a8_mx`` / ``w31``, and the - prepared owner preserves the E4M3 values byte-exact while b12x (not - vLLM) performs the host scale swizzle.""" + prepared owner keeps the E4M3 bytes (w2 verbatim, the gate-first FC1 + with its halves flipped to kernel order) while b12x (not vLLM) performs + the host scale swizzle.""" fused_moe = pytest.importorskip("b12x.moe.fused_moe") from b12x.moe.fused_moe.source import PackedSourceFormat @@ -1907,16 +1962,19 @@ def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( ) num_experts, hidden_size, intermediate_size = 4, 256, 128 - w13 = torch.full( - (num_experts, 2 * intermediate_size, hidden_size), - 0.5, - dtype=torch.float8_e4m3fn, - ) - w2 = torch.full( - (num_experts, hidden_size, intermediate_size), - 0.25, - dtype=torch.float8_e4m3fn, - ) + generator = torch.Generator().manual_seed(0) + + def random_e4m3(*shape: int) -> torch.Tensor: + # Distinct bytes make the FC1 half order observable; 0x7F/0xFF are NaN. + values = torch.randint(0, 256, shape, generator=generator, dtype=torch.uint8) + values[(values & 0x7F) == 0x7F] = 0x38 + return values.view(torch.float8_e4m3fn) + + w13 = random_e4m3(num_experts, 2 * intermediate_size, hidden_size) + w2 = random_e4m3(num_experts, hidden_size, intermediate_size) + # b12x flips the gate-first FC1 in place, so compare against copies. + w13_source = w13.view(torch.uint8).clone() + w2_source = w2.view(torch.uint8).clone() w13_scale = torch.full( (num_experts, 2 * intermediate_size, hidden_size // 32), 130, @@ -1981,8 +2039,12 @@ def test_b12x_mxfp8_preparation_passes_source_tensors_verbatim( w13_prepared = w13_prepared if w13_prepared is not None else impl.w1_fp4 w2_prepared = getattr(impl, "w2_values", None) w2_prepared = w2_prepared if w2_prepared is not None else impl.w2_fp4 - assert torch.equal(w13_prepared.view(torch.uint8), w13.view(torch.uint8)) - assert torch.equal(w2_prepared.view(torch.uint8), w2.view(torch.uint8)) + up_first = torch.cat( + [w13_source[:, intermediate_size:], w13_source[:, :intermediate_size]], + dim=1, + ) + assert torch.equal(w13_prepared.view(torch.uint8).cpu(), up_first) + assert torch.equal(w2_prepared.view(torch.uint8).cpu(), w2_source) # The b12x plan discarded the source parameters, so the raw tensors are # only reachable through the prepared owner. @@ -2048,7 +2110,9 @@ def _modelopt_mxfp8_method(intermediate: int): return method, layer -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.skipif( + not _has_b12x_mxfp8_moe(), reason="requires b12x MXFP8 MoE on SM120" +) @torch.inference_mode() # 640 is the TP=1 expert width; 320 is the TP=2 shard, which B12X prepares # unpadded (split up/gate FC1 descriptors, TMA zero-filled tail tile). From 69ba12d03f825908e41fccd1becdd9690f535a1d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 10:03:49 +0000 Subject: [PATCH 335/365] docs(release): describe B12X MXFP8 MTP draft experts Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-957.json | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 .lil/changes/vllm-957.json diff --git a/.lil/changes/vllm-957.json b/.lil/changes/vllm-957.json new file mode 100644 index 000000000000..4ba9206657f3 --- /dev/null +++ b/.lil/changes/vllm-957.json @@ -0,0 +1,26 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-957", + "category": "feature", + "summary": "MXFP8 MTP draft experts can run on B12X", + "models": [ + "Qwen3.8-Flash-Next" + ], + "compatibility": "No user action required. The default is unchanged: the launcher keeps moe_backend auto for MXFP8 drafters (docker-114), which picks Marlin. To try B12X, pass an explicit --speculative-config with \"moe_backend\": \"b12x\".", + "details": [ + "The Qwen3.8-Flash-Next QAD step 5500 checkpoint (revision qad-step5500-ple1000) ships its MTP draft experts as MXFP8 while the target experts are NVFP4. With moe_backend b12x in the speculative config, ModelOpt MXFP8 experts now run on B12X instead of failing kernel selection.", + "On RTX PRO 6000 TP2 the B12X drafter is correct and accepts slightly more draft tokens (2.34-2.37 against 2.30-2.34 per step with Marlin), but decode is not faster: C1 196-209 against 200-220 tok/s, C8 949-955 against 960-961 tok/s.", + "moe_backend b12x for MXFP8 experts is rejected at startup, with a message, when the installed B12X has no MXFP8 recipe or when the checkpoint is not a ModelOpt MXFP8 export (compressed-tensors, INC, online MXFP8).", + "Contributed by Naadir Jeewa." + ], + "authors": [ + "randomvariable", + "voipmonitor" + ], + "pull_requests": [ + 957 + ], + "requires": [ + "b12x-452" + ] +} From 6f4b64421721e8c891d1597b949f9a7439a30e85 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 10:48:20 +0000 Subject: [PATCH 336/365] fix(scheduler): give prefill lanes to queued requests only while admissible PrefillInterleaveController ranks running and queued requests together, and never-scheduled requests rank first. While every running slot was taken, or new requests were paused, queued requests that the waiting pass cannot admit took the prefill lanes. The running chunked prefills then had none, and every step scheduled nothing: generation stopped for good with idle GPUs. Queued requests now compete for lanes only while the waiting pass can admit one, the same condition it applies itself. Reported, with the root cause and this fix, by netwalker4370 on Discord (DeepSeek-V4.1-Flash TP4, --max-parallel-prefills 2, 40 concurrent 150K-token prompts at --max-num-seqs 32). Co-Authored-By: Claude Opus 5.5 --- .../test_prefill_compute_share_scheduler.py | 33 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 12 +++++-- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_prefill_compute_share_scheduler.py b/tests/v1/core/test_prefill_compute_share_scheduler.py index b20f1a2e4713..87da88448640 100644 --- a/tests/v1/core/test_prefill_compute_share_scheduler.py +++ b/tests/v1/core/test_prefill_compute_share_scheduler.py @@ -7,6 +7,7 @@ from vllm.engine.arg_utils import EngineArgs from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.core.boundary_checkpoint import BoundaryCheckpointCache +from vllm.v1.core.sched.interface import PauseState from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -1000,6 +1001,38 @@ def test_round_robin_prioritizes_never_scheduled_prefills(opt_model_path): assert output.num_scheduled_tokens == {"short": 8, "medium": 8} +@pytest.mark.parametrize("admission_closed_by", ["full-running-set", "paused-new"]) +def test_queued_prefills_leave_lanes_to_running_prefills( + opt_model_path, admission_closed_by +): + """Requests that cannot be admitted must not hold prefill lanes. + + Never-scheduled requests rank first for lanes. While no queued request can + be admitted, lanes given to them left the running chunked prefills + unscheduled, so every step scheduled nothing and the engine stalled. + """ + scheduler = _create_interleaving_scheduler( + opt_model_path, + max_parallel_prefills=2, + max_num_seqs=2 if admission_closed_by == "full-running-set" else 4, + ) + for request in create_requests( + num_requests=2, num_tokens=64, req_ids=["running0", "running1"] + ): + scheduler.add_request(request) + _update(scheduler, scheduler.schedule()) + for request in create_requests( + num_requests=2, num_tokens=64, req_ids=["queued0", "queued1"] + ): + scheduler.add_request(request) + if admission_closed_by == "paused-new": + scheduler.set_pause_state(PauseState.PAUSED_NEW) + + output = scheduler.schedule() + + assert output.num_scheduled_tokens == {"running0": 8, "running1": 8} + + def test_decode_aware_reserves_one_lane_for_nearest_decode(opt_model_path): scheduler = _create_interleaving_scheduler( opt_model_path, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index e9feff11a4a1..dced44674ce6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -924,11 +924,19 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: and step_token_cap < input_budget ): step_token_cap = token_budget = input_budget + # Queued requests compete for prefill lanes only while the waiting pass + # can admit one. Otherwise lanes given to them would idle while running + # chunked prefills wait for a lane, and the step would schedule nothing. + admission_open = ( + self._pause_state == PauseState.UNPAUSED + and len(self.running) + self.num_waiting_for_streaming_input + < self.max_num_running_reqs + ) prefill_interleave_step = ( self.prefill_interleave_controller.begin_step( running=self.running, - waiting=self.waiting, - skipped_waiting=self.skipped_waiting, + waiting=self.waiting if admission_open else (), + skipped_waiting=self.skipped_waiting if admission_open else (), request_lookup=self.requests, is_local_prefill=self._request_has_local_prefill, max_parallel_prefills=self.max_parallel_prefills, From 01ad5ad2f42be5a4c61f6d9e52722ac626646048 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 10:51:05 +0000 Subject: [PATCH 337/365] docs(release): describe the prefill-lane admission stall fix Co-Authored-By: Claude Opus 5.5 --- .../vllm-prefill-lane-admission-stall.json | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 .lil/changes/vllm-prefill-lane-admission-stall.json diff --git a/.lil/changes/vllm-prefill-lane-admission-stall.json b/.lil/changes/vllm-prefill-lane-admission-stall.json new file mode 100644 index 000000000000..9dfb0932cc5c --- /dev/null +++ b/.lil/changes/vllm-prefill-lane-admission-stall.json @@ -0,0 +1,23 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-prefill-lane-admission-stall", + "category": "fix", + "summary": "Servers with --max-parallel-prefills 2 or more no longer stop generating when all running slots hold long prompts and more requests wait", + "models": [ + "all" + ], + "compatibility": "No user action required. Affected only servers started with --max-parallel-prefills 2 or more (or auto). The profiles keep 1, the default.", + "details": [ + "What was wrong: prefill lanes went to never-scheduled requests first, including queued requests that could not be admitted because every running slot (--max-num-seqs) was taken or new requests were paused. The running requests that were still prefilling long prompts got no lane.", + "Consequence: every step scheduled nothing. Running stayed at max-num-seqs with requests waiting, the GPUs idled, and generation stopped for good without an error. netwalker4370 hit it on DeepSeek-V4.1-Flash TP4 with --max-parallel-prefills 2 and 40 concurrent 150K-token prompts at --max-num-seqs 32.", + "Now queued requests compete for prefill lanes only while one of them can be admitted. Reported, with the root cause and this fix, by netwalker4370." + ], + "pull_requests": [ + 959, + 960 + ], + "authors": [ + "Local Inference Lab" + ], + "requires": [] +} From 348f4fa814b4a2d57ed7c5f8803345eb2492a11a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 13:35:23 +0000 Subject: [PATCH 338/365] feat(qwen4exp): share the mapped-host PLE table between same-host processes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit VLLM_PLE_TABLE_MEMORY=shared (or additional_config.ple_table_memory) keeps one host-RAM copy of the Qwen3.8-Flash-Next PLE n-gram table per host for every process serving the same checkpoint, instead of one cudaHostAlloc copy per process. Independent TP1 replicas on separate GPUs then hold the 26.8 GiB NVFP4 table once. The table of each TP rank lives in tmpfs files under VLLM_PLE_SHARED_TABLE_DIR (default /dev/shm/vllm-ple), keyed by the checkpoint identity and the b12x table geometry. The first process takes the key's lock, fills the files through the ordinary weight loader and publishes a manifest and READY marker once _validate_embedding_loaded proves every row loaded; a failed validation abandons the table at once. Later processes wait on the lock, check the manifest against their own geometry, map and cudaHostRegister the files, and route the table shards through the file-backed loader path so their payloads are not read. The mappings reach b12x through TableLayout.allocate_storage(host_allocator=), which needs the matching b12x change. The checkpoint identity is VLLM_PLE_SHARED_TABLE_IDENTITY when set, else the Hugging Face snapshot commit or the names, sizes and modification times of a local directory's checkpoint files. Role, directory, lock timeout and identity also come from additional_config. The module keeps a list/prune CLI. Ported from local-inference-lab/vllm#774 by René Honig, written for the Jovian Judgement layout; awsdl (carrion256) runs the same idea on Karmic Kraken as a 2xTP1 fleet. Co-authored-by: René Honig <5851246+renehonig@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- docs/design/streaming_weight_loading.md | 35 +- docs/features/quantization/b12x.md | 9 +- .../models/qwen4_exp/test_ple_shared_table.py | 726 ++++++++++ vllm/envs.py | 28 +- vllm/models/qwen4_exp/nvidia/b12x_ple.py | 205 ++- vllm/models/qwen4_exp/nvidia/model.py | 18 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 2 + .../qwen4_exp/nvidia/ple_shared_table.py | 1172 +++++++++++++++++ 8 files changed, 2139 insertions(+), 56 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_shared_table.py create mode 100644 vllm/models/qwen4_exp/nvidia/ple_shared_table.py diff --git a/docs/design/streaming_weight_loading.md b/docs/design/streaming_weight_loading.md index 0d4b9761f0df..c80423114e7c 100644 --- a/docs/design/streaming_weight_loading.md +++ b/docs/design/streaming_weight_loading.md @@ -34,6 +34,8 @@ contract below. Choose the PLE offload backing with `VLLM_PLE_TABLE_MEMORY`: - `ram`: CUDA-mapped, pinned host tables. +- `shared`: CUDA-mapped host tables in tmpfs files that every process serving + the same checkpoint on the host maps, instead of one private copy each. - `disk`: bounded `O_DIRECT` reads from checkpoint files using io_uring. Select the public policy in vLLM: @@ -49,8 +51,37 @@ are not valid environment values. When the selector is unset, `VLLM_PLE_CPU_OFFLOAD=1` retains host-RAM offload; otherwise tables remain GPU-resident. Explicit `additional_config.ple_table_memory` -overrides the environment. It accepts `ram`, `disk`, or `device`; `device` -explicitly selects GPU-resident tables. Matching vLLM and b12x packages are required. +overrides the environment. It accepts `ram`, `shared`, `disk`, or `device`; +`device` explicitly selects GPU-resident tables. Matching vLLM and b12x packages +are required. + +### Shared host tables + +With `shared`, the table of each TP rank lives in +`VLLM_PLE_SHARED_TABLE_DIR//` (default `/dev/shm/vllm-ple`), which must be a +tmpfs: the CUDA driver pins only shmem-backed file mappings. The first process +takes `.lock`, fills the files through the ordinary weight loader and +publishes them with a manifest and a `READY` marker once every row is loaded. +Later processes wait on the lock, check the manifest against their own table +geometry and checkpoint identity, then map the files and skip reading the table +shards from the checkpoint. A mismatch fails startup naming the field; there is +no fallback to a private copy. For the 26.8 GiB NVFP4 table of +Qwen3.8-Flash-Next, two TP1 replicas hold 26.8 GiB of host RAM instead of 53.6. + +| Setting | Default | Meaning | +|---|---|---| +| `VLLM_PLE_SHARED_TABLE_DIR` | `/dev/shm/vllm-ple` | tmpfs table root, created `0700` | +| `VLLM_PLE_SHARED_TABLE_ROLE` | `auto` | `auto` attaches or populates, `attach` requires a complete table, `populate` always rewrites | +| `VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S` | `3600` | seconds to wait for another populator | +| `VLLM_PLE_SHARED_TABLE_IDENTITY` | derived | checkpoint identity in the key, for example a content hash from a launcher | + +The `additional_config` keys `ple_shared_table_dir`, `ple_shared_table_role`, +`ple_shared_table_lock_timeout_s` and `ple_shared_table_identity` take precedence +over the environment. Without an explicit identity, a Hugging Face snapshot +resolves to its commit and a local directory to the names, sizes and +modification times of its checkpoint files. Every process sharing a table must +run as the same user. Tables stay in tmpfs until removed; list or prune them +with `python -m vllm.models.qwen4_exp.nvidia.ple_shared_table {list,prune}`. The io_uring reader does not map, pin, or prewarm the entire table. It deduplicates requested 4 KiB blocks, coalesces adjacent blocks into reads of at diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index 2eb9c7ecf44a..651cdf9ff420 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -230,10 +230,15 @@ memory. Disk offload requires safetensors checkpoints and supports the Explicit settings override these defaults, in descending precedence: -1. `--additional-config '{"ple_table_memory":"disk"}'` (`disk`, `ram`, or `device`). -2. `VLLM_PLE_TABLE_MEMORY` (`disk` or `ram`). +1. `--additional-config '{"ple_table_memory":"disk"}'` (`disk`, `ram`, `shared`, + or `device`). +2. `VLLM_PLE_TABLE_MEMORY` (`disk`, `ram`, or `shared`). 3. `VLLM_PLE_CPU_OFFLOAD` (`1` for CUDA-mapped host RAM, `0` for GPU memory). +`shared` keeps one host-RAM copy of the table per host for every process that +serves the same checkpoint, such as independent TP1 replicas on separate GPUs. +See [PLE table storage](../../design/streaming_weight_loading.md#ple-table-storage). + ## PCIe All-Reduce The B12X communication backend provides single-node PCIe all-reduce for SM120 diff --git a/tests/models/qwen4_exp/test_ple_shared_table.py b/tests/models/qwen4_exp/test_ple_shared_table.py new file mode 100644 index 000000000000..1f7ac6b25fcf --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_shared_table.py @@ -0,0 +1,726 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-shared Qwen4Exp PLE tables (CPU unless marked). + +Ported from the Jovian Judgement tests of local-inference-lab/vllm#774. +""" + +import json +import os +import shutil +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.qwen4_exp.nvidia import b12x_ple as b12x_ple_module +from vllm.models.qwen4_exp.nvidia import ple_shared_table as shared +from vllm.models.qwen4_exp.nvidia.ple_shared_table import ( + SharedTableConfig, + SharedTableGeometry, + SharedTableLock, + SharedTableLockTimeout, + SharedTableMismatch, + SharedTableMissing, + check_manifest, + list_tables, + open_shared_table, + prune_tables, + read_manifest, +) + +CPU = torch.device("cpu") + + +def _layout(quant_mode: str = "nvfp4_group16", **overrides) -> SimpleNamespace: + if quant_mode == "nvfp4_group16": + fields = dict( + weight_shape=(8, 4), + weight_dtype=torch.uint8, + weight_scale_shape=(8, 1), + weight_scale_dtype=torch.float8_e4m3fn, + weight_scale_2_shape=(1,), + weight_scale_2_dtype=torch.float32, + ) + elif quant_mode == "fp8_e4m3_per_tensor": + fields = dict( + weight_shape=(8, 8), + weight_dtype=torch.float8_e4m3fn, + weight_scale_shape=(1,), + weight_scale_dtype=torch.bfloat16, + weight_scale_2_shape=None, + weight_scale_2_dtype=None, + ) + else: + fields = dict( + weight_shape=(8, 8), + weight_dtype=torch.bfloat16, + weight_scale_shape=None, + weight_scale_dtype=None, + weight_scale_2_shape=None, + weight_scale_2_dtype=None, + ) + fields.update( + caps=SimpleNamespace( + quant_mode=quant_mode, + tp_size=1, + tp_rank=0, + table_memory="mapped_host", + device=CPU, + ), + padded_vocab_size=8, + shard_start=0, + shard_end=8, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _geometry(quant_mode: str = "nvfp4_group16", **kwargs) -> SharedTableGeometry: + return SharedTableGeometry.from_layout( + _layout(quant_mode), + model_path=kwargs.pop("model_path", "org/model"), + revision=kwargs.pop("revision", "abc123"), + identity=kwargs.pop("identity", "snapshot:abc123"), + embedding_dim=kwargs.pop("embedding_dim", 64), + dense_layer_ordinal=kwargs.pop("dense_layer_ordinal", 0), + ) + + +def _config(tmp_path: Path, **kwargs) -> SharedTableConfig: + kwargs.setdefault("directory", str(tmp_path)) + kwargs.setdefault("model_path", "org/model") + kwargs.setdefault("revision", "abc123") + kwargs.setdefault("identity", "snapshot:abc123") + return SharedTableConfig(**kwargs) + + +def _fill(table: shared.SharedTable) -> None: + for attribute in ("weight", "weight_scale"): + view = table.host_view(attribute) + if view is not None: + raw = view.view(torch.uint8) + raw.copy_( + torch.arange(raw.numel(), dtype=torch.int64) + .remainder(251) + .to(torch.uint8) + .reshape(raw.shape) + ) + + +def test_key_depends_on_every_geometry_field() -> None: + base = _geometry() + assert base.key == _geometry().key + assert len(base.key) == 16 + seen = {base.key} + for field, value in [ + ("revision", "other"), + ("identity", "snapshot:other"), + ("model_path", "org/other"), + ("tp_rank", 1), + ("tp_size", 2), + ("dense_layer_ordinal", 1), + ("shard_end", 4), + ("embedding_dim", 32), + ("weight_dtype", "int8"), + ("layout_version", 2), + ]: + data = base.as_dict() + data[field] = value + key = SharedTableGeometry.from_dict(data).key + assert key not in seen, field + seen.add(key) + + +@pytest.mark.parametrize( + ("quant_mode", "names"), + [ + ("nvfp4_group16", ["weight.bin", "weight_scale.bin"]), + ("fp8_e4m3_per_tensor", ["weight.bin"]), + ("bf16", ["weight.bin"]), + ], +) +def test_only_mapped_host_tensors_become_files(quant_mode, names) -> None: + geometry = _geometry(quant_mode) + assert [spec.name for spec in geometry.files] == names + assert geometry.nbytes == sum(spec.nbytes for spec in geometry.files) + + +def test_populate_then_attach_round_trip(tmp_path: Path) -> None: + geometry = _geometry() + populator = open_shared_table(geometry, _config(tmp_path), device=CPU) + assert populator.populating and not populator.attached + table_dir = tmp_path / geometry.key + assert (table_dir / "weight.bin").stat().st_size == geometry.files[0].nbytes + assert not (table_dir / "READY").exists() + assert list_tables(tmp_path)[0].state == "populating" + + _fill(populator) + populator.publish() + assert (table_dir / "READY").exists() + manifest = read_manifest(table_dir) + assert manifest is not None + assert manifest["key"] == geometry.key + assert manifest["geometry"] == geometry.as_dict() + assert manifest["files"] == {"weight.bin": 32, "weight_scale.bin": 8} + assert manifest["creator"]["pid"] + assert list_tables(tmp_path)[0].state == "ready" + + follower = open_shared_table(geometry, _config(tmp_path), device=CPU) + assert follower.attached and not follower.populating + for attribute in ("weight", "weight_scale"): + torch.testing.assert_close( + follower.host_view(attribute).view(torch.uint8), + populator.host_view(attribute).view(torch.uint8), + ) + assert follower.device_view(attribute) is follower.host_view(attribute) + assert follower.nbytes == populator.nbytes == 40 + follower.close() + populator.close() + assert (table_dir / "READY").exists() + + +def test_attach_role_fails_fast_without_a_table(tmp_path: Path) -> None: + with pytest.raises(SharedTableMissing, match="attach forbids populating"): + open_shared_table(_geometry(), _config(tmp_path, role="attach"), device=CPU) + assert list_tables(tmp_path) == [] + + +def test_manifest_mismatch_names_the_field(tmp_path: Path) -> None: + geometry = _geometry() + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + _fill(table) + table.publish() + table.close() + manifest_path = tmp_path / geometry.key / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["geometry"]["shard_end"] = 4 + manifest_path.write_text(json.dumps(manifest)) + + with pytest.raises(SharedTableMismatch, match="shard_end expected 8") as info: + open_shared_table(geometry, _config(tmp_path, role="attach"), device=CPU) + assert info.value.field == "shard_end" + # auto never silently repopulates over a READY table that disagrees. + with pytest.raises(SharedTableMismatch): + open_shared_table(geometry, _config(tmp_path), device=CPU) + + +def test_manifest_check_covers_file_sizes_and_format(tmp_path: Path) -> None: + geometry = _geometry() + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + _fill(table) + table.publish() + table.close() + table_dir = tmp_path / geometry.key + manifest = json.loads((table_dir / "manifest.json").read_text()) + check_manifest(manifest, geometry, table_dir) + + with pytest.raises(SharedTableMismatch, match="format"): + check_manifest({**manifest, "format": "other"}, geometry, table_dir) + with open(table_dir / "weight.bin", "ab") as handle: + handle.write(b"\0") + with pytest.raises(SharedTableMismatch, match=r"files\.weight\.bin expected 32"): + check_manifest(manifest, geometry, table_dir) + + +def test_stale_directory_without_ready_is_replaced(tmp_path: Path) -> None: + geometry = _geometry() + stale = tmp_path / geometry.key + stale.mkdir() + (stale / "weight.bin").write_bytes(b"garbage") + (stale / "leftover").write_text("x") + + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + assert table.populating + assert not (stale / "leftover").exists() + assert (stale / "weight.bin").stat().st_size == 32 + table.close() + # Closing an unpublished populator removes its partial table. + assert not stale.exists() + assert list_tables(tmp_path) == [] + + +def test_populate_role_always_rewrites(tmp_path: Path) -> None: + geometry = _geometry() + first = open_shared_table(geometry, _config(tmp_path), device=CPU) + _fill(first) + first.publish() + first.close() + created = read_manifest(tmp_path / geometry.key)["created_at"] + + second = open_shared_table(geometry, _config(tmp_path, role="populate"), device=CPU) + assert second.populating + assert not (tmp_path / geometry.key / "READY").exists() + _fill(second) + second.publish() + assert read_manifest(tmp_path / geometry.key)["created_at"] >= created + second.close() + + +def test_lock_timeout_is_reported(tmp_path: Path) -> None: + geometry = _geometry() + holder = SharedTableLock(shared.lock_path(tmp_path, geometry.key), 1) + holder.acquire() + try: + with pytest.raises(SharedTableLockTimeout, match="timed out after 0s"): + open_shared_table(geometry, _config(tmp_path, lock_timeout_s=0), device=CPU) + finally: + holder.release() + table = open_shared_table(geometry, _config(tmp_path, lock_timeout_s=0), device=CPU) + assert table.populating + table.close() + + +def test_populator_holds_the_lock_until_publish(tmp_path: Path) -> None: + geometry = _geometry() + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + probe = SharedTableLock(shared.lock_path(tmp_path, geometry.key), 0) + with pytest.raises(SharedTableLockTimeout): + probe.acquire() + _fill(table) + table.publish() + probe.acquire() + probe.release() + table.close() + + +def test_config_rejects_unknown_role(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="ple_shared_table_role"): + _config(tmp_path, role="maybe") + + +def test_cli_list_and_prune(tmp_path: Path, capsys) -> None: + keep = _geometry() + stale = _geometry(revision="old") + busy = _geometry(revision="busy") + for geometry in (keep, stale): + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + _fill(table) + table.publish() + table.close() + populating = open_shared_table(busy, _config(tmp_path), device=CPU) + + assert shared.main(["--dir", str(tmp_path), "list"]) == 0 + lines = capsys.readouterr().out.strip().splitlines() + assert len(lines) == 3 + states = {line.split()[0]: line.split()[1] for line in lines} + assert states == {keep.key: "ready", stale.key: "ready", busy.key: "populating"} + assert any("revision=old" in line for line in lines) + + assert shared.main(["--dir", str(tmp_path), "list", "--json"]) == 0 + entries = json.loads(capsys.readouterr().out) + assert {entry["key"] for entry in entries} == {keep.key, stale.key, busy.key} + + assert ( + shared.main(["--dir", str(tmp_path), "prune", "--keep", keep.key, "--dry-run"]) + == 0 + ) + assert capsys.readouterr().out.strip() == f"would remove {stale.key}" + assert (tmp_path / stale.key).exists() + + assert prune_tables(tmp_path, keep=[keep.key]) == [stale.key] + assert not (tmp_path / stale.key).exists() + # flock binds to an inode: the lock file must survive the prune so a + # concurrent opener never ends up holding a different inode. + assert shared.lock_path(tmp_path, stale.key).exists() + assert (tmp_path / keep.key / "READY").exists() + # A table being populated is never pruned. + assert (tmp_path / busy.key).exists() + populating.close() + assert shared.main(["--dir", str(tmp_path / "missing"), "list"]) == 0 + + +def test_abort_releases_lock_and_removes_partial_table(tmp_path: Path) -> None: + geometry = _geometry() + table = open_shared_table(geometry, _config(tmp_path), device=CPU) + view = table.host_view("weight") + view.fill_(9) + table.abort() + assert not (tmp_path / geometry.key).exists() + # The mapping stays usable for a model that still references it. + assert int(view.view(torch.uint8)[0, 0]) == 9 + # Another opener is not blocked and repopulates. + follower = open_shared_table( + geometry, _config(tmp_path, lock_timeout_s=0), device=CPU + ) + assert follower.populating + follower.close() + table.close() + + +def test_validation_failure_aborts_populate() -> None: + model = _LoaderAudit(attached=False) + model.ngram_embedding.aborted = 0 + model.ngram_embedding.abort_shared_table = lambda: setattr( + model.ngram_embedding, "aborted", model.ngram_embedding.aborted + 1 + ) + model.load_weights(_shard_weights(meta=False)[:4]) + with pytest.raises(ValueError, match="do not cover the local table"): + model._validate_embedding_loaded() + assert model.ngram_embedding.aborted == 1 + assert model.ngram_embedding.published == 0 + + +def test_root_is_created_private_and_unsafe_paths_are_rejected(tmp_path: Path) -> None: + root = tmp_path / "tables" + geometry = _geometry() + table = open_shared_table( + geometry, _config(tmp_path, directory=str(root)), device=CPU + ) + assert oct(root.stat().st_mode & 0o777) == "0o700" + assert oct((root / geometry.key).stat().st_mode & 0o777) == "0o700" + assert oct((root / geometry.key / "weight.bin").stat().st_mode & 0o777) == "0o600" + _fill(table) + table.publish() + table.close() + + # A group/world-writable table directory is not trusted even with READY. + (root / geometry.key).chmod(0o777) + with pytest.raises(shared.SharedTableError, match="writable by group or others"): + open_shared_table(geometry, _config(tmp_path, directory=str(root)), device=CPU) + (root / geometry.key).chmod(0o700) + open_shared_table( + geometry, _config(tmp_path, directory=str(root)), device=CPU + ).close() + + # A symlinked root is refused. + link = tmp_path / "link" + link.symlink_to(root) + with pytest.raises(shared.SharedTableError, match="symlink"): + open_shared_table(geometry, _config(tmp_path, directory=str(link)), device=CPU) + + # A world-writable root is refused. + other = tmp_path / "loose" + other.mkdir(mode=0o777) + other.chmod(0o777) + with pytest.raises(shared.SharedTableError, match="writable by group or others"): + open_shared_table(geometry, _config(tmp_path, directory=str(other)), device=CPU) + + +@dataclass(kw_only=True) +class _FakeTableStorage: + weight: torch.Tensor + weight_scale: torch.Tensor | None + weight_scale_2: torch.Tensor | None + weight_load_view: torch.Tensor + weight_scale_load_view: torch.Tensor | None + weight_scale_2_load_view: torch.Tensor | None + mapped_host_nbytes: int + _mapped_allocations: tuple + + +def _fake_allocate_storage(layout: SimpleNamespace, host_allocator): + """Mirror b12x allocate_storage for NVFP4 mapped-host tables on CPU.""" + weight = host_allocator("weight", layout.weight_shape, layout.weight_dtype) + scale = host_allocator( + "weight_scale", layout.weight_scale_shape, layout.weight_scale_dtype + ) + scale_2 = torch.empty(layout.weight_scale_2_shape, dtype=torch.float32) + return _FakeTableStorage( + weight=weight.device_view, + weight_scale=scale.device_view, + weight_scale_2=scale_2, + weight_load_view=weight.host_view, + weight_scale_load_view=scale.host_view, + weight_scale_2_load_view=scale_2, + mapped_host_nbytes=weight.nbytes + scale.nbytes, + _mapped_allocations=(weight, scale), + ) + + +def _storage_layout() -> SimpleNamespace: + layout = _layout() + layout.allocate_storage = lambda host_allocator=None: _fake_allocate_storage( + layout, host_allocator + ) + return layout + + +def _storage(config: SharedTableConfig): + return b12x_ple_module._NGramEmbeddingStorage( + _storage_layout(), 2, shared_table=config, embedding_dim=64 + ) + + +def test_storage_populates_then_attaches(tmp_path: Path) -> None: + config = _config(tmp_path) + + populator = _storage(config) + assert not populator.shared_table_attached + assert isinstance(populator.weight, nn.Parameter) + assert populator.weight.shape == (8, 4) and populator.weight.dtype == torch.uint8 + assert populator.mapped_host_nbytes == 40 + assert populator.weight_load_view.data_ptr() == populator.weight.data_ptr() + populator.weight_load_view.fill_(7) + populator.weight_scale_load_view.view(torch.uint8).fill_(3) + populator.publish_shared_table() + + follower = _storage(config) + assert follower.shared_table_attached + assert torch.equal(follower.weight, torch.full((8, 4), 7, dtype=torch.uint8)) + assert torch.equal( + follower.weight_scale.view(torch.uint8), + torch.full((8, 1), 3, dtype=torch.uint8), + ) + # Scale 2 stays a private tensor loaded from the checkpoint. + assert follower.weight_scale_2.data_ptr() != populator.weight_scale_2.data_ptr() + follower.publish_shared_table() # no-op for attachers + for storage in (populator, follower): + for allocation in storage._table_storage._mapped_allocations: + allocation.close() + + +def test_storage_closes_the_table_when_b12x_rejects_it(tmp_path: Path) -> None: + layout = _storage_layout() + + def reject(host_allocator=None): + raise ValueError("b12x rejected the region") + + layout.allocate_storage = reject + with pytest.raises(ValueError, match="b12x rejected the region"): + b12x_ple_module._NGramEmbeddingStorage( + layout, 2, shared_table=_config(tmp_path), embedding_dim=64 + ) + # The aborted populate left no table and released its lock. + assert list_tables(tmp_path) == [] + _storage(_config(tmp_path, role="populate")) + + +def test_shared_policy_maps_to_mapped_host(monkeypatch) -> None: + monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) + monkeypatch.setenv("VLLM_PLE_TABLE_MEMORY", "shared") + assert b12x_ple_module._resolve_ple_table_memory(None) == "mapped_host" + assert ( + b12x_ple_module._resolve_ple_table_memory({"ple_table_memory": "ram"}) + == "mapped_host" + ) + with pytest.raises(ValueError, match="'shared'"): + b12x_ple_module._resolve_ple_table_memory({"ple_table_memory": "shm"}) + + +def _vllm_config(additional_config) -> SimpleNamespace: + return SimpleNamespace( + additional_config=additional_config, + model_config=SimpleNamespace(model="org/model", revision="rev1"), + ) + + +def test_shared_table_settings_from_env_and_additional_config(monkeypatch) -> None: + monkeypatch.setattr( + b12x_ple_module, "checkpoint_identity", lambda model, revision: "snapshot:c1" + ) + monkeypatch.setenv("VLLM_PLE_TABLE_MEMORY", "ram") + assert b12x_ple_module._resolve_ple_shared_table(_vllm_config(None)) is None + + monkeypatch.setenv("VLLM_PLE_TABLE_MEMORY", "shared") + for name in ("DIR", "ROLE", "LOCK_TIMEOUT_S", "IDENTITY"): + monkeypatch.delenv(f"VLLM_PLE_SHARED_TABLE_{name}", raising=False) + config = b12x_ple_module._resolve_ple_shared_table(_vllm_config({})) + assert config == SharedTableConfig( + directory="/dev/shm/vllm-ple", + role="auto", + lock_timeout_s=3600.0, + model_path="org/model", + revision="rev1", + identity="snapshot:c1", + ) + + monkeypatch.setenv("VLLM_PLE_SHARED_TABLE_DIR", "/mnt/tables") + monkeypatch.setenv("VLLM_PLE_SHARED_TABLE_ROLE", "attach") + monkeypatch.setenv("VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S", "5") + monkeypatch.setenv("VLLM_PLE_SHARED_TABLE_IDENTITY", "sha256:abc") + config = b12x_ple_module._resolve_ple_shared_table(_vllm_config(None)) + assert ( + config.directory, + config.role, + config.lock_timeout_s, + config.identity, + ) == ("/mnt/tables", "attach", 5.0, "sha256:abc") + + config = b12x_ple_module._resolve_ple_shared_table( + _vllm_config( + { + "ple_table_memory": "shared", + "ple_shared_table_dir": "/tables", + "ple_shared_table_role": "populate", + "ple_shared_table_lock_timeout_s": 1, + "ple_shared_table_identity": "sha256:def", + } + ) + ) + assert ( + config.directory, + config.role, + config.lock_timeout_s, + config.identity, + ) == ("/tables", "populate", 1.0, "sha256:def") + with pytest.raises(ValueError, match="ple_shared_table_role"): + b12x_ple_module._resolve_ple_shared_table( + _vllm_config({"ple_table_memory": "shared", "ple_shared_table_role": "x"}) + ) + + +def test_checkpoint_identity_tracks_the_checkpoint_files(tmp_path: Path) -> None: + snapshot = tmp_path / "snapshots" / "0123abcd" + snapshot.mkdir(parents=True) + assert shared.checkpoint_identity(str(snapshot), None) == "snapshot:0123abcd" + + local = tmp_path / "local" + local.mkdir() + shard = local / "model-00001.safetensors" + shard.write_bytes(b"a" * 8) + (local / "README.md").write_text("ignored") + first = shared.checkpoint_identity(str(local), None) + assert first.startswith("files:") + (local / "README.md").write_text("still ignored") + assert shared.checkpoint_identity(str(local), None) == first + shard.write_bytes(b"b" * 9) + assert shared.checkpoint_identity(str(local), None) != first + + assert ( + shared.checkpoint_identity("org/not-cached", "rev") == "name:org/not-cached@rev" + ) + + +class _LoaderAudit(b12x_ple_module.B12xNGramEmbedding): + """A minimal NVFP4 embedding whose table may be attached to a shared one.""" + + def __init__(self, attached: bool) -> None: + nn.Module.__init__(self) + self._quant_mode = "nvfp4_group16" + self._embedding_load_ranges = set() + self._scale_load_ranges = set() + self._weight_scale_loaded = False + self._weight_scale_2_loaded = False + self._embedding_validated = False + self.split_ngram_parts = 4 + self._table_layout = SimpleNamespace( + padded_vocab_size=8, + shard_start=0, + shard_end=8, + weight_shape=(8, 8), + weight_dtype=torch.uint8, + weight_scale_shape=(8, 1), + weight_scale_dtype=torch.float8_e4m3fn, + weight_scale_2_shape=(1,), + weight_scale_2_dtype=torch.float32, + ) + self.register_buffer("layer_multipliers", torch.tensor([11, 13])) + self.register_buffer("ngram_heads_offsets", torch.tensor([0, 4])) + self.register_buffer("ngram_heads_vocab_sizes", torch.tensor([4, 4])) + embedding = nn.Module() + embedding.shared_table_attached = attached + embedding.published = 0 + embedding.publish_shared_table = lambda: setattr( + embedding, "published", embedding.published + 1 + ) + for name, shape, dtype in ( + ("weight", (8, 8), torch.uint8), + ("weight_scale", (8, 1), torch.float8_e4m3fn), + ("weight_scale_2", (1,), torch.float32), + ): + embedding.register_parameter( + name, nn.Parameter(torch.zeros(shape, dtype=dtype), requires_grad=False) + ) + self.add_module("ngram_embedding", embedding) + + +def _shard_weights(meta: bool): + device = "meta" if meta else "cpu" + weights = [] + for shard in range(4): + weights.append( + ( + f"ngram_embedding.shard_{shard}.weight", + torch.full((2, 8), shard + 1, dtype=torch.uint8, device=device), + ) + ) + weights.append( + ( + f"ngram_embedding.shard_{shard}.weight_scale", + torch.ones((2, 1), dtype=torch.float8_e4m3fn, device=device), + ) + ) + weights.append(("ngram_embedding.weight_scale_2", torch.tensor(0.5))) + weights.append(("layer_multipliers", torch.tensor([11, 13]))) + return weights + + +def test_attached_loader_skips_shard_copies_but_checks_geometry() -> None: + """File-backed (meta) shards of an attached table are never read, yet + their geometry is checked and their rows count toward coverage.""" + model = _LoaderAudit(attached=True) + loaded = model.load_weights(_shard_weights(meta=True)) + assert loaded == { + "ngram_embedding.weight", + "ngram_embedding.weight_scale", + "ngram_embedding.weight_scale_2", + "layer_multipliers", + } + assert torch.equal( + model.ngram_embedding.weight, torch.zeros(8, 8, dtype=torch.uint8) + ) + assert model.ngram_embedding.weight_scale_2.item() == 0.5 + model._validate_embedding_loaded() + assert model.ngram_embedding.published == 1 + model._validate_embedding_loaded() + assert model.ngram_embedding.published == 1 + + with pytest.raises(ValueError, match="shape mismatch for PLE shard 1 weight"): + _LoaderAudit(attached=True).load_weights( + [("ngram_embedding.shard_1.weight", torch.zeros(3, 8, dtype=torch.uint8))] + ) + + +def test_populating_loader_copies_shards_and_publishes_once_complete() -> None: + model = _LoaderAudit(attached=False) + model.load_weights(_shard_weights(meta=False)) + expected = torch.cat( + [torch.full((2, 8), shard + 1, dtype=torch.uint8) for shard in range(4)] + ) + assert torch.equal(model.ngram_embedding.weight, expected) + assert model.ngram_embedding.published == 0 + model._validate_embedding_loaded() + assert model.ngram_embedding.published == 1 + + partial = _LoaderAudit(attached=False) + partial.load_weights(_shard_weights(meta=False)[:4]) + with pytest.raises(ValueError, match="do not cover the local table"): + partial._validate_embedding_loaded() + assert partial.ngram_embedding.published == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_registered_mapping_exposes_the_device_alias() -> None: + import tempfile + + from cuda.bindings import runtime as cudart + + # The driver only pins tmpfs-backed mappings; pytest's tmp_path may not be. + if not os.path.isdir("/dev/shm"): + pytest.skip("needs a tmpfs at /dev/shm") + tmp_path = Path(tempfile.mkdtemp(prefix="vllm-ple-test-", dir="/dev/shm")) + geometry = _geometry() + device = torch.device("cuda", torch.cuda.current_device()) + populator = open_shared_table(geometry, _config(tmp_path), device=device) + _fill(populator) + populator.publish() + follower = open_shared_table(geometry, _config(tmp_path), device=device) + for table in (populator, follower): + for attribute in ("weight", "weight_scale"): + host = table.host_view(attribute) + dev = table.device_view(attribute) + assert dev.device == device and host.device.type == "cpu" + error, attributes = cudart.cudaPointerGetAttributes(dev.data_ptr()) + assert error == cudart.cudaError_t.cudaSuccess + assert attributes.type == cudart.cudaMemoryType.cudaMemoryTypeHost + assert int(attributes.devicePointer) == dev.data_ptr() + torch.testing.assert_close( + dev.cpu().view(torch.uint8), host.view(torch.uint8) + ) + follower.close() + populator.close() + shutil.rmtree(tmp_path) diff --git a/vllm/envs.py b/vllm/envs.py index 61c6282fd1fe..4906740a2806 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -225,7 +225,11 @@ VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 VLLM_B12X_PAGED_DECODE: Literal["auto", "0", "1"] = "auto" - VLLM_PLE_TABLE_MEMORY: Literal["ram", "disk"] | None = None + VLLM_PLE_TABLE_MEMORY: Literal["ram", "disk", "shared"] | None = None + VLLM_PLE_SHARED_TABLE_DIR: str = "/dev/shm/vllm-ple" + VLLM_PLE_SHARED_TABLE_ROLE: Literal["auto", "populate", "attach"] = "auto" + VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S: float = 3600.0 + VLLM_PLE_SHARED_TABLE_IDENTITY: str | None = None VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False VLLM_TPU_USING_PATHWAYS: bool = False @@ -1892,10 +1896,30 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_B12X_PAGED_DECODE", "auto", ["auto", "0", "1"] ), # Qwen3.8-Flash-Next PLE offload policy, resolved by vLLM for b12x. + # "shared" keeps one host-RAM table per host in VLLM_PLE_SHARED_TABLE_DIR + # for every process serving the same checkpoint, such as TP1 replicas. "VLLM_PLE_TABLE_MEMORY": env_with_choices( "VLLM_PLE_TABLE_MEMORY", None, - ["ram", "disk"], + ["ram", "disk", "shared"], + ), + # tmpfs directory of the shared PLE tables. + "VLLM_PLE_SHARED_TABLE_DIR": lambda: os.getenv( + "VLLM_PLE_SHARED_TABLE_DIR", "/dev/shm/vllm-ple" + ), + # "auto" attaches to a complete table or populates it, "attach" requires + # one, "populate" always rewrites it. + "VLLM_PLE_SHARED_TABLE_ROLE": env_with_choices( + "VLLM_PLE_SHARED_TABLE_ROLE", "auto", ["auto", "populate", "attach"] + ), + # Seconds to wait while another process populates the shared table. + "VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S": lambda: float( + os.getenv("VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S", "3600") + ), + # Checkpoint identity in the shared table key, for example a content hash + # supplied by a launcher; derived from the checkpoint location when unset. + "VLLM_PLE_SHARED_TABLE_IDENTITY": lambda: os.getenv( + "VLLM_PLE_SHARED_TABLE_IDENTITY" ), # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( diff --git a/vllm/models/qwen4_exp/nvidia/b12x_ple.py b/vllm/models/qwen4_exp/nvidia/b12x_ple.py index f8e37466797f..a4345079c42e 100644 --- a/vllm/models/qwen4_exp/nvidia/b12x_ple.py +++ b/vllm/models/qwen4_exp/nvidia/b12x_ple.py @@ -36,6 +36,12 @@ from vllm.utils.torch_utils import direct_register_custom_op from ..config import Qwen4ExpTextConfig +from .ple_shared_table import ( + SharedTableConfig, + SharedTableGeometry, + checkpoint_identity, + open_shared_table, +) logger = init_logger(__name__) @@ -94,7 +100,7 @@ def _resolve_ple_table_memory( if envs.is_set("VLLM_PLE_CPU_OFFLOAD"): return "mapped_host" if envs.VLLM_PLE_CPU_OFFLOAD else "device" return "io_uring" if embedding_dtype == "bfloat16" else "device" - if table_memory == "ram": + if table_memory in ("ram", "shared"): return "mapped_host" if table_memory == "disk": return "io_uring" @@ -102,7 +108,37 @@ def _resolve_ple_table_memory( return "device" raise ValueError( "additional_config.ple_table_memory must be 'device', " - f"'ram', or 'disk', got {table_memory!r}" + f"'ram', 'shared', or 'disk', got {table_memory!r}" + ) + + +def _resolve_ple_shared_table(vllm_config: Any) -> SharedTableConfig | None: + """Return the host-shared table settings when ``shared`` is selected.""" + additional_config = vllm_config.additional_config + config = additional_config if isinstance(additional_config, dict) else {} + table_memory = config.get("ple_table_memory", envs.VLLM_PLE_TABLE_MEMORY) + if table_memory != "shared": + return None + model_config = vllm_config.model_config + identity = config.get( + "ple_shared_table_identity", envs.VLLM_PLE_SHARED_TABLE_IDENTITY + ) + return SharedTableConfig( + directory=str( + config.get("ple_shared_table_dir", envs.VLLM_PLE_SHARED_TABLE_DIR) + ), + role=str(config.get("ple_shared_table_role", envs.VLLM_PLE_SHARED_TABLE_ROLE)), + lock_timeout_s=float( + config.get( + "ple_shared_table_lock_timeout_s", + envs.VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S, + ) + ), + model_path=str(model_config.model), + revision=model_config.revision, + identity=str(identity) + if identity + else checkpoint_identity(str(model_config.model), model_config.revision), ) @@ -128,10 +164,19 @@ def _copy_embedding_shard( class _NGramEmbeddingStorage(nn.Module): - def __init__(self, layout: Any, shard_rows: int) -> None: + def __init__( + self, + layout: Any, + shard_rows: int, + *, + shared_table: SharedTableConfig | None = None, + embedding_dim: int = 0, + dense_layer_ordinal: int = 0, + ) -> None: super().__init__() self.disk_table = None self._table_storage = None + self._shared_table = None if layout.caps.table_memory == "io_uring": api = _b12x_module("ple_embedding") self.disk_table = api.DiskTable(layout, shard_rows) @@ -147,6 +192,29 @@ def __init__(self, layout: Any, shard_rows: int) -> None: if shape == (1,) else None ) + elif shared_table is not None: + geometry = SharedTableGeometry.from_layout( + layout, + model_path=shared_table.model_path, + revision=shared_table.revision, + identity=shared_table.identity, + embedding_dim=embedding_dim, + dense_layer_ordinal=dense_layer_ordinal, + ) + self._shared_table = open_shared_table( + geometry, shared_table, device=layout.caps.device + ) + try: + self._table_storage = layout.allocate_storage( + host_allocator=self._shared_table.region + ) + except BaseException: + self._shared_table.close() + raise + tensors = { + name: getattr(self._table_storage, name) + for name in ("weight", "weight_scale", "weight_scale_2") + } else: self._table_storage = layout.allocate_storage() tensors = { @@ -181,6 +249,24 @@ def weight_scale_2_load_view(self) -> torch.Tensor | None: def mapped_host_nbytes(self) -> int: return int(self._table_storage.mapped_host_nbytes) if self._table_storage else 0 + @property + def shared_table_attached(self) -> bool: + """Whether another process already filled every row of this table.""" + return self._shared_table is not None and self._shared_table.attached + + def abort_shared_table(self) -> None: + """Give up a shared table this process was populating.""" + if self._shared_table is not None and self._shared_table.populating: + self._shared_table.abort() + + def publish_shared_table(self) -> None: + """Mark a populated shared table complete for other processes.""" + if self._shared_table is None or not self._shared_table.populating: + return + flush_weight_transfers() + torch.cuda.synchronize() + self._shared_table.publish() + class B12xNGramEmbedding(nn.Module): """Prime-hashed learned n-gram embedding with fixed b12x storage.""" @@ -205,6 +291,7 @@ def __init__( prefix: str, dtype: torch.dtype, table_memory: str, + shared_table: SharedTableConfig | None = None, ) -> None: super().__init__() self.embedding_dim = int(embedding_dim) @@ -268,12 +355,19 @@ def __init__( shard_rows = ( self._table_layout.padded_vocab_size + self.split_ngram_parts - 1 ) // self.split_ngram_parts - self.ngram_embedding = _NGramEmbeddingStorage(self._table_layout, shard_rows) + self.ngram_embedding = _NGramEmbeddingStorage( + self._table_layout, + shard_rows, + shared_table=shared_table, + embedding_dim=self.embedding_dim, + dense_layer_ordinal=int(ple_dense_layer_id), + ) if self.ngram_embedding.mapped_host_nbytes: logger.info( - "Using %.2f GiB of CUDA-mapped host memory for this TP rank's " + "Using %.2f GiB of %s CUDA-mapped host memory for this TP rank's " "PLE table", self.ngram_embedding.mapped_host_nbytes / (1 << 30), + "shared" if shared_table is not None else "private", ) scratch_spec = self._table_layout.scratch_specs()[0] self.register_buffer( @@ -546,49 +640,22 @@ def prepare_dummy_output(self, num_tokens: int) -> None: def _validate_embedding_loaded(self) -> None: if self._embedding_validated: return - - covered_until = self._table_layout.shard_start - for start, end in sorted(self._embedding_load_ranges): - if start > covered_until: - raise ValueError( - "PLE embedding shards do not cover the local table: " - f"expected row {covered_until}, got {start}" - ) - covered_until = max(covered_until, end) - if covered_until != self._table_layout.shard_end: - raise ValueError( - "PLE embedding shards do not cover the local table: " - f"stopped at row {covered_until}, " - f"expected {self._table_layout.shard_end}" + try: + _require_embedding_coverage(self) + except ValueError: + # Waiting replicas must not attach to a table this process cannot fill. + abort_shared_table = getattr( + self.ngram_embedding, "abort_shared_table", None ) - - if self._table_layout.weight_scale_shape is not None: - if self._quant_mode == "fp8_e4m3_per_tensor": - if not self._weight_scale_loaded: - raise ValueError( - "FP8 PLE embedding checkpoint is missing weight_scale" - ) - else: - scale_covered_until = self._table_layout.shard_start - for start, end in sorted(self._scale_load_ranges): - if start > scale_covered_until: - raise ValueError( - "NVFP4 PLE scale shards do not cover the local table: " - f"expected row {scale_covered_until}, got {start}" - ) - scale_covered_until = max(scale_covered_until, end) - if scale_covered_until != self._table_layout.shard_end: - raise ValueError( - "NVFP4 PLE scale shards do not cover the local table: " - f"stopped at row {scale_covered_until}, expected " - f"{self._table_layout.shard_end}" - ) - if ( - self._table_layout.weight_scale_2_shape is not None - and not self._weight_scale_2_loaded - ): - raise ValueError("NVFP4 PLE embedding checkpoint is missing weight_scale_2") + if abort_shared_table is not None: + abort_shared_table() + raise self._embedding_validated = True + publish_shared_table = getattr( + self.ngram_embedding, "publish_shared_table", None + ) + if publish_shared_table is not None: + publish_shared_table() def forward( self, @@ -797,6 +864,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: source.offset, scale=suffix == "weight_scale", ) + elif getattr(embedding, "shared_table_attached", False): + # The attached table already holds every row of this shard. + pass else: parameter = getattr(embedding, suffix) destination = getattr( @@ -837,6 +907,49 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded +def _require_embedding_coverage(embedding: B12xNGramEmbedding) -> None: + """Raise when the loaded shards or scales do not cover the local table.""" + covered_until = embedding._table_layout.shard_start + for start, end in sorted(embedding._embedding_load_ranges): + if start > covered_until: + raise ValueError( + "PLE embedding shards do not cover the local table: " + f"expected row {covered_until}, got {start}" + ) + covered_until = max(covered_until, end) + if covered_until != embedding._table_layout.shard_end: + raise ValueError( + "PLE embedding shards do not cover the local table: " + f"stopped at row {covered_until}, " + f"expected {embedding._table_layout.shard_end}" + ) + + if embedding._table_layout.weight_scale_shape is not None: + if embedding._quant_mode == "fp8_e4m3_per_tensor": + if not embedding._weight_scale_loaded: + raise ValueError("FP8 PLE embedding checkpoint is missing weight_scale") + else: + scale_covered_until = embedding._table_layout.shard_start + for start, end in sorted(embedding._scale_load_ranges): + if start > scale_covered_until: + raise ValueError( + "NVFP4 PLE scale shards do not cover the local table: " + f"expected row {scale_covered_until}, got {start}" + ) + scale_covered_until = max(scale_covered_until, end) + if scale_covered_until != embedding._table_layout.shard_end: + raise ValueError( + "NVFP4 PLE scale shards do not cover the local table: " + f"stopped at row {scale_covered_until}, expected " + f"{embedding._table_layout.shard_end}" + ) + if ( + embedding._table_layout.weight_scale_2_shape is not None + and not embedding._weight_scale_2_loaded + ): + raise ValueError("NVFP4 PLE embedding checkpoint is missing weight_scale_2") + + def _ple_embedding_op( input_ids: torch.Tensor, query_start_loc: torch.Tensor, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index e0991cfcc5c7..5646fa1abd7c 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -77,7 +77,7 @@ from vllm.v1.kv_cache_interface import MambaSpec from ..config import Qwen4ExpConfig -from .b12x_ple import _resolve_ple_table_memory +from .b12x_ple import B12xNGramEmbedding, _resolve_ple_table_memory from .b12x_qsa import Qwen4ExpQSAAttention as B12xQSAAttention from .backend import uses_b12x from .hyperconnection import ( @@ -888,15 +888,25 @@ def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: @property def checkpoint_file_weight_filter(self) -> Callable[[str], bool] | None: + if not (uses_b12x(self.vllm_config) and self.config.ple_layer_ids): + return None if ( - uses_b12x(self.vllm_config) - and self.config.ple_layer_ids - and _resolve_ple_table_memory( + _resolve_ple_table_memory( self.vllm_config.additional_config, self.config.ple_embedding_dtype ) == "io_uring" ): return _is_file_backed_ple_weight + embeddings = [ + module + for module in self.modules() + if isinstance(module, B12xNGramEmbedding) + ] + # Attached shared tables already hold every shard row: skip reading them. + if embeddings and all( + embedding.ngram_embedding.shared_table_attached for embedding in embeddings + ): + return _is_file_backed_ple_weight return None diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 87f3cae6ccc9..ff641d112c7f 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -46,6 +46,7 @@ B12xNGramEmbedding, _b12x_module, _register_ple_compilation_context, + _resolve_ple_shared_table, _resolve_ple_table_memory, ) from .backend import uses_b12x @@ -145,6 +146,7 @@ def __init__( _resolve_ple_table_memory( vllm_config.additional_config, config.ple_embedding_dtype ), + shared_table=_resolve_ple_shared_table(vllm_config), ) else: self.ple_embedding = Qwen4ExpNGramEmbedding( diff --git a/vllm/models/qwen4_exp/nvidia/ple_shared_table.py b/vllm/models/qwen4_exp/nvidia/ple_shared_table.py new file mode 100644 index 000000000000..2ad6940a11cb --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ple_shared_table.py @@ -0,0 +1,1172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-shared PLE table files for independent Qwen4Exp processes. + +Several TP1 replicas of Qwen3.8-Flash-Next on one host can serve one +CPU-resident copy of the n-gram table instead of one ``cudaHostAlloc`` copy per +process (26.8 GiB for the NVFP4 checkpoint). The table lives in a tmpfs +directory as contiguous files in the b12x storage layout. The first process to +take the directory lock writes them from the checkpoint through the ordinary +weight-loading path; later processes ``mmap(MAP_SHARED)`` the same files and +``cudaHostRegister`` the mapping, which pins each page once no matter how many +processes map it. The mappings reach b12x through +``TableLayout.allocate_storage(host_allocator=...)``. + +The design follows local-inference-lab/vllm#774 by René Honig, written for the +Jovian Judgement layout of this model. + +Layout under ````:: + + /manifest.json geometry and provenance; written before READY + /weight.bin plan.weight_shape rows, plan.weight_dtype, exact size + /weight_scale.bin only for mapped-host scales (nvfp4_group16) + /READY empty marker created last + .lock flock target, separate so / can be replaced + +The directory must be tmpfs (or another shmem-backed file system): the CUDA +driver refuses to pin ``MAP_SHARED`` pages of ordinary file systems. It is +private to one user: the root is created ``0700``, and a root, table +directory or file that is a symlink, owned by another user, or writable by +group/others is rejected, so every replica sharing a table must run as the +same user. + +The key identifies a checkpoint by ``model_path``, ``revision`` and an identity +string: the Hugging Face snapshot commit for a cached repository, or the names, +sizes and modification times of a local directory's checkpoint files, unless an +explicit identity is configured. Contents changed behind an unchanged identity +are not detected; prune the key (see the CLI) after such a change. + +``key`` derives from the checkpoint identity and the planned table geometry, +so TP ranks and checkpoint revisions never share a directory. A manifest that +disagrees with the live plan in any field fails startup; nothing here falls +back to a private allocation. +""" + +from __future__ import annotations + +import argparse +import ctypes +import datetime +import fcntl +import hashlib +import json +import math +import mmap +import os +import shutil +import socket +import stat +import sys +import time +from collections.abc import Sequence +from contextlib import suppress +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +LAYOUT_VERSION = 1 +MANIFEST_FORMAT = "vllm-qwen4exp-ple-shared-table/1" +MANIFEST_FILE = "manifest.json" +READY_FILE = "READY" +WEIGHT_FILE = "weight.bin" +WEIGHT_SCALE_FILE = "weight_scale.bin" +DEFAULT_DIRECTORY = "/dev/shm/vllm-ple" +ROLES = ("auto", "populate", "attach") +_KEY_HEX_DIGITS = 16 +_LOCK_POLL_S = 0.5 +_MS_SYNC = 4 # ; the mmap module does not export it. +_MAPPED_HOST_QUANT_MODES = frozenset({"nvfp4_group16"}) + + +class SharedTableError(RuntimeError): + """Base class for shared-table failures.""" + + +class SharedTableMissing(SharedTableError): + """No complete table exists for the key and the role forbids populating.""" + + +class SharedTableMismatch(SharedTableError): + """A manifest field disagrees with the live plan.""" + + def __init__(self, field: str, expected: Any, found: Any, path: Path) -> None: + self.field = field + self.expected = expected + self.found = found + super().__init__( + f"shared PLE table at {path} does not match the live plan: " + f"{field} expected {expected!r}, manifest has {found!r}" + ) + + +class SharedTableLockTimeout(SharedTableError, TimeoutError): + """The directory lock was not acquired within the timeout.""" + + +@dataclass(frozen=True, kw_only=True) +class SharedTableConfig: + """Operator-facing settings resolved by the PLE layer.""" + + directory: str = DEFAULT_DIRECTORY + role: str = "auto" + lock_timeout_s: float = 3600.0 + model_path: str + revision: str | None + identity: str + + def __post_init__(self) -> None: + if self.role not in ROLES: + raise ValueError( + f"ple_shared_table_role must be one of {ROLES}, got {self.role!r}" + ) + if self.lock_timeout_s < 0: + raise ValueError("ple_shared_table_lock_timeout_s must be non-negative") + + +@dataclass(frozen=True) +class TableFileSpec: + name: str + attribute: str + shape: tuple[int, ...] + dtype: torch.dtype + + @property + def nbytes(self) -> int: + return math.prod(self.shape) * self.dtype.itemsize + + +def _dtype_name(dtype: torch.dtype | None) -> str | None: + return None if dtype is None else str(dtype).removeprefix("torch.") + + +def _dtype_from_name(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if not isinstance(dtype, torch.dtype): + raise ValueError(f"unknown tensor dtype {name!r}") + return dtype + + +@dataclass(frozen=True, kw_only=True) +class SharedTableGeometry: + """Every input that decides the bytes of one TP rank's table files.""" + + model_path: str + revision: str | None + identity: str + quant_mode: str + tp_size: int + tp_rank: int + dense_layer_ordinal: int + padded_vocab_size: int + shard_start: int + shard_end: int + embedding_dim: int + weight_shape: tuple[int, ...] + weight_dtype: str + weight_scale_shape: tuple[int, ...] | None + weight_scale_dtype: str | None + layout_version: int = LAYOUT_VERSION + + @classmethod + def from_layout( + cls, + layout: Any, + *, + model_path: str, + revision: str | None, + identity: str, + embedding_dim: int, + dense_layer_ordinal: int, + ) -> SharedTableGeometry: + """Read the geometry from a b12x storage layout. + + ``dense_layer_ordinal`` identifies the PLE layer: each one owns its + own checkpoint shards and therefore its own table. + """ + caps = layout.caps + return cls( + model_path=str(model_path), + revision=None if revision is None else str(revision), + identity=str(identity), + quant_mode=str(caps.quant_mode), + tp_size=int(caps.tp_size), + tp_rank=int(caps.tp_rank), + dense_layer_ordinal=int(dense_layer_ordinal), + padded_vocab_size=int(layout.padded_vocab_size), + shard_start=int(layout.shard_start), + shard_end=int(layout.shard_end), + embedding_dim=int(embedding_dim), + weight_shape=tuple(int(extent) for extent in layout.weight_shape), + weight_dtype=_dtype_name(layout.weight_dtype), + weight_scale_shape=( + None + if layout.weight_scale_shape is None + else tuple(int(extent) for extent in layout.weight_scale_shape) + ), + weight_scale_dtype=_dtype_name(layout.weight_scale_dtype), + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> SharedTableGeometry: + fields = dict(data) + for name in ("weight_shape", "weight_scale_shape"): + if fields.get(name) is not None: + fields[name] = tuple(int(extent) for extent in fields[name]) + return cls(**fields) + + def as_dict(self) -> dict[str, Any]: + data = { + "model_path": self.model_path, + "revision": self.revision, + "identity": self.identity, + "quant_mode": self.quant_mode, + "tp_size": self.tp_size, + "tp_rank": self.tp_rank, + "dense_layer_ordinal": self.dense_layer_ordinal, + "padded_vocab_size": self.padded_vocab_size, + "shard_start": self.shard_start, + "shard_end": self.shard_end, + "embedding_dim": self.embedding_dim, + "weight_shape": list(self.weight_shape), + "weight_dtype": self.weight_dtype, + "weight_scale_shape": ( + None + if self.weight_scale_shape is None + else list(self.weight_scale_shape) + ), + "weight_scale_dtype": self.weight_scale_dtype, + "layout_version": self.layout_version, + } + return data + + @property + def key(self) -> str: + canonical = json.dumps(self.as_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).hexdigest()[:_KEY_HEX_DIGITS] + + @property + def files(self) -> tuple[TableFileSpec, ...]: + """Files backing the mapped-host tensors, in b12x storage order.""" + specs = [ + TableFileSpec( + WEIGHT_FILE, + "weight", + self.weight_shape, + _dtype_from_name(self.weight_dtype), + ) + ] + if ( + self.weight_scale_shape is not None + and self.quant_mode in _MAPPED_HOST_QUANT_MODES + ): + assert self.weight_scale_dtype is not None + specs.append( + TableFileSpec( + WEIGHT_SCALE_FILE, + "weight_scale", + self.weight_scale_shape, + _dtype_from_name(self.weight_scale_dtype), + ) + ) + return tuple(specs) + + @property + def nbytes(self) -> int: + return sum(spec.nbytes for spec in self.files) + + +def checkpoint_identity(model_path: str, revision: str | None) -> str: + """Name the checkpoint bytes behind ``model_path`` without reading them. + + A Hugging Face snapshot directory or cached repository resolves to its + snapshot commit; another local directory to the names, sizes and + modification times of its checkpoint files. Anything else falls back to + ``model_path@revision``. + """ + path = Path(model_path) + if path.is_dir(): + if path.parent.name == "snapshots": + return f"snapshot:{path.name}" + entries = [] + for entry in sorted(path.iterdir()): + if entry.suffix in (".safetensors", ".json") and entry.is_file(): + info = entry.stat() + entries.append(f"{entry.name}:{info.st_size}:{info.st_mtime_ns}") + digest = hashlib.sha256("\n".join(entries).encode()).hexdigest() + return f"files:{digest[:_KEY_HEX_DIGITS]}" + with suppress(Exception): + from huggingface_hub import try_to_load_from_cache + + cached = try_to_load_from_cache(model_path, "config.json", revision=revision) + if isinstance(cached, str) and Path(cached).parent.parent.name == "snapshots": + return f"snapshot:{Path(cached).parent.name}" + logger.warning( + "Cannot resolve the checkpoint behind %s; the shared PLE table key uses " + "only its name and revision, so prune the table after updating it", + model_path, + ) + return f"name:{model_path}@{revision or 'main'}" + + +def table_directory(directory: str | os.PathLike[str], key: str) -> Path: + return Path(directory) / key + + +def lock_path(directory: str | os.PathLike[str], key: str) -> Path: + return Path(directory) / f"{key}.lock" + + +class SharedTableLock: + """Exclusive ``flock`` on ``/.lock`` with a polling timeout.""" + + def __init__(self, path: Path, timeout_s: float) -> None: + self.path = path + self.timeout_s = float(timeout_s) + self._fd: int | None = None + + @property + def held(self) -> bool: + return self._fd is not None + + def acquire(self) -> None: + if self._fd is not None: + return + _ensure_private_root(self.path.parent) + fd = os.open( + self.path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600 + ) + deadline = time.monotonic() + self.timeout_s + logged = False + try: + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + pass + remaining = deadline - time.monotonic() + if remaining <= 0: + raise SharedTableLockTimeout( + f"timed out after {self.timeout_s:.0f}s waiting for the " + f"shared PLE table lock {self.path}" + ) + if not logged: + logger.info( + "Waiting for the shared PLE table lock %s (another process " + "is populating; timeout %.0fs)", + self.path, + self.timeout_s, + ) + logged = True + time.sleep(min(_LOCK_POLL_S, remaining)) + except BaseException: + os.close(fd) + raise + self._fd = fd + + def release(self) -> None: + if self._fd is None: + return + fd, self._fd = self._fd, None + try: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + def __enter__(self) -> SharedTableLock: + self.acquire() + return self + + def __exit__(self, *exc: object) -> None: + self.release() + + +def _require_private(path: Path, *, directory: bool) -> None: + """Refuse a table path another local user could have planted or altered.""" + info = os.lstat(path) + if stat.S_ISLNK(info.st_mode): + raise SharedTableError(f"shared PLE table path {path} is a symlink") + if directory and not stat.S_ISDIR(info.st_mode): + raise SharedTableError(f"shared PLE table path {path} is not a directory") + if not directory and not stat.S_ISREG(info.st_mode): + raise SharedTableError(f"shared PLE table path {path} is not a regular file") + if info.st_uid != os.geteuid(): + raise SharedTableError( + f"shared PLE table path {path} is owned by uid {info.st_uid}, not " + f"{os.geteuid()}; every replica sharing a table must run as one user" + ) + if info.st_mode & (stat.S_IWGRP | stat.S_IWOTH): + raise SharedTableError( + f"shared PLE table path {path} is writable by group or others " + f"(mode {stat.S_IMODE(info.st_mode):04o})" + ) + + +def _ensure_private_root(directory: Path) -> None: + """Create the table root ``0700`` or verify an existing one is private.""" + with suppress(FileExistsError): + os.mkdir(directory, 0o700) + _require_private(directory, directory=True) + + +def _try_lock(path: Path) -> tuple[int | None, bool]: + """Return ``(fd, busy)``: a locked descriptor, or ``busy`` when another + process holds the lock. A missing lock file yields ``(None, False)``.""" + try: + fd = os.open(path, os.O_RDWR | os.O_CLOEXEC) + except FileNotFoundError: + return None, False + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + os.close(fd) + return None, True + return fd, False + + +def _libc() -> ctypes.CDLL: + libc = ctypes.CDLL(None, use_errno=True) + libc.mmap.restype = ctypes.c_void_p + libc.mmap.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_long, + ] + libc.munmap.restype = ctypes.c_int + libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + libc.msync.restype = ctypes.c_int + libc.msync.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + return libc + + +def _cudart() -> Any: + from cuda.bindings import runtime as cudart + + return cudart + + +def _check_cuda(error: Any, operation: str) -> None: + cudart = _cudart() + if error != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"{operation} failed: {error}") + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + stride = 1 + result = [] + for extent in reversed(shape): + result.append(stride) + stride *= int(extent) + return tuple(reversed(result)) + + +def _tensor_from_pointer( + pointer: int, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + nbytes: int, +) -> torch.Tensor: + # Same construction as b12x MappedHostAllocation; copied so this module + # does not import a private b12x helper. + constructor = getattr(torch._C, "_construct_storage_from_data_pointer", None) + if constructor is None: + raise RuntimeError( + "mapped-host storage requires torch._C._construct_storage_from_data_pointer" + ) + storage = constructor(int(pointer), device, int(nbytes)) + return torch.empty(0, dtype=dtype, device=device).set_( + storage, 0, shape, _contiguous_strides(shape) + ) + + +def _host_register_read_only_supported(device: torch.device) -> bool: + cudart = _cudart() + error, value = cudart.cudaDeviceGetAttribute( + cudart.cudaDeviceAttr.cudaDevAttrHostRegisterReadOnlySupported, device.index + ) + return error == cudart.cudaError_t.cudaSuccess and bool(value) + + +class SharedTableMapping: + """Own one ``MAP_SHARED`` file mapping, its CUDA registration and aliases. + + ``host_view`` is a CPU tensor over the mapping; ``device_view`` is the CUDA + tensor b12x binds. With a CPU ``device`` (tests) nothing is registered + and ``device_view`` is ``host_view``. + """ + + def __init__( + self, + path: Path, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + writable: bool, + ) -> None: + self.path = Path(path) + self.shape = tuple(int(extent) for extent in shape) + self.dtype = dtype + self.device = device + self.nbytes = math.prod(self.shape) * dtype.itemsize + self.writable = bool(writable) + self.read_only_registration = False + self._fd: int | None = None + self._pointer = 0 + self._registered = False + self._libc = _libc() + if self.nbytes <= 0: + raise ValueError(f"shared table file {path} must be non-empty") + if device.type == "cuda" and device.index is None: + raise ValueError( + f"shared table mapping requires an indexed CUDA device, got {device}" + ) + # Attachers still open read-write when the file allows it, so the + # registration fallback below can map with PROT_WRITE. + _require_private(self.path, directory=False) + open_flags = os.O_RDWR if writable or os.access(path, os.W_OK) else os.O_RDONLY + self._fd = os.open(path, open_flags | os.O_NOFOLLOW | os.O_CLOEXEC) + try: + size = os.fstat(self._fd).st_size + if size != self.nbytes: + raise SharedTableError( + f"shared PLE table file {path} has {size} bytes, expected " + f"{self.nbytes}" + ) + if self.writable: + self._map(mmap.PROT_READ | mmap.PROT_WRITE) + self._register(read_only=False) + else: + self._map(mmap.PROT_READ) + try: + self._register(read_only=True) + except RuntimeError as exc: + if open_flags != os.O_RDWR or device.type != "cuda": + raise + logger.warning( + "Read-only registration of the shared PLE table %s was " + "rejected (%s); mapping it writable instead", + path, + exc, + ) + self._unmap() + self._map(mmap.PROT_READ | mmap.PROT_WRITE) + self._register(read_only=False) + self.host_view = _tensor_from_pointer( + self._pointer, + shape=self.shape, + dtype=dtype, + device=torch.device("cpu"), + nbytes=self.nbytes, + ) + if device.type == "cuda": + cudart = _cudart() + with torch.cuda.device(device): + error, device_pointer = cudart.cudaHostGetDevicePointer( + self._pointer, 0 + ) + _check_cuda(error, "cudaHostGetDevicePointer") + self.device_view = _tensor_from_pointer( + int(device_pointer), + shape=self.shape, + dtype=dtype, + device=device, + nbytes=self.nbytes, + ) + else: + self.device_view = self.host_view + except BaseException: + self.close() + raise + + def _map(self, prot: int) -> None: + assert self._fd is not None + pointer = self._libc.mmap(None, self.nbytes, prot, mmap.MAP_SHARED, self._fd, 0) + if pointer is None or pointer == ctypes.c_void_p(-1).value: + code = ctypes.get_errno() + raise OSError(code, f"mmap of {self.path} failed: {os.strerror(code)}") + self._pointer = int(pointer) + + def _unmap(self) -> None: + if self._pointer: + self._libc.munmap(self._pointer, self.nbytes) + self._pointer = 0 + + def _register(self, *, read_only: bool) -> None: + if self.device.type != "cuda": + return + cudart = _cudart() + flags = cudart.cudaHostRegisterMapped + if read_only: + if not _host_register_read_only_supported(self.device): + raise RuntimeError( + "cudaHostRegisterReadOnly is not supported on this device" + ) + flags |= cudart.cudaHostRegisterReadOnly + with torch.cuda.device(self.device): + (error,) = cudart.cudaHostRegister(self._pointer, self.nbytes, flags) + if error != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"cudaHostRegister failed for {self.path}: {error}; the table " + "directory must be on tmpfs (for example /dev/shm), the driver " + "does not pin mappings of ordinary file systems" + ) + self._registered = True + self.read_only_registration = read_only + + def sync(self) -> None: + """Flush the mapping and file so another process sees every byte.""" + if not self.writable: + return + if self._pointer and self._libc.msync(self._pointer, self.nbytes, _MS_SYNC): + code = ctypes.get_errno() + raise OSError(code, f"msync of {self.path} failed: {os.strerror(code)}") + if self._fd is not None: + os.fsync(self._fd) + + def close(self) -> None: + if self._registered: + cudart = _cudart() + torch.cuda.synchronize(self.device) + cudart.cudaHostUnregister(self._pointer) + self._registered = False + self._unmap() + if self._fd is not None: + os.close(self._fd) + self._fd = None + + def __del__(self) -> None: + if sys is None or sys.is_finalizing(): + return + with suppress(Exception): + self.close() + + +def _write_json(path: Path, data: dict[str, Any]) -> None: + tmp = path.with_name(path.name + ".tmp") + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(data, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def read_manifest(table_dir: Path) -> dict[str, Any] | None: + """Return the manifest of a complete table, or None when READY is absent.""" + if not (table_dir / READY_FILE).exists(): + return None + _require_private(table_dir, directory=True) + manifest_path = table_dir / MANIFEST_FILE + _require_private(manifest_path, directory=False) + try: + with open(manifest_path, encoding="utf-8") as handle: + manifest = json.load(handle) + except (OSError, ValueError) as exc: + raise SharedTableError( + f"shared PLE table {table_dir} is marked READY but its manifest " + f"cannot be read: {exc}" + ) from exc + if not isinstance(manifest, dict): + raise SharedTableError( + f"shared PLE table manifest {manifest_path} is not an object" + ) + return manifest + + +def check_manifest( + manifest: dict[str, Any], geometry: SharedTableGeometry, table_dir: Path +) -> None: + """Fail with the offending field when a manifest disagrees with the plan.""" + if manifest.get("format") != MANIFEST_FORMAT: + raise SharedTableMismatch( + "format", MANIFEST_FORMAT, manifest.get("format"), table_dir + ) + recorded = manifest.get("geometry") + if not isinstance(recorded, dict): + raise SharedTableMismatch("geometry", geometry.as_dict(), recorded, table_dir) + expected = geometry.as_dict() + for field, value in expected.items(): + found = recorded.get(field) + if found != value: + raise SharedTableMismatch(field, value, found, table_dir) + extra = sorted(set(recorded) - set(expected)) + if extra: + raise SharedTableMismatch(extra[0], None, recorded[extra[0]], table_dir) + files = manifest.get("files") + if not isinstance(files, dict): + raise SharedTableMismatch("files", {}, files, table_dir) + expected_files = {spec.name: spec.nbytes for spec in geometry.files} + if files != expected_files: + raise SharedTableMismatch("files", expected_files, files, table_dir) + for name, nbytes in expected_files.items(): + path = table_dir / name + try: + size = path.stat().st_size + except FileNotFoundError: + raise SharedTableMismatch( + f"files.{name}", nbytes, "missing", table_dir + ) from None + if size != nbytes: + raise SharedTableMismatch(f"files.{name}", nbytes, size, table_dir) + + +def _provenance() -> dict[str, Any]: + import importlib.metadata + + def version(distribution: str) -> str | None: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + ple_module = Path(__file__).with_name("b12x_ple.py") + try: + ple_module_sha256 = hashlib.sha256(ple_module.read_bytes()).hexdigest() + except OSError: + ple_module_sha256 = None + return { + "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "creator": {"hostname": socket.gethostname(), "pid": os.getpid()}, + "vllm_version": version("vllm"), + "b12x_version": version("b12x"), + "b12x_ple_sha256": ple_module_sha256, + } + + +def _remove_table_dir(table_dir: Path) -> None: + if table_dir.exists(): + shutil.rmtree(table_dir) + + +class SharedTable: + """One process's view of a shared table: mappings plus, when populating, + the held directory lock until :meth:`publish`. + + :meth:`region` is the ``host_allocator`` for b12x + ``TableLayout.allocate_storage``; the storage then owns the mappings. + """ + + def __init__( + self, + *, + geometry: SharedTableGeometry, + directory: Path, + mappings: dict[str, SharedTableMapping], + lock: SharedTableLock | None, + ) -> None: + self.geometry = geometry + self.key = geometry.key + self.directory = directory + self.table_dir = table_directory(directory, self.key) + self._mappings = mappings + self._lock = lock + self._published = lock is None + self._closed = False + self._started = time.monotonic() + self.nbytes = sum(mapping.nbytes for mapping in mappings.values()) + + @property + def populating(self) -> bool: + return self._lock is not None and not self._published + + @property + def attached(self) -> bool: + """True when the bytes came from another process and are complete.""" + return self._lock is None + + @property + def read_only(self) -> bool: + return all( + mapping.read_only_registration for mapping in self._mappings.values() + ) + + def host_view(self, attribute: str) -> torch.Tensor | None: + mapping = self._mappings.get(attribute) + return None if mapping is None else mapping.host_view + + def device_view(self, attribute: str) -> torch.Tensor | None: + mapping = self._mappings.get(attribute) + return None if mapping is None else mapping.device_view + + def region( + self, name: str, shape: tuple[int, ...], dtype: torch.dtype + ) -> SharedTableMapping: + """Return the mapping b12x requests for one table plane.""" + mapping = self._mappings.get(name) + if mapping is None: + raise SharedTableError( + f"shared PLE table {self.key} has no {name!r} plane; it holds " + f"{sorted(self._mappings)}" + ) + if mapping.shape != tuple(shape) or mapping.dtype != dtype: + raise SharedTableError( + f"shared PLE table {self.key} holds {name} as {mapping.shape} " + f"{mapping.dtype}, b12x requested {tuple(shape)} {dtype}" + ) + return mapping + + def publish(self) -> None: + """Durably mark the table complete and release the lock. + + Call exactly once, after the checkpoint loader has filled every row of + every host view and their coverage has been validated. + """ + if self._published: + return + assert self._lock is not None + try: + for mapping in self._mappings.values(): + mapping.sync() + manifest = { + "format": MANIFEST_FORMAT, + "key": self.key, + "geometry": self.geometry.as_dict(), + "files": {spec.name: spec.nbytes for spec in self.geometry.files}, + **_provenance(), + } + _write_json(self.table_dir / MANIFEST_FILE, manifest) + ready = self.table_dir / READY_FILE + fd = os.open( + ready, os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600 + ) + try: + os.fsync(fd) + finally: + os.close(fd) + _fsync_directory(self.table_dir) + except BaseException: + _remove_table_dir(self.table_dir) + self._lock.release() + self._published = True + raise + self._published = True + self._lock.release() + logger.info( + "Populated shared PLE table %s (%.2f GiB) in %s after %.1fs", + self.key, + self.nbytes / (1 << 30), + self.directory, + time.monotonic() - self._started, + ) + + def abort(self) -> None: + """Give up populating: remove the partial directory and release the lock. + + The mappings stay valid (tmpfs keeps unlinked pages alive for them), so + a model that still references the tensors cannot fault; a waiting + follower repopulates instead of timing out on the lock. + """ + if self._lock is None or self._published: + return + self._published = True + _remove_table_dir(self.table_dir) + self._lock.release() + + def close(self) -> None: + if self._closed: + return + self._closed = True + try: + for mapping in reversed(list(self._mappings.values())): + mapping.close() + finally: + # Never leave a half-written table behind for followers. + self.abort() + + def __del__(self) -> None: + if sys is None or sys.is_finalizing(): + return + with suppress(Exception): + self.abort() + + +def _open_mappings( + geometry: SharedTableGeometry, + table_dir: Path, + *, + device: torch.device, + writable: bool, +) -> dict[str, SharedTableMapping]: + mappings: dict[str, SharedTableMapping] = {} + try: + for spec in geometry.files: + mappings[spec.attribute] = SharedTableMapping( + table_dir / spec.name, + shape=spec.shape, + dtype=spec.dtype, + device=device, + writable=writable, + ) + except BaseException: + for mapping in mappings.values(): + mapping.close() + raise + return mappings + + +def _create_table_files(geometry: SharedTableGeometry, table_dir: Path) -> None: + _remove_table_dir(table_dir) + os.mkdir(table_dir, 0o700) + for spec in geometry.files: + fd = os.open( + table_dir / spec.name, + os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC, + 0o600, + ) + try: + # Reserve every page now so a full tmpfs fails here with ENOSPC + # instead of raising SIGBUS in the checkpoint copy later. + os.posix_fallocate(fd, 0, spec.nbytes) + finally: + os.close(fd) + + +def open_shared_table( + geometry: SharedTableGeometry, + config: SharedTableConfig, + *, + device: torch.device, +) -> SharedTable: + """Attach to a complete table or become its populator, per ``config.role``. + + A populator returns holding the directory lock; the caller must fill the + host views, then :meth:`SharedTable.publish`. Any failure before that + removes the directory, so a waiting follower repopulates instead of + attaching to partial bytes. + """ + directory = Path(config.directory) + key = geometry.key + table_dir = table_directory(directory, key) + lock = SharedTableLock(lock_path(directory, key), config.lock_timeout_s) + lock.acquire() + try: + manifest = None if config.role == "populate" else read_manifest(table_dir) + if manifest is not None: + check_manifest(manifest, geometry, table_dir) + mappings = _open_mappings( + geometry, table_dir, device=device, writable=False + ) + lock.release() + table = SharedTable( + geometry=geometry, directory=directory, mappings=mappings, lock=None + ) + logger.info( + "Attached shared PLE table %s (%.2f GiB, %s) from %s", + key, + table.nbytes / (1 << 30), + "read-only" if table.read_only else "writable mapping", + directory, + ) + return table + if config.role == "attach": + raise SharedTableMissing( + f"no complete shared PLE table {key} under {directory} and " + "ple_shared_table_role=attach forbids populating one; start a " + "populate/auto replica first or check the directory" + ) + logger.info( + "Populating shared PLE table %s (%.2f GiB) in %s", + key, + geometry.nbytes / (1 << 30), + directory, + ) + _create_table_files(geometry, table_dir) + try: + mappings = _open_mappings(geometry, table_dir, device=device, writable=True) + except BaseException: + _remove_table_dir(table_dir) + raise + return SharedTable( + geometry=geometry, directory=directory, mappings=mappings, lock=lock + ) + except BaseException: + lock.release() + raise + + +# --- maintenance CLI ------------------------------------------------------- + + +@dataclass(frozen=True) +class TableStatus: + key: str + state: str # "ready", "incomplete" or "populating" + nbytes: int + manifest: dict[str, Any] | None + + +def _directory_nbytes(path: Path) -> int: + # A concurrent prune may remove the directory or its files at any point. + total = 0 + try: + for entry in path.iterdir(): + with suppress(OSError): + if entry.is_file(): + total += entry.stat().st_size + except OSError: + pass + return total + + +def list_tables(directory: str | os.PathLike[str]) -> list[TableStatus]: + root = Path(directory) + if not root.is_dir(): + return [] + statuses = [] + for entry in sorted(root.iterdir()): + if not entry.is_dir(): + continue + fd, busy = _try_lock(lock_path(root, entry.name)) + state = "populating" if busy else "incomplete" + if fd is not None: + os.close(fd) + manifest = None + if state != "populating": + try: + manifest = read_manifest(entry) + except (SharedTableError, OSError): + manifest = None + if manifest is not None: + state = "ready" + statuses.append( + TableStatus(entry.name, state, _directory_nbytes(entry), manifest) + ) + return statuses + + +def prune_tables( + directory: str | os.PathLike[str], + *, + keep: Sequence[str], + dry_run: bool = False, +) -> list[str]: + """Remove every table directory whose key is not in ``keep``. + + Tables whose lock is held (being populated) are left alone. Removing a + table that running processes have mapped is safe: tmpfs keeps the pages + until the last mapping closes, and no new process can attach to it. The + empty ``.lock`` files stay: ``flock`` binds to an inode, and + unlinking one while another process opens the key would leave two owners. + """ + root = Path(directory) + kept = set(keep) + removed = [] + for status in list_tables(root): + if status.key in kept or status.state == "populating": + continue + fd, busy = _try_lock(lock_path(root, status.key)) + if busy: + continue + try: + if not dry_run: + _remove_table_dir(root / status.key) + removed.append(status.key) + finally: + if fd is not None: + os.close(fd) + return removed + + +def _format_status(status: TableStatus) -> str: + manifest = status.manifest or {} + geometry = manifest.get("geometry", {}) + creator = manifest.get("creator", {}) + return " ".join( + [ + status.key, + f"{status.state:<11}", + f"{status.nbytes / (1 << 30):7.2f} GiB", + f"model={geometry.get('model_path', '?')}", + f"revision={geometry.get('revision', '?')}", + f"tp={geometry.get('tp_rank', '?')}/{geometry.get('tp_size', '?')}", + f"layer={geometry.get('dense_layer_ordinal', '?')}", + f"created={manifest.get('created_at', '?')}", + f"by={creator.get('hostname', '?')}:{creator.get('pid', '?')}", + ] + ) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m vllm.models.qwen4_exp.nvidia.ple_shared_table", + description="Inspect or prune shared PLE table directories.", + ) + parser.add_argument( + "--dir", + default=os.environ.get("VLLM_PLE_SHARED_TABLE_DIR", DEFAULT_DIRECTORY), + help="table root directory (default: $VLLM_PLE_SHARED_TABLE_DIR or " + f"{DEFAULT_DIRECTORY})", + ) + commands = parser.add_subparsers(dest="command", required=True) + list_parser = commands.add_parser("list", help="list tables and their state") + list_parser.add_argument("--json", action="store_true", help="emit JSON") + prune_parser = commands.add_parser( + "prune", help="remove every table except the keys given with --keep" + ) + prune_parser.add_argument( + "--keep", + action="append", + default=[], + metavar="KEY", + help="key to keep; repeat for several", + ) + prune_parser.add_argument( + "--dry-run", action="store_true", help="print what would be removed" + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _build_parser().parse_args(argv) + if args.command == "list": + statuses = list_tables(args.dir) + if args.json: + print( + json.dumps( + [ + { + "key": status.key, + "state": status.state, + "nbytes": status.nbytes, + "manifest": status.manifest, + } + for status in statuses + ], + indent=2, + sort_keys=True, + ) + ) + else: + for status in statuses: + print(_format_status(status)) + return 0 + removed = prune_tables(args.dir, keep=args.keep, dry_run=args.dry_run) + verb = "would remove" if args.dry_run else "removed" + for key in removed: + print(f"{verb} {key}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 5daabd5b0b6b993db4366e46d81be29c8096574e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 14:52:01 +0000 Subject: [PATCH 339/365] docs(release): describe the shared PLE table Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-961.json | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 .lil/changes/vllm-961.json diff --git a/.lil/changes/vllm-961.json b/.lil/changes/vllm-961.json new file mode 100644 index 000000000000..f4279424f035 --- /dev/null +++ b/.lil/changes/vllm-961.json @@ -0,0 +1,25 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-961", + "category": "feature", + "summary": "Qwen3.8-Flash-Next servers on one host can share one host-RAM copy of the 26.8 GiB PLE table", + "models": [ + "Qwen3.8-Flash-Next" + ], + "compatibility": "No user action required. Opt in with VLLM_PLE_TABLE_MEMORY=shared; the launcher sets it for REPLICAS=2 or more (docker replicas mode). Tables stay in /dev/shm until removed, so with --ipc host they outlive the container; python -m vllm.models.qwen4_exp.nvidia.ple_shared_table list|prune manages them.", + "details": [ + "With VLLM_PLE_TABLE_MEMORY=shared the PLE table lives in tmpfs files keyed by the checkpoint and table geometry. The first server loads and validates it, and the others map the same bytes instead of keeping their own 26.8 GiB copy and reading the table again.", + "A failed table check abandons the shared copy, so no server maps an incomplete table.", + "Ported from René Honig's vLLM #774; awsdl (carrion256) showed the 2xTP1 setup it serves." + ], + "authors": [ + "renehonig", + "voipmonitor" + ], + "pull_requests": [ + 961 + ], + "requires": [ + "b12x-454" + ] +} From 4d11fdcff1198a4f19b1c87d3cb1647754bf7f6f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 16:06:18 +0000 Subject: [PATCH 340/365] Own CSF checkpoint discovery in vLLM and supply generic tensor sources Validate model inventories and map source tensor names in the CSF loaders. Pass lazy projection views to B12X for TP slicing and GPU preparation, preserving calibration and bounded weight materialization. Requires the matching B12X tensor-source API; checkpoint contents and serving flags are unchanged. Validation: 30 loader tests, staged-file hooks and manual mypy-3.12 pass; real-layer tensor parity covers all four model families. Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-956-csf-tensor-sources.json | 15 +++ docs/features/quantization/fp4_csf.md | 5 + tests/quantization/test_mxfp4_csf.py | 95 +++++++++++++++++-- tests/quantization/test_nvfp4_csf.py | 3 +- .../layers/quantization/mxfp4_csf.py | 4 +- .../layers/quantization/nvfp4_csf.py | 5 +- vllm/model_executor/model_loader/csf_utils.py | 76 +++++++++++++++ .../model_loader/mxfp4_csf_loader.py | 76 ++++++++++++++- .../model_loader/nvfp4_csf_loader.py | 72 +++++++++++++- vllm/models/deepseek_v4_1/mxfp4_csf.py | 5 +- 10 files changed, 341 insertions(+), 15 deletions(-) create mode 100644 .lil/changes/vllm-956-csf-tensor-sources.json create mode 100644 vllm/model_executor/model_loader/csf_utils.py diff --git a/.lil/changes/vllm-956-csf-tensor-sources.json b/.lil/changes/vllm-956-csf-tensor-sources.json new file mode 100644 index 000000000000..2702fee78df0 --- /dev/null +++ b/.lil/changes/vllm-956-csf-tensor-sources.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-956-csf-tensor-sources", + "category": "compatibility", + "summary": "Keep CSF checkpoint discovery and model tensor mapping in vLLM", + "models": ["Kimi-K3", "DeepSeek-V4.1-Flash", "GLM-5.3-Flash", "Qwen3.8-Flash-Next"], + "compatibility": "Requires the B12X CsfMatrix tensor-source loader API. Canonical CSF checkpoints and launch commands remain compatible.", + "details": [ + "vLLM validates manifests and model geometry, then passes lazy expert projection sources to B12X.", + "Shared shard handling preserves TP-bounded weight reads and closes files after loading or failure.", + "Retained tensors, FP32 calibration and native expert execution preserve their existing behavior." + ], + "pull_requests": [956], + "requires": ["b12x-450-csf-tensor-sources"] +} diff --git a/docs/features/quantization/fp4_csf.md b/docs/features/quantization/fp4_csf.md index 1b77cdff32a2..9d0c3f65801f 100644 --- a/docs/features/quantization/fp4_csf.md +++ b/docs/features/quantization/fp4_csf.md @@ -78,6 +78,11 @@ local alignment requirement. ## Compatibility +Install matching vLLM and B12X CSF support. vLLM owns the checkpoint manifest, +model tensor names and file access; B12X receives tensor sources for TP slicing +and GPU preparation. This API boundary does not change checkpoint contents, +serving flags or native expert arithmetic. + Readers require `lil-mxfp4-csf-checkpoint/1` or `lil-nvfp4-csf-checkpoint/1` and their `.mxfp4_csf_*` or `.nvfp4_csf_*` tensor components. Predecessor X4T/LSC names, schemas and suffixes are not diff --git a/tests/quantization/test_mxfp4_csf.py b/tests/quantization/test_mxfp4_csf.py index 47f9fe683c3e..c8a1d0d2ae65 100644 --- a/tests/quantization/test_mxfp4_csf.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -77,8 +77,7 @@ def test_exact_config_and_loader_are_registered(): def test_csf_cpu_shard_reader_bounds_threads_and_restores_on_failure( monkeypatch, moe, owner ): - from b12x.moe.checkpoints import mxfp4_csf as checkpoint - + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint from vllm.models.deepseek_v4_1 import mxfp4_csf from vllm.utils.torch_utils import set_default_torch_num_threads @@ -193,7 +192,7 @@ def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): - from b12x.moe.checkpoints.mxfp4_csf import CODEC, SCHEMA + from vllm.model_executor.model_loader.mxfp4_csf_loader import CODEC, SCHEMA tensor_dir = tmp_path / "tensors" tensor_dir.mkdir() @@ -283,7 +282,7 @@ def test_kimi_mxfp4_csf_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): def test_kimi_loader_skips_compressed_experts_with_nested_text_config( tmp_path, quant_method, load_format ): - from b12x.moe.checkpoints.mxfp4_csf import CODEC, SCHEMA + from vllm.model_executor.model_loader.mxfp4_csf_loader import CODEC, SCHEMA tensor_dir = tmp_path / "tensors" tensor_dir.mkdir() @@ -323,8 +322,6 @@ def test_kimi_loader_skips_compressed_experts_with_nested_text_config( def test_mxfp4_csf_selects_retained_precision_and_rejects_other_configs( monkeypatch, moe, family ): - from b12x.moe.checkpoints import mxfp4_csf as checkpoint - from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import ( @@ -332,6 +329,7 @@ def test_mxfp4_csf_selects_retained_precision_and_rejects_other_configs( kimi_mxfp4_csf, ) from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint from vllm.models.deepseek_v4_1 import mxfp4_csf monkeypatch.setattr(checkpoint, "checkpoint_contract", lambda _: {"family": family}) @@ -373,3 +371,88 @@ def test_mxfp4_csf_selects_retained_precision_and_rejects_other_configs( ) is None ) + + +@pytest.mark.parametrize("codec", ["mxfp4", "nvfp4"]) +@pytest.mark.parametrize("fault", ["schema", "codec", "family", "inventory", "path"]) +def test_csf_contract_rejects_mismatched_identity_and_shard_inventory( + tmp_path, codec, fault +): + from vllm.model_executor.model_loader import mxfp4_csf_loader, nvfp4_csf_loader + + reader = mxfp4_csf_loader if codec == "mxfp4" else nvfp4_csf_loader + identity = { + "schema": reader.SCHEMA, + "codec": reader.CODEC, + "family": next(iter(reader.FAMILIES)), + } + manifest = {**identity, "shards": [{"file": "weights.safetensors"}]} + contract = {**identity, "source_names": {"weight": "weights.safetensors"}} + if fault in ("schema", "codec", "family"): + contract[fault] = "unrecognized" + elif fault == "inventory": + contract["source_names"]["weight"] = "missing.safetensors" + else: + contract["source_names"]["weight"] = "../weights.safetensors" + manifest["shards"][0]["file"] = "../weights.safetensors" + (tmp_path / "manifest.json").write_text(json.dumps(manifest)) + (tmp_path / "build-contract.json").write_text(json.dumps(contract)) + with pytest.raises(ValueError): + reader.checkpoint_contract(str(tmp_path)) + + +def test_csf_tensor_reader_resolves_separate_shards_and_closes_after_failure( + tmp_path, monkeypatch +): + from vllm.model_executor.model_loader import csf_utils + + tensors = tmp_path / "tensors" + tensors.mkdir() + weight = torch.arange(128, dtype=torch.uint8).reshape(16, 8) + fixed = torch.arange(32, dtype=torch.uint8).reshape(1, 32) + exceptions = torch.tensor([0xFE000009], dtype=torch.uint32) + save_file({"arbitrary.packed": weight}, tensors / "weights.safetensors") + save_file( + { + "arbitrary.scales.mxfp4_csf_fixed": fixed, + "arbitrary.scales.mxfp4_csf_exceptions": exceptions, + }, + tensors / "scales.safetensors", + ) + index = { + "arbitrary.packed": "weights.safetensors", + "arbitrary.scales": "scales.safetensors", + } + opened, closed = [], [] + original_open = csf_utils.safe_open + + class TrackedShard: + def __init__(self, path, **kwargs): + self.path = path + self.shard = original_open(path, **kwargs) + + def __enter__(self): + opened.append(self.path.name) + return self.shard.__enter__() + + def __exit__(self, *exc): + closed.append(self.path.name) + return self.shard.__exit__(*exc) + + monkeypatch.setattr(csf_utils, "safe_open", TrackedShard) + with ( + pytest.raises(RuntimeError, match="preparation failed"), + csf_utils.CsfTensorReader(tmp_path, index, "mxfp4") as reader, + ): + for _ in range(2): + source = reader.matrix("arbitrary.packed", "arbitrary.scales") + assert source.weight.get_shape() == [16, 8] + assert torch.equal(source.weight[8:16, 2:6], weight[8:16, 2:6]) + assert torch.equal(source.fixed, fixed) + assert torch.equal(source.exceptions, exceptions) + raise RuntimeError("preparation failed") + assert ( + sorted(opened) + == sorted(closed) + == ["scales.safetensors", "weights.safetensors"] + ) diff --git a/tests/quantization/test_nvfp4_csf.py b/tests/quantization/test_nvfp4_csf.py index dae9a4c54920..f8144f304088 100644 --- a/tests/quantization/test_nvfp4_csf.py +++ b/tests/quantization/test_nvfp4_csf.py @@ -117,10 +117,9 @@ def test_prepared_nvfp4_accepts_kernel_order_and_rejects_activation_change(monke def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( tmp_path, quant_method, load_format ): - from b12x.moe.checkpoints.nvfp4_csf import CODEC, SCHEMA - from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.model_loader import get_model_loader + from vllm.model_executor.model_loader.nvfp4_csf_loader import CODEC, SCHEMA directory = tmp_path / "tensors" directory.mkdir() diff --git a/vllm/model_executor/layers/quantization/mxfp4_csf.py b/vllm/model_executor/layers/quantization/mxfp4_csf.py index 80c0e1446ed7..c0e2bf95eade 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_csf.py +++ b/vllm/model_executor/layers/quantization/mxfp4_csf.py @@ -23,7 +23,9 @@ def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): @classmethod def from_config(cls, config): - from b12x.moe.checkpoints.mxfp4_csf import checkpoint_contract + from vllm.model_executor.model_loader.mxfp4_csf_loader import ( + checkpoint_contract, + ) if config.get("format_version") != 1 or not config.get("checkpoint_root"): raise ValueError("MXFP4-CSF requires format_version=1 and checkpoint_root") diff --git a/vllm/model_executor/layers/quantization/nvfp4_csf.py b/vllm/model_executor/layers/quantization/nvfp4_csf.py index b73e38d500f6..55398c32102c 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_csf.py +++ b/vllm/model_executor/layers/quantization/nvfp4_csf.py @@ -154,7 +154,10 @@ def get_fused_moe_quant_config(self, layer): def process_weights_after_loading(self, layer): from b12x.moe import fused_moe - from b12x.moe.checkpoints.nvfp4_csf import read_nvfp4_csf_layer + + from vllm.model_executor.model_loader.nvfp4_csf_loader import ( + read_nvfp4_csf_layer, + ) tp, rank = ( get_tensor_model_parallel_world_size(), diff --git a/vllm/model_executor/model_loader/csf_utils.py b/vllm/model_executor/model_loader/csf_utils.py new file mode 100644 index 000000000000..8ad5b0a8ec7d --- /dev/null +++ b/vllm/model_executor/model_loader/csf_utils.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Manifest validation and bounded file access for compressed FP4 checkpoints.""" + +import json +from contextlib import ExitStack +from pathlib import Path + +from safetensors import safe_open + + +def read_csf_contract(root, *, schema, codec, families): + """Check storage identity and the source tensor-to-shard inventory.""" + root = Path(root) + manifest = json.loads((root / "manifest.json").read_text()) + contract = json.loads((root / "build-contract.json").read_text()) + for key, expected in (("schema", schema), ("codec", codec)): + if manifest.get(key) != expected or contract.get(key) != expected: + raise ValueError(f"CSF requires {key}={expected!r}") + if ( + contract.get("family") not in families + or manifest.get("family") != contract["family"] + ): + raise ValueError("CSF manifest and contract must name a supported model family") + names = contract["source_names"] + files = {item["file"] for item in manifest["shards"]} + if set(names.values()) != files: + raise ValueError("CSF manifest and source tensor inventory disagree") + for name in files: + if Path(name).name != name or not name.endswith(".safetensors"): + raise ValueError( + "CSF shard names must be checkpoint-local safetensors files" + ) + return contract + + +class CsfTensorReader(ExitStack): + """Keep a layer's lazily opened shards alive while B12X consumes its tensors.""" + + def __init__(self, root, source_names, codec): + super().__init__() + self.root = Path(root) + self.source_names = source_names + self.codec = codec + self.handles = {} + + def _handle(self, name): + filename = self.source_names[name] + if filename not in self.handles: + self.handles[filename] = self.enter_context( + safe_open( + self.root / "tensors" / filename, framework="pt", device="cpu" + ) + ) + return self.handles[filename] + + def matrix(self, weight, scale, *, global_scale=None, input_scale=None): + """Expose a lazy packed weight view and one projection's CPU scale data.""" + from b12x.moe.checkpoints.csf import CsfMatrix + + scales = self._handle(scale) + return CsfMatrix( + weight=self._handle(weight).get_slice(weight), + fixed=scales.get_tensor(f"{scale}.{self.codec}_csf_fixed"), + exceptions=scales.get_tensor(f"{scale}.{self.codec}_csf_exceptions"), + global_scale=( + self._handle(global_scale).get_tensor(global_scale) + if global_scale is not None + else None + ), + input_scale=( + self._handle(input_scale).get_tensor(input_scale) + if input_scale is not None + else None + ), + ) diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py index 7ed3c7ab6e32..0e1dd4152ad0 100644 --- a/vllm/model_executor/model_loader/mxfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -8,22 +8,94 @@ """ import time +from functools import lru_cache from pathlib import Path import regex as re from safetensors import safe_open +from vllm.model_executor.model_loader.csf_utils import ( + CsfTensorReader, + read_csf_contract, +) from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( file_source_tensor, safetensors_file_sources, ) +SCHEMA = "lil-mxfp4-csf-checkpoint/1" +CODEC = "row-base-offset1-u24-exceptions/1" +FAMILIES = { + "deepseek_v41": (384, 5120, 2304), + "kimi_k3": (896, 3584, 3072), +} + + +@lru_cache(maxsize=4) +def checkpoint_contract(root: str) -> dict: + """Validate the MXFP4-CSF container before loading model tensors.""" + return read_csf_contract(root, schema=SCHEMA, codec=CODEC, families=FAMILIES) + + +def read_mxfp4_csf_layer( + root, + layer_index, + *, + num_experts, + hidden_size, + intermediate_size, + tp_rank, + tp_size, + device, + w13_scale_scratch, + w2_scale_scratch, +): + """Resolve model tensor names and pass gate/up/down sources to B12X.""" + from b12x.moe.checkpoints.mxfp4_csf import load_mxfp4_csf_weights + + contract = checkpoint_contract(str(Path(root).resolve())) + family = contract["family"] + if (num_experts, hidden_size, intermediate_size) != FAMILIES[family]: + raise ValueError("MXFP4-CSF expert geometry differs from the checkpoint family") + supported_tp = (1, 2, 4, 8) if family == "deepseek_v41" else (1, 2, 4, 8, 12, 16) + if tp_size not in supported_tp or not 0 <= tp_rank < tp_size: + raise ValueError( + f"MXFP4-CSF {family} supports TP {supported_tp} with a valid rank" + ) + with CsfTensorReader(root, contract["source_names"], "mxfp4") as reader: + + def experts(): + for expert in range(num_experts): + if family == "kimi_k3": + prefix = ( + f"language_model.model.layers.{layer_index}." + f"block_sparse_moe.experts.{expert}" + ) + weight, scale = "weight_packed", "weight_scale" + else: + prefix = f"layers.{layer_index}.ffn.experts.{expert}" + weight, scale = "weight", "scale" + yield tuple( + reader.matrix(f"{prefix}.{p}.{weight}", f"{prefix}.{p}.{scale}") + for p in ("w1", "w3", "w2") + ) + + return load_mxfp4_csf_weights( + experts(), + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_rank=tp_rank, + tp_size=tp_size, + device=device, + w13_scale_scratch=w13_scale_scratch, + w2_scale_scratch=w2_scale_scratch, + ) + class Mxfp4CsfModelLoader(DefaultModelLoader): def _root(self, model_config): - from b12x.moe.checkpoints.mxfp4_csf import checkpoint_contract - text_config = getattr(model_config, "hf_text_config", model_config.hf_config) quant = getattr(model_config.hf_config, "quantization_config", None) quant = quant or text_config.quantization_config diff --git a/vllm/model_executor/model_loader/nvfp4_csf_loader.py b/vllm/model_executor/model_loader/nvfp4_csf_loader.py index a8b8c01dbcde..b2f60ebf1185 100644 --- a/vllm/model_executor/model_loader/nvfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/nvfp4_csf_loader.py @@ -3,22 +3,90 @@ """Load retained tensors while the quantization method owns compressed experts.""" import time +from functools import lru_cache from pathlib import Path import regex as re from safetensors import safe_open +from vllm.model_executor.model_loader.csf_utils import ( + CsfTensorReader, + read_csf_contract, +) from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( file_source_tensor, safetensors_file_sources, ) +SCHEMA = "lil-nvfp4-csf-checkpoint/1" +CODEC = "byte-window4-fixed-stream-u24-exceptions/1" +FAMILIES = { + "glm53_nvfp4": (288, 4096, 2048, range(3, 45)), + "qwen38_flash_next_nvfp4": (512, 2560, 640, range(48)), +} + + +@lru_cache(maxsize=4) +def checkpoint_contract(root: str) -> dict: + """Validate the NVFP4-CSF container before loading model tensors.""" + return read_csf_contract(root, schema=SCHEMA, codec=CODEC, families=FAMILIES) + + +def read_nvfp4_csf_layer( + root, + layer_index, + *, + num_experts, + hidden_size, + intermediate_size, + tp_rank, + tp_size, + device, + w13_scale_scratch, + w2_scale_scratch, +): + """Resolve model tensor names and pass up/gate/down sources to B12X.""" + from b12x.moe.checkpoints.nvfp4_csf import load_nvfp4_csf_weights + + contract = checkpoint_contract(str(Path(root).resolve())) + e, h, n, layers = FAMILIES[contract["family"]] + if (num_experts, hidden_size, intermediate_size) != (e, h, n): + raise ValueError("NVFP4-CSF expert geometry differs from the checkpoint family") + if layer_index not in layers: + raise ValueError("NVFP4-CSF layer is outside the compressed expert inventory") + with CsfTensorReader(root, contract["source_names"], "nvfp4") as reader: + + def experts(): + for expert in range(num_experts): + prefix = ( + f"model.language_model.layers.{layer_index}.mlp.experts.{expert}" + ) + yield tuple( + reader.matrix( + f"{prefix}.{p}.weight", + f"{prefix}.{p}.weight_scale", + global_scale=f"{prefix}.{p}.weight_scale_2", + input_scale=f"{prefix}.{p}.input_scale", + ) + for p in ("up_proj", "gate_proj", "down_proj") + ) + + return load_nvfp4_csf_weights( + experts(), + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_rank=tp_rank, + tp_size=tp_size, + device=device, + w13_scale_scratch=w13_scale_scratch, + w2_scale_scratch=w2_scale_scratch, + ) + class Nvfp4CsfModelLoader(DefaultModelLoader): def _root(self, model_config): - from b12x.moe.checkpoints.nvfp4_csf import checkpoint_contract - quant = getattr(model_config.hf_config, "quantization_config", None) quant = quant or model_config.hf_text_config.quantization_config if quant.get("quant_method") != "nvfp4_csf": diff --git a/vllm/models/deepseek_v4_1/mxfp4_csf.py b/vllm/models/deepseek_v4_1/mxfp4_csf.py index 6bd454a56885..40290f7da1a1 100644 --- a/vllm/models/deepseek_v4_1/mxfp4_csf.py +++ b/vllm/models/deepseek_v4_1/mxfp4_csf.py @@ -152,7 +152,10 @@ def get_fused_moe_quant_config(self, layer): def process_weights_after_loading(self, layer): from b12x.moe import fused_moe - from b12x.moe.checkpoints.mxfp4_csf import read_mxfp4_csf_layer + + from vllm.model_executor.model_loader.mxfp4_csf_loader import ( + read_mxfp4_csf_layer, + ) tp, rank = ( get_tensor_model_parallel_world_size(), From 80c4201c6cdc724fb3c361e88454eadde1269b9c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 16:30:30 +0000 Subject: [PATCH 341/365] Own CSF checkpoint slicing and use B12X weight preparation Read manifests, resolve model tensor names and assemble TP-local FP4 weights in vLLM. Pass canonical CPU scale planes to B12X prepare_weights, which owns scale upload and swizzling. Requires paired CsfScalePlanes support; checkpoint schemas and launch flags are unchanged. Validation: 56 loader tests, staged-file hooks and manual mypy-3.12 pass. Ten complete layer/rank comparisons across Kimi, DeepSeek, GLM and Qwen yield byte-identical prepared state. Full-model serving was not rerun. Co-authored-by: Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-956-csf-preparation.json | 24 +++ docs/features/quantization/fp4_csf.md | 5 +- tests/quantization/csf_fixtures.py | 77 ++++++++ tests/quantization/test_mxfp4_csf.py | 146 +++++++++++++++ tests/quantization/test_nvfp4_csf.py | 120 +++++++++++++ vllm/model_executor/model_loader/csf_utils.py | 49 ++++- .../model_loader/mxfp4_csf_loader.py | 142 ++++++++++++++- .../model_loader/nvfp4_csf_loader.py | 170 +++++++++++++++++- 8 files changed, 720 insertions(+), 13 deletions(-) create mode 100644 .lil/changes/vllm-956-csf-preparation.json create mode 100644 tests/quantization/csf_fixtures.py diff --git a/.lil/changes/vllm-956-csf-preparation.json b/.lil/changes/vllm-956-csf-preparation.json new file mode 100644 index 000000000000..d3fcbc2026f8 --- /dev/null +++ b/.lil/changes/vllm-956-csf-preparation.json @@ -0,0 +1,24 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-956-csf-preparation", + "category": "compatibility", + "summary": "Own CSF file loading and pass raw scale planes to B12X preparation", + "models": [ + "Kimi-K3", + "DeepSeek-V4.1-Flash", + "GLM-5.3-Flash", + "Qwen3.8-Flash-Next" + ], + "compatibility": "Install the paired vLLM/B12X CSF preparation changes. Canonical CSF checkpoint files and launch commands remain compatible.", + "details": [ + "vLLM owns CSF manifests, model inventories, shard lifetimes, TP slicing and expert assembly.", + "The quantization method supplies CsfScalePlanes to the standard B12X prepare_weights API; vLLM does not construct kernel-specific GPU scale batches.", + "Canonical checkpoint schemas, native tensor precision and launch flags are unchanged." + ], + "pull_requests": [ + 956 + ], + "requires": [ + "b12x-450-csf-preparation" + ] +} diff --git a/docs/features/quantization/fp4_csf.md b/docs/features/quantization/fp4_csf.md index 9d0c3f65801f..b7449ca4708f 100644 --- a/docs/features/quantization/fp4_csf.md +++ b/docs/features/quantization/fp4_csf.md @@ -79,8 +79,9 @@ local alignment requirement. ## Compatibility Install matching vLLM and B12X CSF support. vLLM owns the checkpoint manifest, -model tensor names and file access; B12X receives tensor sources for TP slicing -and GPU preparation. This API boundary does not change checkpoint contents, +model tensor names, file access and TP slicing. B12X receives rank-local weights +and compressed CPU scale planes; its ordinary `prepare_weights` post-processing +uploads and rearranges them for the selected kernels. This API boundary does not change checkpoint contents, serving flags or native expert arithmetic. Readers require `lil-mxfp4-csf-checkpoint/1` or diff --git a/tests/quantization/csf_fixtures.py b/tests/quantization/csf_fixtures.py new file mode 100644 index 000000000000..da39252f4ec8 --- /dev/null +++ b/tests/quantization/csf_fixtures.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small compressed matrices with known bytes across TP partition boundaries.""" + +import numpy as np +import torch + +from vllm.model_executor.model_loader.csf_utils import CsfMatrix + + +class TensorView: + def __init__(self, tensor): + self.tensor = tensor + + def get_shape(self): + return list(self.tensor.shape) + + def get_dtype(self): + return "I8" if self.tensor.dtype == torch.int8 else "U8" + + def __getitem__(self, extent): + return self.tensor[extent] + + +def matrix(rows, channels, *, group_size, seed): + """Return a source projection and its uncompressed reference scale bytes.""" + columns = channels // group_size + bases = ((np.arange(rows) + seed) % 190 + 20).astype(np.uint8) + offsets = (np.arange(rows * columns) + seed) % (2 if group_size == 32 else 16) + offsets = offsets.reshape(rows, columns).astype(np.uint8) + logical = bases[:, None] + offsets + positions = np.unique( + [ + 0, + columns - 1, + rows // 2 * columns - 1, + rows // 2 * columns, + rows * columns - 1, + ] + ).astype(np.uint32) + values = (np.arange(len(positions), dtype=np.uint32) * 7 + seed) % 256 + logical.flat[positions] = values.astype(np.uint8) + packed = ( + np.packbits(offsets, axis=1, bitorder="little") + if group_size == 32 + else offsets[:, ::2] | (offsets[:, 1::2] << 4) + ) + fixed = np.concatenate( + (bases.reshape(rows // 16, 16), packed.reshape(rows // 16, -1)), axis=1 + ) + weight = (torch.arange(rows * channels // 2) + seed).to(torch.uint8) + source = CsfMatrix( + weight=TensorView(weight.reshape(rows, channels // 2)), + fixed=torch.from_numpy(fixed), + exceptions=torch.from_numpy(positions | (values << 24)), + ) + return source, torch.from_numpy(logical.copy()) + + +def decode_planes(planes, rows, columns, group_size): + """Reconstruct CPU scale bytes independently of B12X preparation.""" + result = [] + for fixed, exceptions in zip(planes.fixed, planes.exceptions, strict=True): + selectors = (columns + 7) // 8 if group_size == 32 else columns // 2 + slabs = fixed.numpy().reshape(rows // 16, 16 * (1 + selectors)) + packed = slabs[:, 16:].reshape(rows, selectors) + if group_size == 32: + offsets = np.unpackbits(packed, axis=1, bitorder="little")[:, :columns] + else: + offsets = np.stack((packed & 15, packed >> 4), axis=-1).reshape( + rows, columns + ) + logical = (slabs[:, :16].reshape(rows, 1) + offsets).astype(np.uint8) + words = exceptions.numpy() + logical.flat[words & 0xFFFFFF] = (words >> 24).astype(np.uint8) + result.append(torch.from_numpy(logical.copy())) + return torch.stack(result) diff --git a/tests/quantization/test_mxfp4_csf.py b/tests/quantization/test_mxfp4_csf.py index c8a1d0d2ae65..72e1c45526a5 100644 --- a/tests/quantization/test_mxfp4_csf.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -6,6 +6,7 @@ from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest import torch from safetensors.torch import save_file @@ -18,11 +19,16 @@ RoutingMethodType, ) from vllm.model_executor.model_loader.mxfp4_csf_loader import Mxfp4CsfModelLoader +from vllm.model_executor.model_loader.mxfp4_csf_loader import ( + _slice_scale_plane as slice_scale_plane, +) from vllm.models.deepseek_v4_1.mxfp4_csf import ( DeepseekV41Mxfp4CsfConfig, Mxfp4CsfMoEMethod, ) +from .csf_fixtures import decode_planes, matrix + @pytest.fixture def moe(): @@ -456,3 +462,143 @@ def __exit__(self, *exc): == sorted(closed) == ["scales.safetensors", "weights.safetensors"] ) + + +def decode(fixed, exceptions, rows, columns): + selectors = (columns + 7) // 8 + stream = fixed.numpy().reshape(rows // 16, 16 * (1 + selectors)) + bits = np.unpackbits( + stream[:, 16:].reshape(rows, selectors), axis=1, bitorder="little" + )[:, :columns] + result = (stream[:, :16].reshape(rows, 1) + bits).astype(np.uint8) + words = exceptions.numpy() + result.flat[words & 0xFFFFFF] = words >> 24 + return result + + +@pytest.mark.parametrize("columns", [9, 18, 36, 72, 160]) +@pytest.mark.parametrize("row_slice", [(0, 64), (16, 48), (48, 64)]) +def test_slices_preserve_all_scale_bytes(columns, row_slice): + rng = np.random.default_rng(78) + rows = 64 + bits = rng.integers(0, 2, (rows, columns), dtype=np.uint8) + bases = rng.integers(0, 254, (rows // 16, 16), dtype=np.uint8) + fixed = torch.from_numpy( + np.concatenate( + ( + bases, + np.packbits(bits, axis=1, bitorder="little").reshape(rows // 16, -1), + ), + 1, + ) + ) + positions = np.unique(rng.integers(0, rows * columns, 50)).astype(np.uint32) + values = rng.integers(0, 256, len(positions), dtype=np.uint32) + exceptions = torch.from_numpy(positions | values << 24) + reference = decode(fixed, exceptions, rows, columns) + for c0, c1 in ((0, columns), (1, columns), (columns // 2, columns)): + sliced = slice_scale_plane( + fixed, exceptions, rows, columns, row_slice, (c0, c1) + ) + actual = decode(sliced[0], sliced[1], row_slice[1] - row_slice[0], c1 - c0) + assert np.array_equal(actual, reference[row_slice[0] : row_slice[1], c0:c1]) + + +def test_rejects_unaligned_rows_and_bad_padding(): + fixed = torch.zeros((4, 48), dtype=torch.uint8) + exceptions = torch.empty(0, dtype=torch.uint32) + with pytest.raises(ValueError, match="16-row"): + slice_scale_plane(fixed, exceptions, 64, 9, (1, 17), (0, 9)) + fixed[0, 17] = 128 + with pytest.raises(ValueError, match="unused selector"): + slice_scale_plane(fixed, exceptions, 64, 9, (0, 64), (0, 9)) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_tensor_sources_preserve_projection_order_and_tp_bytes(rank): + """An arbitrary geometry loads without any model identity or checkpoint path.""" + from vllm.model_executor.model_loader.mxfp4_csf_loader import ( + _load_mxfp4_csf_weights as load_mxfp4_csf_weights, + ) + + device = "cpu" + experts, hidden, intermediate, local = 3, 256, 192, 96 + sources, scales = [], [] + for expert in range(experts): + pairs = [ + matrix(r, c, group_size=32, seed=expert * 17 + projection) + for projection, (r, c) in enumerate( + ((intermediate, hidden), (intermediate, hidden), (hidden, intermediate)) + ) + ] + sources.append(tuple(pair[0] for pair in pairs)) + scales.append(tuple(pair[1] for pair in pairs)) + scratch13 = torch.empty( + (experts, hidden // 32, 2 * local), dtype=torch.uint8, device=device + ) + scratch2 = torch.empty( + (experts, local // 32, hidden), dtype=torch.uint8, device=device + ) + weights = load_mxfp4_csf_weights( + iter(sources), + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + tp_rank=rank, + tp_size=2, + device=device, + w13_scale_scratch=scratch13, + w2_scale_scratch=scratch2, + ) + first, last = rank * local, (rank + 1) * local + expected13 = torch.stack( + [ + torch.cat((gate.weight[first:last, :], up.weight[first:last, :])) + for gate, up, _ in sources + ] + ) + expected2 = torch.stack( + [down.weight[:, first // 2 : last // 2] for _, _, down in sources] + ) + assert torch.equal(weights.w13.cpu(), expected13) + assert torch.equal(weights.w2.cpu(), expected2) + assert ( + weights.w13_scale_scratch is scratch13 and weights.w2_scale_scratch is scratch2 + ) + for batch, expected in ( + ( + weights.w13_scales, + torch.stack( + [torch.cat((s[0][first:last], s[1][first:last])) for s in scales] + ), + ), + ( + weights.w2_scales, + torch.stack([s[2][:, first // 32 : last // 32] for s in scales]), + ), + ): + output = decode_planes( + batch, expected.shape[1], expected.shape[2], group_size=32 + ) + assert torch.equal(output, expected) + + +@pytest.mark.parametrize("count", [0, 2]) +def test_expert_inventory_must_match_declared_count(count): + from vllm.model_executor.model_loader.mxfp4_csf_loader import ( + _load_mxfp4_csf_weights as load_mxfp4_csf_weights, + ) + + source, _ = matrix(128, 128, group_size=32, seed=0) + with pytest.raises(ValueError, match="zip"): + load_mxfp4_csf_weights( + [(source, source, source)] * count, + num_experts=1, + hidden_size=128, + intermediate_size=128, + tp_rank=0, + tp_size=1, + device="cpu", + w13_scale_scratch=None, + w2_scale_scratch=None, + ) diff --git a/tests/quantization/test_nvfp4_csf.py b/tests/quantization/test_nvfp4_csf.py index f8144f304088..f988b64d69e4 100644 --- a/tests/quantization/test_nvfp4_csf.py +++ b/tests/quantization/test_nvfp4_csf.py @@ -25,6 +25,8 @@ ) from vllm.model_executor.model_loader.nvfp4_csf_loader import Nvfp4CsfModelLoader +from .csf_fixtures import decode_planes, matrix + @pytest.mark.parametrize("config_cls", [Nvfp4CsfConfig]) def test_csf_config_preserves_source_recipes_and_avoids_expert_allocations(config_cls): @@ -161,3 +163,121 @@ def test_loader_excludes_compressed_main_experts_and_retains_native_tensors( assert set(actual) == {name for name in tensors if ".layers.3." not in name} for name, tensor in actual.items(): assert torch.equal(tensor, tensors[name]) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_tensor_sources_preserve_tp_bytes_and_fp32_with_bf16_default(rank): + from vllm.model_executor.model_loader.nvfp4_csf_loader import ( + _load_nvfp4_csf_weights as load_nvfp4_csf_weights, + ) + + device = "cpu" + experts, hidden, intermediate, local = 3, 256, 256, 128 + calibration = torch.linspace(0.013123, 0.056789, experts, dtype=torch.float32) + gate_input = calibration * 19.321 + up_input = calibration * 20.123 + down_input = calibration * 31.345 + sources, scales = [], [] + for expert in range(experts): + pairs = [ + matrix(r, c, group_size=16, seed=expert * 17 + projection) + for projection, (r, c) in enumerate( + ((intermediate, hidden), (intermediate, hidden), (hidden, intermediate)) + ) + ] + sources.append( + tuple( + replace(p[0], global_scale=calibration[expert], input_scale=a[expert]) + for p, a in zip(pairs, (up_input, gate_input, down_input), strict=True) + ) + ) + scales.append(tuple(pair[1] for pair in pairs)) + scratch13 = torch.empty( + (experts, 2 * local, hidden // 16), device=device, dtype=torch.float8_e4m3fn + ) + scratch2 = torch.empty( + (experts, hidden, local // 16), device=device, dtype=torch.float8_e4m3fn + ) + default = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + weights = load_nvfp4_csf_weights( + iter(sources), + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + tp_rank=rank, + tp_size=2, + device=device, + w13_scale_scratch=scratch13, + w2_scale_scratch=scratch2, + ) + finally: + torch.set_default_dtype(default) + for name, expected in ( + ("w13_global_scales", calibration.to(device)), + ("w2_global_scales", calibration.to(device)), + ("input_scale", up_input.to(device).reciprocal()), + ("intermediate_scale", down_input.to(device).reciprocal()), + ): + actual = getattr(weights.packed, name) + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + first, last = rank * local, (rank + 1) * local + assert torch.equal( + weights.packed.w13.cpu(), + torch.stack( + [ + torch.cat((up.weight[first:last, :], gate.weight[first:last, :])) + for up, gate, _ in sources + ] + ), + ) + assert torch.equal( + weights.packed.w2.cpu(), + torch.stack([down.weight[:, first // 2 : last // 2] for _, _, down in sources]), + ) + assert weights.packed.w13_block_scales is scratch13 + assert weights.packed.w2_block_scales is scratch2 + expected13 = torch.stack( + [torch.cat((s[0][first:last], s[1][first:last])) for s in scales] + ) + expected2 = torch.stack([s[2][:, first // 16 : last // 16] for s in scales]) + for planes, expected in ( + (weights.w13_scales, expected13), + (weights.w2_scales, expected2), + ): + actual = decode_planes( + planes, expected.shape[1], expected.shape[2], group_size=16 + ) + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize( + "bad", + [ + None, + torch.tensor(1.0, dtype=torch.bfloat16), + torch.tensor(float("nan")), + torch.tensor(0.0), + ], +) +def test_missing_or_invalid_calibration_is_rejected(bad): + from vllm.model_executor.model_loader.nvfp4_csf_loader import ( + _load_nvfp4_csf_weights as load_nvfp4_csf_weights, + ) + + source, _ = matrix(128, 128, group_size=16, seed=0) + source = replace(source, global_scale=bad, input_scale=torch.tensor(1.0)) + with pytest.raises(ValueError, match="global_scale"): + load_nvfp4_csf_weights( + [(source, source, source)], + num_experts=1, + hidden_size=128, + intermediate_size=128, + tp_rank=0, + tp_size=1, + device="cpu", + w13_scale_scratch=None, + w2_scale_scratch=None, + ) diff --git a/vllm/model_executor/model_loader/csf_utils.py b/vllm/model_executor/model_loader/csf_utils.py index 8ad5b0a8ec7d..52926815d1d7 100644 --- a/vllm/model_executor/model_loader/csf_utils.py +++ b/vllm/model_executor/model_loader/csf_utils.py @@ -4,11 +4,56 @@ import json from contextlib import ExitStack +from dataclasses import dataclass from pathlib import Path +from typing import Protocol +import torch from safetensors import safe_open +class PackedTensorView(Protocol): + """CPU tensor slicing interface provided by safetensors and tensor stores.""" + + def get_shape(self) -> list[int]: ... + + def get_dtype(self) -> str: ... + + def __getitem__(self, extent: tuple[slice, slice]) -> torch.Tensor: ... + + +@dataclass(frozen=True) +class CsfMatrix: + """One projection's lazy FP4 bytes and resident CPU scale components. + + NVFP4 also supplies the original FP32 scalar weight and activation calibration. + MXFP4 leaves those fields unset. Keep the backing tensor store open until the + layer reader returns. + """ + + weight: PackedTensorView + fixed: torch.Tensor + exceptions: torch.Tensor + global_scale: torch.Tensor | None = None + input_scale: torch.Tensor | None = None + + +def tp_extent(intermediate_size: int, tp_rank: int, tp_size: int, alignment: int): + """Validate an equal TP partition and return its intermediate-axis bounds.""" + if ( + tp_size <= 0 + or not 0 <= tp_rank < tp_size + or intermediate_size <= 0 + or intermediate_size % (alignment * tp_size) + ): + raise ValueError( + f"CSF requires a valid TP rank and extents aligned to {alignment} " + "intermediate channels" + ) + local = intermediate_size // tp_size + return tp_rank * local, (tp_rank + 1) * local + + def read_csf_contract(root, *, schema, codec, families): """Check storage identity and the source tensor-to-shard inventory.""" root = Path(root) @@ -35,7 +80,7 @@ def read_csf_contract(root, *, schema, codec, families): class CsfTensorReader(ExitStack): - """Keep a layer's lazily opened shards alive while B12X consumes its tensors.""" + """Keep a layer's lazily opened shards alive while slicing its tensors.""" def __init__(self, root, source_names, codec): super().__init__() @@ -56,8 +101,6 @@ def _handle(self, name): def matrix(self, weight, scale, *, global_scale=None, input_scale=None): """Expose a lazy packed weight view and one projection's CPU scale data.""" - from b12x.moe.checkpoints.csf import CsfMatrix - scales = self._handle(scale) return CsfMatrix( weight=self._handle(weight).get_slice(weight), diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py index 0e1dd4152ad0..e060dbec4b2c 100644 --- a/vllm/model_executor/model_loader/mxfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -8,15 +8,20 @@ """ import time +from collections.abc import Iterable from functools import lru_cache from pathlib import Path +import numpy as np import regex as re +import torch from safetensors import safe_open from vllm.model_executor.model_loader.csf_utils import ( + CsfMatrix, CsfTensorReader, read_csf_contract, + tp_extent, ) from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( @@ -51,9 +56,7 @@ def read_mxfp4_csf_layer( w13_scale_scratch, w2_scale_scratch, ): - """Resolve model tensor names and pass gate/up/down sources to B12X.""" - from b12x.moe.checkpoints.mxfp4_csf import load_mxfp4_csf_weights - + """Read rank-local gate/up/down tensors and unprepared compressed scales.""" contract = checkpoint_contract(str(Path(root).resolve())) family = contract["family"] if (num_experts, hidden_size, intermediate_size) != FAMILIES[family]: @@ -81,7 +84,7 @@ def experts(): for p in ("w1", "w3", "w2") ) - return load_mxfp4_csf_weights( + return _load_mxfp4_csf_weights( experts(), num_experts=num_experts, hidden_size=hidden_size, @@ -94,6 +97,137 @@ def experts(): ) +POSITION_MASK = (1 << 24) - 1 + + +def _slice_scale_plane(fixed, exceptions, rows, columns, row_slice, column_slice): + """Slice compressed bytes without floating-point reconstruction or fitting.""" + r0, r1 = row_slice + c0, c1 = column_slice + if not (0 <= r0 < r1 <= rows and r0 % 16 == r1 % 16 == 0): + raise ValueError("MXFP4-CSF row slices must contain complete 16-row slabs") + if not 0 <= c0 < c1 <= columns: + raise ValueError("MXFP4-CSF column slice is outside the scale plane") + if fixed.dtype != torch.uint8 or exceptions.dtype != torch.uint32: + raise TypeError("MXFP4-CSF requires uint8 fixed bytes and uint32 exceptions") + selectors = (columns + 7) // 8 + stream = fixed.numpy().reshape(rows // 16, 16 * (1 + selectors)) + bases = stream[:, :16] + if (bases > 254).any(): + raise ValueError("MXFP4-CSF palette bases must be in 0..254") + bits = np.unpackbits( + stream[:, 16:].reshape(rows, selectors), axis=1, bitorder="little" + ) + if bits[:, columns:].any(): + raise ValueError("MXFP4-CSF unused selector bits must be zero") + selected = np.packbits(bits[r0:r1, c0:c1], axis=1, bitorder="little") + result = np.concatenate( + (bases[r0 // 16 : r1 // 16], selected.reshape((r1 - r0) // 16, -1)), 1 + ) + words = exceptions.numpy().reshape(-1) + positions = words & POSITION_MASK + if len(words) and ( + positions[-1] >= rows * columns or (positions[1:] <= positions[:-1]).any() + ): + raise ValueError( + "MXFP4-CSF exception positions must be unique, sorted and in range" + ) + rr, cc = positions // columns, positions % columns + keep = (rr >= r0) & (rr < r1) & (cc >= c0) & (cc < c1) + positions = (rr[keep] - r0) * (c1 - c0) + cc[keep] - c0 + words = (words[keep] & np.uint32(0xFF000000)) | positions + return torch.from_numpy(result.copy()), torch.from_numpy(words.astype(np.uint32)) + + +def _load_mxfp4_csf_weights( + experts: Iterable[tuple[CsfMatrix, CsfMatrix, CsfMatrix]], + *, + num_experts, + hidden_size, + intermediate_size, + tp_rank, + tp_size, + device, + w13_scale_scratch, + w2_scale_scratch, +): + """Slice and upload gate/up/down-ordered expert projections. + + The iterable must yield exactly ``num_experts`` projection triples. Tensor + stores, manifests and model-specific tensor names belong to the caller. + Expanded scale buffers remain caller-owned for serialized layer execution. + """ + from b12x.moe.fused_moe import CsfScalePlanes, Mxfp4CsfWeights + + if num_experts <= 0 or hidden_size <= 0 or hidden_size % 64: + raise ValueError("MXFP4-CSF requires experts and 64-aligned hidden channels") + first, last = tp_extent(intermediate_size, tp_rank, tp_size, 32) + local = last - first + w13 = torch.empty( + (num_experts, 2 * local, hidden_size // 2), dtype=torch.uint8, device="cpu" + ) + w2 = torch.empty( + (num_experts, hidden_size, local // 2), dtype=torch.uint8, device="cpu" + ) + fixed13, fixed2, exceptions13, exceptions2 = [], [], [], [] + for expert, (first_projection, second_projection, down) in zip( + range(num_experts), experts, strict=True + ): + f13, e13 = [], [] + for matrix, projection in enumerate( + (first_projection, second_projection, down) + ): + view = projection.weight + expected = ( + [intermediate_size, hidden_size // 2] + if matrix < 2 + else [hidden_size, intermediate_size // 2] + ) + if view.get_shape() != expected or view.get_dtype() not in ("I8", "U8"): + raise ValueError( + f"MXFP4-CSF nibble geometry/dtype mismatch: " + f"expert={expert}, projection={matrix}" + ) + if matrix < 2: + w13[expert, matrix * local : (matrix + 1) * local].copy_( + view[first:last, :].view(torch.uint8) + ) + rows, columns = intermediate_size, hidden_size // 32 + row_slice, column_slice = (first, last), (0, columns) + else: + w2[expert].copy_(view[:, first // 2 : last // 2].view(torch.uint8)) + rows, columns = hidden_size, intermediate_size // 32 + row_slice, column_slice = (0, rows), (first // 32, last // 32) + fixed, exceptions = _slice_scale_plane( + projection.fixed, + projection.exceptions, + rows, + columns, + row_slice, + column_slice, + ) + if matrix < 2: + f13.append(fixed) + if matrix: + words = exceptions.numpy().copy() + words += np.uint32(local * columns) + exceptions = torch.from_numpy(words) + e13.append(exceptions) + else: + fixed2.append(fixed) + exceptions2.append(exceptions) + fixed13.append(torch.cat(f13)) + exceptions13.append(torch.cat(e13)) + return Mxfp4CsfWeights( + w13=w13.to(device), + w2=w2.to(device), + w13_scales=CsfScalePlanes(tuple(fixed13), tuple(exceptions13)), + w2_scales=CsfScalePlanes(tuple(fixed2), tuple(exceptions2)), + w13_scale_scratch=w13_scale_scratch, + w2_scale_scratch=w2_scale_scratch, + ) + + class Mxfp4CsfModelLoader(DefaultModelLoader): def _root(self, model_config): text_config = getattr(model_config, "hf_text_config", model_config.hf_config) diff --git a/vllm/model_executor/model_loader/nvfp4_csf_loader.py b/vllm/model_executor/model_loader/nvfp4_csf_loader.py index b2f60ebf1185..392d159c8cfd 100644 --- a/vllm/model_executor/model_loader/nvfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/nvfp4_csf_loader.py @@ -3,15 +3,20 @@ """Load retained tensors while the quantization method owns compressed experts.""" import time +from collections.abc import Iterable from functools import lru_cache from pathlib import Path +import numpy as np import regex as re +import torch from safetensors import safe_open from vllm.model_executor.model_loader.csf_utils import ( + CsfMatrix, CsfTensorReader, read_csf_contract, + tp_extent, ) from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( @@ -46,9 +51,7 @@ def read_nvfp4_csf_layer( w13_scale_scratch, w2_scale_scratch, ): - """Resolve model tensor names and pass up/gate/down sources to B12X.""" - from b12x.moe.checkpoints.nvfp4_csf import load_nvfp4_csf_weights - + """Read rank-local up/gate/down tensors and unprepared compressed scales.""" contract = checkpoint_contract(str(Path(root).resolve())) e, h, n, layers = FAMILIES[contract["family"]] if (num_experts, hidden_size, intermediate_size) != (e, h, n): @@ -72,7 +75,7 @@ def experts(): for p in ("up_proj", "gate_proj", "down_proj") ) - return load_nvfp4_csf_weights( + return _load_nvfp4_csf_weights( experts(), num_experts=num_experts, hidden_size=hidden_size, @@ -85,6 +88,165 @@ def experts(): ) +def _slice_scale_plane(fixed, exceptions, rows, columns, row_slice, column_slice): + """Copy aligned compressed extents and rebase exception positions.""" + r0, r1 = row_slice + c0, c1 = column_slice + if not (0 <= r0 < r1 <= rows and r0 % 16 == r1 % 16 == 0): + raise ValueError("NVFP4-CSF row slices require complete 16-row slabs") + if not (0 <= c0 < c1 <= columns and c0 % 2 == c1 % 2 == columns % 2 == 0): + raise ValueError("NVFP4-CSF column slices require complete nibble pairs") + if fixed.dtype != torch.uint8 or exceptions.dtype != torch.uint32: + raise TypeError("NVFP4-CSF requires uint8 fixed and uint32 exception tensors") + stream = fixed.numpy().reshape(rows // 16, 16 * (1 + columns // 2)) + bases = stream[:, :16] + if np.any(bases > 240): + raise ValueError("NVFP4-CSF row bases must be in 0..240") + packed = stream[:, 16:].reshape(rows, columns // 2) + selected = packed[r0:r1, c0 // 2 : c1 // 2] + result = np.concatenate( + (bases[r0 // 16 : r1 // 16], selected.reshape((r1 - r0) // 16, -1)), 1 + ) + words = exceptions.numpy().reshape(-1) + positions = words & np.uint32(0xFFFFFF) + if len(words) and ( + positions[-1] >= rows * columns or np.any(positions[1:] <= positions[:-1]) + ): + raise ValueError("NVFP4-CSF exception positions must be sorted and unique") + rr, cc = positions // columns, positions % columns + keep = (rr >= r0) & (rr < r1) & (cc >= c0) & (cc < c1) + local_positions = (rr[keep] - r0) * (c1 - c0) + cc[keep] - c0 + selected_words = (words[keep] & np.uint32(0xFF000000)) | local_positions + return result.copy(), selected_words.astype(np.uint32) + + +def _load_nvfp4_csf_weights( + experts: Iterable[tuple[CsfMatrix, CsfMatrix, CsfMatrix]], + *, + num_experts, + hidden_size, + intermediate_size, + tp_rank, + tp_size, + device, + w13_scale_scratch, + w2_scale_scratch, +): + """Slice and upload up/gate/down-ordered expert projections. + + The iterable must yield exactly ``num_experts`` projection triples. Tensor + stores, manifests and model-specific tensor names belong to the caller. + Expanded scale buffers remain caller-owned for serialized layer execution. + """ + from b12x.moe.fused_moe import CsfScalePlanes, Nvfp4CsfWeights, PackedWeights + + if num_experts <= 0 or hidden_size <= 0 or hidden_size % 128: + raise ValueError("NVFP4-CSF requires experts and 128-aligned hidden channels") + first, last = tp_extent(intermediate_size, tp_rank, tp_size, 64) + local = last - first + w13 = torch.empty( + (num_experts, 2 * local, hidden_size // 2), dtype=torch.uint8, device="cpu" + ) + w2 = torch.empty( + (num_experts, hidden_size, local // 2), dtype=torch.uint8, device="cpu" + ) + # Model loaders may set BF16 as the default dtype. Calibration belongs to + # the source FP32 contract and must not be rounded with the model weights. + g13, g2 = ( + torch.empty(num_experts, dtype=torch.float32, device="cpu") for _ in range(2) + ) + a13, a2 = ( + torch.empty(num_experts, dtype=torch.float32, device="cpu") for _ in range(2) + ) + fixed13, fixed2, exceptions13, exceptions2 = [], [], [], [] + + def scalar(value, role): + if value is None or value.numel() != 1 or value.dtype != torch.float32: + raise ValueError(f"NVFP4 {role} must be one FP32 value") + if not bool(torch.isfinite(value).all() and (value > 0).all()): + raise ValueError(f"NVFP4 {role} must be positive and finite") + return value.reshape(()) + + for expert, (first_projection, second_projection, down) in zip( + range(num_experts), experts, strict=True + ): + f13, e13, global13, input13 = [], [], [], [] + for matrix, projection in enumerate( + (first_projection, second_projection, down) + ): + view = projection.weight + expected = ( + [intermediate_size, hidden_size // 2] + if matrix < 2 + else [hidden_size, intermediate_size // 2] + ) + if view.get_shape() != expected or view.get_dtype() != "U8": + raise ValueError( + f"NVFP4 nibble geometry/dtype mismatch: " + f"expert={expert}, projection={matrix}" + ) + global_scale, input_scale = ( + scalar(projection.global_scale, "global_scale"), + scalar(projection.input_scale, "input_scale"), + ) + if matrix < 2: + w13[expert, matrix * local : (matrix + 1) * local].copy_( + view[first:last, :] + ) + rows, columns = intermediate_size, hidden_size // 16 + row_slice, column_slice = (first, last), (0, columns) + global13.append(global_scale) + input13.append(input_scale) + else: + w2[expert].copy_(view[:, first // 2 : last // 2]) + rows, columns = hidden_size, intermediate_size // 16 + row_slice, column_slice = (0, rows), (first // 16, last // 16) + g2[expert], a2[expert] = global_scale, input_scale + fixed, exceptions = _slice_scale_plane( + projection.fixed, + projection.exceptions, + rows, + columns, + row_slice, + column_slice, + ) + if matrix < 2: + f13.append(fixed) + if matrix: + exceptions += np.uint32(local * columns) + e13.append(exceptions) + else: + fixed2.append(fixed) + exceptions2.append(exceptions) + if not torch.equal(global13[0], global13[1]): + raise ValueError("NVFP4-CSF requires equal gate/up global weight scales") + g13[expert] = global13[0] + a13[expert] = torch.stack(input13).amax() + fixed13.append(np.concatenate(f13)) + exceptions13.append(np.concatenate(e13)) + return Nvfp4CsfWeights( + packed=PackedWeights( + w13=w13.to(device), + w2=w2.to(device), + w13_block_scales=w13_scale_scratch, + w2_block_scales=w2_scale_scratch, + w13_global_scales=g13.to(device), + w2_global_scales=g2.to(device), + input_scale=a13.to(device).reciprocal(), + intermediate_scale=a2.to(device).reciprocal(), + immutable_input_scales=True, + ), + w13_scales=CsfScalePlanes( + tuple(torch.from_numpy(p) for p in fixed13), + tuple(torch.from_numpy(p) for p in exceptions13), + ), + w2_scales=CsfScalePlanes( + tuple(torch.from_numpy(p) for p in fixed2), + tuple(torch.from_numpy(p) for p in exceptions2), + ), + ) + + class Nvfp4CsfModelLoader(DefaultModelLoader): def _root(self, model_config): quant = getattr(model_config.hf_config, "quantization_config", None) From a5430f5f25b681cf0b3f35bd339e1d1f21d47fb1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 1 Oct 2026 20:57:48 +0000 Subject: [PATCH 342/365] Match DS4.1 CSF expert activations to native MXFP8 policy Signed-off-by: Martin Vit --- .lil/changes/vllm-956-csf-mxfp8.json | 15 ++++++++ tests/quantization/test_mxfp4_csf.py | 38 ++++++++++++++------ vllm/model_executor/layers/fused_moe/b12x.py | 8 ++++- vllm/models/deepseek_v4_1/mxfp4_csf.py | 22 ++++++++---- 4 files changed, 65 insertions(+), 18 deletions(-) create mode 100644 .lil/changes/vllm-956-csf-mxfp8.json diff --git a/.lil/changes/vllm-956-csf-mxfp8.json b/.lil/changes/vllm-956-csf-mxfp8.json new file mode 100644 index 000000000000..a051274d23aa --- /dev/null +++ b/.lil/changes/vllm-956-csf-mxfp8.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-956-csf-mxfp8", + "category": "fix", + "summary": "Preserve native MXFP8 expert activations for DS4.1 CSF checkpoints", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "Install the paired B12X MXFP8-CSF preparation update. Checkpoint files and launch flags are unchanged. Expert results can differ from the previous BF16 CSF execution because the native activation quantizer is restored.", + "details": [ + "DS4.1 CSF uses MXFP8 expert activations unless VLLM_B12X_MOE_FP4_FORCE_A16=1 requests BF16, matching the native MXFP4 backend policy.", + "Attention precision, logits precision, dense layers and Kimi BF16 expert activations are unchanged.", + "Checkpoint loading remains in vLLM and invokes ordinary B12X weight preparation." + ], + "pull_requests": [956], + "requires": ["b12x-450-csf-mxfp8", "vllm-956-csf-preparation"] +} diff --git a/tests/quantization/test_mxfp4_csf.py b/tests/quantization/test_mxfp4_csf.py index 72e1c45526a5..dedcc69c9b6a 100644 --- a/tests/quantization/test_mxfp4_csf.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -106,12 +106,15 @@ def fail_reader(*args, **kwargs): @pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) +@pytest.mark.parametrize("force_a16", [False, True]) def test_target_uses_csf_while_dense_and_draft_keep_native_methods( - monkeypatch, moe, owner, mode + monkeypatch, moe, owner, mode, force_a16 ): from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.models.deepseek_v4_1 import mxfp4_csf + monkeypatch.setattr(mxfp4_csf.envs, "VLLM_B12X_MOE_FP4_FORCE_A16", force_a16) + config = SimpleNamespace( model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), parallel_config=SimpleNamespace( @@ -132,9 +135,11 @@ def test_target_uses_csf_while_dense_and_draft_keep_native_methods( layer = Mock(spec=RoutedExperts) layer.moe_config = moe if mode == "normal": - assert isinstance( - owner.get_quant_method(layer, "model.layers.39.ffn.experts"), - Mxfp4CsfMoEMethod, + method = owner.get_quant_method(layer, "model.layers.39.ffn.experts") + assert isinstance(method, Mxfp4CsfMoEMethod) + assert method.activation_mode == ("a16" if force_a16 else "a8") + assert method.get_fused_moe_quant_config(layer).quant_dtype == ( + None if force_a16 else "mxfp8" ) elif mode == "wrong_loader": with pytest.raises(ValueError, match="load-format"): @@ -155,8 +160,9 @@ def test_exact_experts_reject_non_tp_execution(moe, owner, field, value): Mxfp4CsfMoEMethod(replace(moe, moe_parallel_config=parallel), owner) +@pytest.mark.parametrize("activation_mode", ["a16", "a8"]) def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( - monkeypatch, moe, owner + monkeypatch, moe, owner, activation_mode ): from b12x.moe import fused_moe @@ -165,11 +171,16 @@ def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( layer = torch.nn.Module() layer.activation = moe.activation layer.apply_router_weight_on_input = False - quant = Mxfp4CsfMoEMethod(moe, owner).get_fused_moe_quant_config(layer) + quant = Mxfp4CsfMoEMethod( + moe, owner, activation_mode=activation_mode + ).get_fused_moe_quant_config(layer) backend = b12x.B12xExperts(moe, quant) prepared = Mock(spec=fused_moe.PreparedExperts) activation = fused_moe.ActivationSpec( - mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16, swiglu_limit=10.0 + mode=activation_mode, + nonlinearity="silu", + io_dtype=torch.bfloat16, + swiglu_limit=10.0, ) prepared.plan = SimpleNamespace( source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), @@ -184,6 +195,11 @@ def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( backend.install_prepared_experts(layer, prepared) assert layer._b12x_prepared_experts is prepared assert backend._prepared_experts is prepared + prepared.plan.activation = replace( + activation, mode="a16" if activation_mode == "a8" else "a8" + ) + with pytest.raises(ValueError, match="activation"): + backend.install_prepared_experts(layer, prepared) prepared.plan.activation = replace(activation, swiglu_limit=None) with pytest.raises(ValueError, match="activation"): backend.install_prepared_experts(layer, prepared) @@ -270,10 +286,10 @@ def test_kimi_mxfp4_csf_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): activation_situ_linear_beta=25.0, swiglu_limit=None, ) - assert isinstance( - owner.get_quant_method(experts, "model.layers.1.mlp.experts"), - Mxfp4CsfMoEMethod, - ) + method = owner.get_quant_method(experts, "model.layers.1.mlp.experts") + assert isinstance(method, Mxfp4CsfMoEMethod) + assert method.activation_mode == "a16" + assert method.get_fused_moe_quant_config(experts).quant_dtype is None assert isinstance( owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), UnquantizedLinearMethod, diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 7031d61d906d..b12b33e12c1d 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -526,7 +526,13 @@ def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> Non or prepared.intermediate_size != self.moe_config.intermediate_size_per_partition or prepared.plan.activation.mode - != ("a4" if nvfp4 and self._quant_mode == "nvfp4" else "a16") + != ( + "a4" + if nvfp4 and self._quant_mode == "nvfp4" + else "a8" + if packed and self._quant_mode == "w4a8_mx" + else "a16" + ) or (trellis and prepared.plan.activation.rotation_dtype != torch.float16) or prepared.plan.activation.io_dtype != self.moe_config.in_dtype or prepared.plan.activation.nonlinearity diff --git a/vllm/models/deepseek_v4_1/mxfp4_csf.py b/vllm/models/deepseek_v4_1/mxfp4_csf.py index 40290f7da1a1..a156b664ef33 100644 --- a/vllm/models/deepseek_v4_1/mxfp4_csf.py +++ b/vllm/models/deepseek_v4_1/mxfp4_csf.py @@ -5,6 +5,7 @@ import regex as re import torch +import vllm.envs as envs from vllm.config import get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -87,16 +88,23 @@ def get_quant_method(self, layer, prefix): ) if config.load_config.load_format not in ("mxfp4_csf",): raise ValueError("MXFP4-CSF requires --load-format mxfp4_csf ") - return Mxfp4CsfMoEMethod(layer.moe_config, self) + return Mxfp4CsfMoEMethod( + layer.moe_config, + self, + activation_mode="a16" if envs.VLLM_B12X_MOE_FP4_FORCE_A16 else "a8", + ) return super().get_quant_method(layer, prefix) class Mxfp4CsfMoEMethod(FusedMoEMethodBase): """Expand scales for routed experts into serialized, model-owned scratch.""" - def __init__(self, moe, owner): + def __init__(self, moe, owner, *, activation_mode="a16"): super().__init__(moe) self.owner = owner + if activation_mode not in ("a16", "a8"): + raise ValueError("MXFP4-CSF requires A16 or MXFP8 activations") + self.activation_mode = activation_mode parallel = moe.moe_parallel_config if ( parallel.use_ep @@ -143,9 +151,10 @@ def create_weights( ) def get_fused_moe_quant_config(self, layer): + activation = "mxfp8" if self.activation_mode == "a8" else None return FusedMoEQuantConfig( - _a1=FusedMoEQuantDesc(), - _a2=FusedMoEQuantDesc(), + _a1=FusedMoEQuantDesc(dtype=activation), + _a2=FusedMoEQuantDesc(dtype=activation), _w1=FusedMoEQuantDesc(dtype="mxfp4"), _w2=FusedMoEQuantDesc(dtype="mxfp4"), ) @@ -193,7 +202,7 @@ def process_weights_after_loading(self, layer): plan = fused_moe.plan_weights( source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), activation=fused_moe.ActivationSpec( - mode="a16", + mode=self.activation_mode, nonlinearity="situ" if self.moe.activation == MoEActivation.SITU else "silu", @@ -212,11 +221,12 @@ def process_weights_after_loading(self, layer): MoEPrepareAndFinalizeNoDPEPModular(), backend ) logger.info( - "MXFP4-CSF lossless MXFP4 layer %d rank %d/%d: BF16 activations, " + "MXFP4-CSF lossless MXFP4 layer %d rank %d/%d: %s activations, " "compressed scales, shared scratch %d bytes", self.layer_index, rank, tp, + "MXFP8" if self.activation_mode == "a8" else "BF16", sum(t.numel() for t in scratch), ) From 09da98a7160870d201ca185e11e8d5b40599a632 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 2 Oct 2026 17:18:17 +0000 Subject: [PATCH 343/365] Quantize exact GLM Spark projections and MTP experts on load Add strict projection coverage, NVFP4 W4A16 draft quantization and bounded weight staging. Preserve source checkpoints. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit Co-Authored-By: Claude Opus 5.5 --- .../vllm-glm-spark-online-weights.json | 20 + .../quantization/glm53_spark_online.md | 66 +++ docs/features/quantization/online.md | 11 + .../quantization/glm53_spark_online.py | 66 +++ .../quantization/glm53_spark_targets.json | 544 ++++++++++++++++++ tests/config/test_config_utils.py | 6 +- .../model_loader/test_reload.py | 28 +- .../model_executor/test_eagle_quantization.py | 32 +- tests/models/test_glm5next_pooled_indexer.py | 27 + tests/quantization/test_online.py | 119 +++- .../test_quantization_config_args.py | 18 + vllm/config/quantization.py | 30 +- .../layers/quantization/online/base.py | 79 ++- .../layers/quantization/online/fp8.py | 1 + .../layers/quantization/online/moe_base.py | 3 + .../layers/quantization/online/nvfp4.py | 71 ++- .../model_loader/base_loader.py | 8 +- .../model_loader/reload/layerwise.py | 6 + vllm/model_executor/models/utils.py | 15 + vllm/models/glm5next/nvidia/pooled_indexer.py | 11 +- vllm/v1/spec_decode/llm_base_proposer.py | 3 +- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 3 +- 22 files changed, 1112 insertions(+), 55 deletions(-) create mode 100644 .lil/changes/vllm-glm-spark-online-weights.json create mode 100644 docs/features/quantization/glm53_spark_online.md create mode 100644 examples/features/quantization/glm53_spark_online.py create mode 100644 examples/features/quantization/glm53_spark_targets.json diff --git a/.lil/changes/vllm-glm-spark-online-weights.json b/.lil/changes/vllm-glm-spark-online-weights.json new file mode 100644 index 000000000000..788d782004d5 --- /dev/null +++ b/.lil/changes/vllm-glm-spark-online-weights.json @@ -0,0 +1,20 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-glm-spark-online-weights", + "category": "feature", + "summary": "Quantize the GLM Spark projection set and BF16 MTP experts during loading", + "models": [ + "GLM-5.3-Flash" + ], + "compatibility": "Opt-in exact target configurations; existing online quantization defaults are unchanged. NVFP4 W4A16 requires a supported BF16-activation MoE backend.", + "details": [ + "Strict targets reject missing projections, incomplete fused groups and independently enabled online quantizers outside the declared set.", + "The GLM example selects 531 MXFP8 projections and NVFP4 W4A16 MTP routed experts.", + "Layerwise loading releases replaced full-precision parameters before packed-weight preparation, and NVFP4 encoding writes directly into bounded output storage." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [], + "requires": [] +} diff --git a/docs/features/quantization/glm53_spark_online.md b/docs/features/quantization/glm53_spark_online.md new file mode 100644 index 000000000000..9a8d0d24cb3d --- /dev/null +++ b/docs/features/quantization/glm53_spark_online.md @@ -0,0 +1,66 @@ +# GLM-5.3-Flash online quantization with Spark projection selection + +Use two 96 GB RTX PRO 6000 Blackwell GPUs and a vLLM build containing online +`nvfp4_a16` and strict quantization targets. Run the following from the vLLM +repository inside that environment: + +```bash +python examples/features/quantization/glm53_spark_online.py \ + --output-dir glm-spark-config + +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +VLLM_USE_V2_MODEL_RUNNER=1 \ +VLLM_B12X_MOE_FP4_FORCE_A16=0 \ +VLLM_ENABLE_PCIE_ALLREDUCE=1 VLLM_PCIE_ALLREDUCE_BACKEND=b12x \ +VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE=2048 \ +VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE=auto \ +VLLM_GLM53_KDA_GATE_SIDE_STREAM=1 VLLM_DISABLE_SHARED_EXPERTS_STREAM=0 \ +VLLM_GLM53_L2_PREFETCH=1 VLLM_LM_HEAD_A16=1 \ +VLLM_GLM53_MTP_DRAFT_HEAD=bf16 \ +VLLM_MTP_NVFP4_LM_HEAD=0 VLLM_MXFP8_LM_HEAD=0 \ +VLLM_GLM53_VISION_MXFP8=0 VLLM_GLM53_FP8_DENSE=0 \ +VLLM_GLM53_ONLINE_DENSE_MXFP8=0 \ +vllm serve local-inference-lab/GLM-5.3-Flash-NVFP4 \ + --revision cfd47bd7680e68408924df09b179d5bed25b2ae9 \ + --host 0.0.0.0 --port 8000 \ + --tensor-parallel-size 2 --pipeline-parallel-size 1 \ + --quantization modelopt_mixed --load-format safetensors --dtype bfloat16 \ + --moe-backend b12x --linear-backend b12x --attention-backend B12X \ + --additional-config '{"glm53_kda_decode_backend":"auto","kda_prefill_backend":"b12x"}' \ + --quantization-config "$(cat glm-spark-config/main.json)" \ + --speculative-config "$(cat glm-spark-config/mtp.json)" \ + --kv-cache-dtype fp8 --kv-cache-memory-bytes 3221225472 \ + --gpu-memory-utilization 0.93 \ + --block-size 256 --mamba-cache-mode align \ + --max-model-len 8192 --max-num-batched-tokens 1024 --max-num-seqs 8 \ + --max-parallel-prefills 1 --enable-chunked-prefill --enable-prefix-caching \ + --recurrent-checkpoint-policy request_boundaries \ + --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' \ + --cudagraph-capture-sizes 1 2 4 8 16 32 64 --max-cudagraph-capture-size 64 \ + --no-enable-flashinfer-autotune --reasoning-parser glm45 +``` + +The configuration generator creates a directory and refuses to overwrite one +that already exists. The pinned source revision belongs to the `mtp-bf16` +branch. vLLM reads the checkpoint and converts weights in memory; it does not +write a quantized checkpoint. + +The adjacent projection manifest selects exactly 531 BF16 matrices from the +Spark recipe: 520 in the main model and 11 in MTP. These use MXFP8 weights and +the selected MXFP8 linear backend. MTP routed experts use NVFP4 weights with +BF16 activations. The LM head, embeddings, vision tower, routers, norms, +convolutions and the first three dense feed-forward layers retain their source +representation. Existing main-model NVFP4 experts retain their activation +recipe. + +`strict_targets` defaults to `false` in vLLM; both generated target maps set it +to `true`. A missing projection, incomplete fused group, or independently +enabled online quantizer causes startup to fail. Keep the head, vision and +dense quantization flags shown above disabled when reproducing this selection. +The recipe matches the Spark **layer selection**; it does not claim byte parity +with Spark's serialized quantized weights. + +MTP uses three speculative tokens by default in the generator. Change this +with `--num-speculative-tokens 1`, or omit `--speculative-config` from the serve +command to disable MTP. The 3 GiB cache allocation is per GPU; it is a bounded +8K-context example, not a million-token capacity configuration. diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index d3a4edd2710f..807875fa2a34 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -51,6 +51,7 @@ vllm serve Qwen/Qwen3.5-35B-A3B --quantization mxfp4 \ | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | | `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback | +| `nvfp4_a16` | fp4_e2m1 data, fp8_e4m3 per-1x16-block scale, FP32 global scale per expert | BF16 | MoE only; select a compatible W4A16 backend. B12X supports SM120. | | `mxfp4` | fp4_e2m1 data, e8m0 per-1x32-block scale ([OCP MX specs](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf)) | - linear: fp4_e2m1 data, e8m0 per-1x32-block scale in some backends, or BF16.
- MOE: fp4_e2m1 data, e8m0 per-1x32-block scale. | Linear MXFP4 backend is auto-selected per platform, not enforcing activation dtype. Some use BF16 activation. Use `--linear-backend` to pin one (e.g. `--linear-backend flashinfer`). | ## Advanced Configuration @@ -177,6 +178,16 @@ llm = LLM( ### Fine-Grained Per-Layer Quantization Schemes +For a fixed model recipe, set `"strict_targets": true` alongside `targets`. +The default is `false`. Strict mode accepts exact names, requires every target +to resolve before loading weights, and rejects independently enabled online +quantizers outside that set. Fused linears can be named by all their source +projections. Keep separate target maps for the main model and the draft; +targets outside a local pipeline partition also cause an error. + +The [GLM Spark online-quantization guide](glm53_spark_online.md) includes an +exact projection manifest and a complete TP2 launch example. + Use the `targets` parameter to apply different online shorthands to different layers, instead of one scheme applied everywhere via `linear`/`moe`. Keys are exact layer names, regex patterns (prefixed with `re:`), or patterns understood by [`fnmatch.fnmatch`](https://docs.python.org/3/library/fnmatch.html#fnmatch.fnmatch); values are shorthand names (`fp8_per_tensor`, `fp8_per_block`, `fp8_per_channel`, `mxfp8`, `int8_per_channel_weight_only`, `nvfp4_per_token`). Example: diff --git a/examples/features/quantization/glm53_spark_online.py b/examples/features/quantization/glm53_spark_online.py new file mode 100644 index 000000000000..5cf606d26d30 --- /dev/null +++ b/examples/features/quantization/glm53_spark_online.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Write exact GLM Spark projection and MTP online-quantization configurations. + +The adjacent manifest records the selected projections in the pinned Spark +checkpoint. This generates configuration only; vLLM quantizes BF16 weights +during loading and does not rewrite checkpoint files. +""" + +import argparse +import json +from pathlib import Path + +SOURCE_MODEL = "local-inference-lab/GLM-5.3-Flash-NVFP4" +SOURCE_REVISION = "cfd47bd7680e68408924df09b179d5bed25b2ae9" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model", default=SOURCE_MODEL) + parser.add_argument("--revision", default=SOURCE_REVISION) + parser.add_argument("--num-speculative-tokens", type=int, default=3) + args = parser.parse_args() + manifest = json.loads( + Path(__file__).with_name("glm53_spark_targets.json").read_text() + ) + main_targets, mtp_targets = {}, {} + source_prefix = "model.language_model." + for name in manifest["targets"]: + if not name.startswith(source_prefix): + raise ValueError(f"Unexpected source projection name: {name}") + if name.startswith(source_prefix + "layers.45."): + mtp_targets[name.replace(source_prefix, "model.", 1)] = "mxfp8" + else: + main_targets[name.replace(source_prefix, "language_model.model.", 1)] = ( + "mxfp8" + ) + if len(main_targets) != 520 or len(mtp_targets) != 11: + raise ValueError( + "The Spark manifest must define 520 main and 11 MTP projections" + ) + for name in manifest["nvfp4_mtp_targets"]: + mtp_targets[name.replace(source_prefix, "model.", 1)] = "nvfp4_a16" + configurations = { + "main.json": {"targets": main_targets, "strict_targets": True}, + "mtp.json": { + "method": "mtp", + "model": args.model, + "revision": args.revision, + "num_speculative_tokens": args.num_speculative_tokens, + "draft_sample_method": "probabilistic", + "rejection_sample_method": "standard", + "moe_backend": "b12x", + "attention_backend": "B12X", + "quantization_config": {"targets": mtp_targets, "strict_targets": True}, + }, + } + args.output_dir.mkdir(parents=True, exist_ok=False) + for filename, value in configurations.items(): + (args.output_dir / filename).write_text(json.dumps(value, indent=2) + "\n") + print(f"Wrote {args.output_dir}: 531 MXFP8 projections and NVFP4 W4A16 MTP experts") + + +if __name__ == "__main__": + main() diff --git a/examples/features/quantization/glm53_spark_targets.json b/examples/features/quantization/glm53_spark_targets.json new file mode 100644 index 000000000000..28ab696c23c5 --- /dev/null +++ b/examples/features/quantization/glm53_spark_targets.json @@ -0,0 +1,544 @@ +{ + "source": { + "repository": "local-inference-lab/GLM-5.3-Flash-NVFP4-Spark", + "revision": "a608241037e4c2565356bff7ca293f2133888f88" + }, + "config_sha256": "e1c0246a44ebefb5fd6383fb57aebbf7ac69ff6e7b23e989c0571b279a0eca23", + "targets": [ + "model.language_model.layers.0.self_attn.b_proj", + "model.language_model.layers.0.self_attn.f_a_proj", + "model.language_model.layers.0.self_attn.f_b_proj", + "model.language_model.layers.0.self_attn.g_a_proj", + "model.language_model.layers.0.self_attn.g_b_proj", + "model.language_model.layers.0.self_attn.k_proj", + "model.language_model.layers.0.self_attn.o_proj", + "model.language_model.layers.0.self_attn.q_proj", + "model.language_model.layers.0.self_attn.v_proj", + "model.language_model.layers.1.self_attn.b_proj", + "model.language_model.layers.1.self_attn.f_a_proj", + "model.language_model.layers.1.self_attn.f_b_proj", + "model.language_model.layers.1.self_attn.g_a_proj", + "model.language_model.layers.1.self_attn.g_b_proj", + "model.language_model.layers.1.self_attn.k_proj", + "model.language_model.layers.1.self_attn.o_proj", + "model.language_model.layers.1.self_attn.q_proj", + "model.language_model.layers.1.self_attn.v_proj", + "model.language_model.layers.10.mlp.shared_experts.down_proj", + "model.language_model.layers.10.mlp.shared_experts.gate_proj", + "model.language_model.layers.10.mlp.shared_experts.up_proj", + "model.language_model.layers.10.self_attn.b_proj", + "model.language_model.layers.10.self_attn.f_a_proj", + "model.language_model.layers.10.self_attn.f_b_proj", + "model.language_model.layers.10.self_attn.g_a_proj", + "model.language_model.layers.10.self_attn.g_b_proj", + "model.language_model.layers.10.self_attn.k_proj", + "model.language_model.layers.10.self_attn.o_proj", + "model.language_model.layers.10.self_attn.q_proj", + "model.language_model.layers.10.self_attn.v_proj", + "model.language_model.layers.11.mlp.shared_experts.down_proj", + "model.language_model.layers.11.mlp.shared_experts.gate_proj", + "model.language_model.layers.11.mlp.shared_experts.up_proj", + "model.language_model.layers.11.self_attn.indexer.weights_proj", + "model.language_model.layers.11.self_attn.indexer.wk", + "model.language_model.layers.11.self_attn.indexer.wq_b", + "model.language_model.layers.11.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.11.self_attn.kv_b_proj", + "model.language_model.layers.11.self_attn.o_proj", + "model.language_model.layers.11.self_attn.q_a_proj", + "model.language_model.layers.11.self_attn.q_b_proj", + "model.language_model.layers.12.mlp.shared_experts.down_proj", + "model.language_model.layers.12.mlp.shared_experts.gate_proj", + "model.language_model.layers.12.mlp.shared_experts.up_proj", + "model.language_model.layers.12.self_attn.b_proj", + "model.language_model.layers.12.self_attn.f_a_proj", + "model.language_model.layers.12.self_attn.f_b_proj", + "model.language_model.layers.12.self_attn.g_a_proj", + "model.language_model.layers.12.self_attn.g_b_proj", + "model.language_model.layers.12.self_attn.k_proj", + "model.language_model.layers.12.self_attn.o_proj", + "model.language_model.layers.12.self_attn.q_proj", + "model.language_model.layers.12.self_attn.v_proj", + "model.language_model.layers.13.mlp.shared_experts.down_proj", + "model.language_model.layers.13.mlp.shared_experts.gate_proj", + "model.language_model.layers.13.mlp.shared_experts.up_proj", + "model.language_model.layers.13.self_attn.b_proj", + "model.language_model.layers.13.self_attn.f_a_proj", + "model.language_model.layers.13.self_attn.f_b_proj", + "model.language_model.layers.13.self_attn.g_a_proj", + "model.language_model.layers.13.self_attn.g_b_proj", + "model.language_model.layers.13.self_attn.k_proj", + "model.language_model.layers.13.self_attn.o_proj", + "model.language_model.layers.13.self_attn.q_proj", + "model.language_model.layers.13.self_attn.v_proj", + "model.language_model.layers.14.mlp.shared_experts.down_proj", + "model.language_model.layers.14.mlp.shared_experts.gate_proj", + "model.language_model.layers.14.mlp.shared_experts.up_proj", + "model.language_model.layers.14.self_attn.b_proj", + "model.language_model.layers.14.self_attn.f_a_proj", + "model.language_model.layers.14.self_attn.f_b_proj", + "model.language_model.layers.14.self_attn.g_a_proj", + "model.language_model.layers.14.self_attn.g_b_proj", + "model.language_model.layers.14.self_attn.k_proj", + "model.language_model.layers.14.self_attn.o_proj", + "model.language_model.layers.14.self_attn.q_proj", + "model.language_model.layers.14.self_attn.v_proj", + "model.language_model.layers.15.mlp.shared_experts.down_proj", + "model.language_model.layers.15.mlp.shared_experts.gate_proj", + "model.language_model.layers.15.mlp.shared_experts.up_proj", + "model.language_model.layers.15.self_attn.indexer.weights_proj", + "model.language_model.layers.15.self_attn.indexer.wk", + "model.language_model.layers.15.self_attn.indexer.wq_b", + "model.language_model.layers.15.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.15.self_attn.kv_b_proj", + "model.language_model.layers.15.self_attn.o_proj", + "model.language_model.layers.15.self_attn.q_a_proj", + "model.language_model.layers.15.self_attn.q_b_proj", + "model.language_model.layers.16.mlp.shared_experts.down_proj", + "model.language_model.layers.16.mlp.shared_experts.gate_proj", + "model.language_model.layers.16.mlp.shared_experts.up_proj", + "model.language_model.layers.16.self_attn.b_proj", + "model.language_model.layers.16.self_attn.f_a_proj", + "model.language_model.layers.16.self_attn.f_b_proj", + "model.language_model.layers.16.self_attn.g_a_proj", + "model.language_model.layers.16.self_attn.g_b_proj", + "model.language_model.layers.16.self_attn.k_proj", + "model.language_model.layers.16.self_attn.o_proj", + "model.language_model.layers.16.self_attn.q_proj", + "model.language_model.layers.16.self_attn.v_proj", + "model.language_model.layers.17.mlp.shared_experts.down_proj", + "model.language_model.layers.17.mlp.shared_experts.gate_proj", + "model.language_model.layers.17.mlp.shared_experts.up_proj", + "model.language_model.layers.17.self_attn.b_proj", + "model.language_model.layers.17.self_attn.f_a_proj", + "model.language_model.layers.17.self_attn.f_b_proj", + "model.language_model.layers.17.self_attn.g_a_proj", + "model.language_model.layers.17.self_attn.g_b_proj", + "model.language_model.layers.17.self_attn.k_proj", + "model.language_model.layers.17.self_attn.o_proj", + "model.language_model.layers.17.self_attn.q_proj", + "model.language_model.layers.17.self_attn.v_proj", + "model.language_model.layers.18.mlp.shared_experts.down_proj", + "model.language_model.layers.18.mlp.shared_experts.gate_proj", + "model.language_model.layers.18.mlp.shared_experts.up_proj", + "model.language_model.layers.18.self_attn.b_proj", + "model.language_model.layers.18.self_attn.f_a_proj", + "model.language_model.layers.18.self_attn.f_b_proj", + "model.language_model.layers.18.self_attn.g_a_proj", + "model.language_model.layers.18.self_attn.g_b_proj", + "model.language_model.layers.18.self_attn.k_proj", + "model.language_model.layers.18.self_attn.o_proj", + "model.language_model.layers.18.self_attn.q_proj", + "model.language_model.layers.18.self_attn.v_proj", + "model.language_model.layers.19.mlp.shared_experts.down_proj", + "model.language_model.layers.19.mlp.shared_experts.gate_proj", + "model.language_model.layers.19.mlp.shared_experts.up_proj", + "model.language_model.layers.19.self_attn.indexer.weights_proj", + "model.language_model.layers.19.self_attn.indexer.wk", + "model.language_model.layers.19.self_attn.indexer.wq_b", + "model.language_model.layers.19.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.19.self_attn.kv_b_proj", + "model.language_model.layers.19.self_attn.o_proj", + "model.language_model.layers.19.self_attn.q_a_proj", + "model.language_model.layers.19.self_attn.q_b_proj", + "model.language_model.layers.2.self_attn.b_proj", + "model.language_model.layers.2.self_attn.f_a_proj", + "model.language_model.layers.2.self_attn.f_b_proj", + "model.language_model.layers.2.self_attn.g_a_proj", + "model.language_model.layers.2.self_attn.g_b_proj", + "model.language_model.layers.2.self_attn.k_proj", + "model.language_model.layers.2.self_attn.o_proj", + "model.language_model.layers.2.self_attn.q_proj", + "model.language_model.layers.2.self_attn.v_proj", + "model.language_model.layers.20.mlp.shared_experts.down_proj", + "model.language_model.layers.20.mlp.shared_experts.gate_proj", + "model.language_model.layers.20.mlp.shared_experts.up_proj", + "model.language_model.layers.20.self_attn.b_proj", + "model.language_model.layers.20.self_attn.f_a_proj", + "model.language_model.layers.20.self_attn.f_b_proj", + "model.language_model.layers.20.self_attn.g_a_proj", + "model.language_model.layers.20.self_attn.g_b_proj", + "model.language_model.layers.20.self_attn.k_proj", + "model.language_model.layers.20.self_attn.o_proj", + "model.language_model.layers.20.self_attn.q_proj", + "model.language_model.layers.20.self_attn.v_proj", + "model.language_model.layers.21.mlp.shared_experts.down_proj", + "model.language_model.layers.21.mlp.shared_experts.gate_proj", + "model.language_model.layers.21.mlp.shared_experts.up_proj", + "model.language_model.layers.21.self_attn.b_proj", + "model.language_model.layers.21.self_attn.f_a_proj", + "model.language_model.layers.21.self_attn.f_b_proj", + "model.language_model.layers.21.self_attn.g_a_proj", + "model.language_model.layers.21.self_attn.g_b_proj", + "model.language_model.layers.21.self_attn.k_proj", + "model.language_model.layers.21.self_attn.o_proj", + "model.language_model.layers.21.self_attn.q_proj", + "model.language_model.layers.21.self_attn.v_proj", + "model.language_model.layers.22.mlp.shared_experts.down_proj", + "model.language_model.layers.22.mlp.shared_experts.gate_proj", + "model.language_model.layers.22.mlp.shared_experts.up_proj", + "model.language_model.layers.22.self_attn.b_proj", + "model.language_model.layers.22.self_attn.f_a_proj", + "model.language_model.layers.22.self_attn.f_b_proj", + "model.language_model.layers.22.self_attn.g_a_proj", + "model.language_model.layers.22.self_attn.g_b_proj", + "model.language_model.layers.22.self_attn.k_proj", + "model.language_model.layers.22.self_attn.o_proj", + "model.language_model.layers.22.self_attn.q_proj", + "model.language_model.layers.22.self_attn.v_proj", + "model.language_model.layers.23.mlp.shared_experts.down_proj", + "model.language_model.layers.23.mlp.shared_experts.gate_proj", + "model.language_model.layers.23.mlp.shared_experts.up_proj", + "model.language_model.layers.23.self_attn.indexer.weights_proj", + "model.language_model.layers.23.self_attn.indexer.wk", + "model.language_model.layers.23.self_attn.indexer.wq_b", + "model.language_model.layers.23.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.23.self_attn.kv_b_proj", + "model.language_model.layers.23.self_attn.o_proj", + "model.language_model.layers.23.self_attn.q_a_proj", + "model.language_model.layers.23.self_attn.q_b_proj", + "model.language_model.layers.24.mlp.shared_experts.down_proj", + "model.language_model.layers.24.mlp.shared_experts.gate_proj", + "model.language_model.layers.24.mlp.shared_experts.up_proj", + "model.language_model.layers.24.self_attn.b_proj", + "model.language_model.layers.24.self_attn.f_a_proj", + "model.language_model.layers.24.self_attn.f_b_proj", + "model.language_model.layers.24.self_attn.g_a_proj", + "model.language_model.layers.24.self_attn.g_b_proj", + "model.language_model.layers.24.self_attn.k_proj", + "model.language_model.layers.24.self_attn.o_proj", + "model.language_model.layers.24.self_attn.q_proj", + "model.language_model.layers.24.self_attn.v_proj", + "model.language_model.layers.25.mlp.shared_experts.down_proj", + "model.language_model.layers.25.mlp.shared_experts.gate_proj", + "model.language_model.layers.25.mlp.shared_experts.up_proj", + "model.language_model.layers.25.self_attn.b_proj", + "model.language_model.layers.25.self_attn.f_a_proj", + "model.language_model.layers.25.self_attn.f_b_proj", + "model.language_model.layers.25.self_attn.g_a_proj", + "model.language_model.layers.25.self_attn.g_b_proj", + "model.language_model.layers.25.self_attn.k_proj", + "model.language_model.layers.25.self_attn.o_proj", + "model.language_model.layers.25.self_attn.q_proj", + "model.language_model.layers.25.self_attn.v_proj", + "model.language_model.layers.26.mlp.shared_experts.down_proj", + "model.language_model.layers.26.mlp.shared_experts.gate_proj", + "model.language_model.layers.26.mlp.shared_experts.up_proj", + "model.language_model.layers.26.self_attn.b_proj", + "model.language_model.layers.26.self_attn.f_a_proj", + "model.language_model.layers.26.self_attn.f_b_proj", + "model.language_model.layers.26.self_attn.g_a_proj", + "model.language_model.layers.26.self_attn.g_b_proj", + "model.language_model.layers.26.self_attn.k_proj", + "model.language_model.layers.26.self_attn.o_proj", + "model.language_model.layers.26.self_attn.q_proj", + "model.language_model.layers.26.self_attn.v_proj", + "model.language_model.layers.27.mlp.shared_experts.down_proj", + "model.language_model.layers.27.mlp.shared_experts.gate_proj", + "model.language_model.layers.27.mlp.shared_experts.up_proj", + "model.language_model.layers.27.self_attn.indexer.weights_proj", + "model.language_model.layers.27.self_attn.indexer.wk", + "model.language_model.layers.27.self_attn.indexer.wq_b", + "model.language_model.layers.27.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.27.self_attn.kv_b_proj", + "model.language_model.layers.27.self_attn.o_proj", + "model.language_model.layers.27.self_attn.q_a_proj", + "model.language_model.layers.27.self_attn.q_b_proj", + "model.language_model.layers.28.mlp.shared_experts.down_proj", + "model.language_model.layers.28.mlp.shared_experts.gate_proj", + "model.language_model.layers.28.mlp.shared_experts.up_proj", + "model.language_model.layers.28.self_attn.b_proj", + "model.language_model.layers.28.self_attn.f_a_proj", + "model.language_model.layers.28.self_attn.f_b_proj", + "model.language_model.layers.28.self_attn.g_a_proj", + "model.language_model.layers.28.self_attn.g_b_proj", + "model.language_model.layers.28.self_attn.k_proj", + "model.language_model.layers.28.self_attn.o_proj", + "model.language_model.layers.28.self_attn.q_proj", + "model.language_model.layers.28.self_attn.v_proj", + "model.language_model.layers.29.mlp.shared_experts.down_proj", + "model.language_model.layers.29.mlp.shared_experts.gate_proj", + "model.language_model.layers.29.mlp.shared_experts.up_proj", + "model.language_model.layers.29.self_attn.b_proj", + "model.language_model.layers.29.self_attn.f_a_proj", + "model.language_model.layers.29.self_attn.f_b_proj", + "model.language_model.layers.29.self_attn.g_a_proj", + "model.language_model.layers.29.self_attn.g_b_proj", + "model.language_model.layers.29.self_attn.k_proj", + "model.language_model.layers.29.self_attn.o_proj", + "model.language_model.layers.29.self_attn.q_proj", + "model.language_model.layers.29.self_attn.v_proj", + "model.language_model.layers.3.mlp.shared_experts.down_proj", + "model.language_model.layers.3.mlp.shared_experts.gate_proj", + "model.language_model.layers.3.mlp.shared_experts.up_proj", + "model.language_model.layers.3.self_attn.indexer.weights_proj", + "model.language_model.layers.3.self_attn.indexer.wk", + "model.language_model.layers.3.self_attn.indexer.wq_b", + "model.language_model.layers.3.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.3.self_attn.kv_b_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.self_attn.q_a_proj", + "model.language_model.layers.3.self_attn.q_b_proj", + "model.language_model.layers.30.mlp.shared_experts.down_proj", + "model.language_model.layers.30.mlp.shared_experts.gate_proj", + "model.language_model.layers.30.mlp.shared_experts.up_proj", + "model.language_model.layers.30.self_attn.b_proj", + "model.language_model.layers.30.self_attn.f_a_proj", + "model.language_model.layers.30.self_attn.f_b_proj", + "model.language_model.layers.30.self_attn.g_a_proj", + "model.language_model.layers.30.self_attn.g_b_proj", + "model.language_model.layers.30.self_attn.k_proj", + "model.language_model.layers.30.self_attn.o_proj", + "model.language_model.layers.30.self_attn.q_proj", + "model.language_model.layers.30.self_attn.v_proj", + "model.language_model.layers.31.mlp.shared_experts.down_proj", + "model.language_model.layers.31.mlp.shared_experts.gate_proj", + "model.language_model.layers.31.mlp.shared_experts.up_proj", + "model.language_model.layers.31.self_attn.indexer.weights_proj", + "model.language_model.layers.31.self_attn.indexer.wk", + "model.language_model.layers.31.self_attn.indexer.wq_b", + "model.language_model.layers.31.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.31.self_attn.kv_b_proj", + "model.language_model.layers.31.self_attn.o_proj", + "model.language_model.layers.31.self_attn.q_a_proj", + "model.language_model.layers.31.self_attn.q_b_proj", + "model.language_model.layers.32.mlp.shared_experts.down_proj", + "model.language_model.layers.32.mlp.shared_experts.gate_proj", + "model.language_model.layers.32.mlp.shared_experts.up_proj", + "model.language_model.layers.32.self_attn.b_proj", + "model.language_model.layers.32.self_attn.f_a_proj", + "model.language_model.layers.32.self_attn.f_b_proj", + "model.language_model.layers.32.self_attn.g_a_proj", + "model.language_model.layers.32.self_attn.g_b_proj", + "model.language_model.layers.32.self_attn.k_proj", + "model.language_model.layers.32.self_attn.o_proj", + "model.language_model.layers.32.self_attn.q_proj", + "model.language_model.layers.32.self_attn.v_proj", + "model.language_model.layers.33.mlp.shared_experts.down_proj", + "model.language_model.layers.33.mlp.shared_experts.gate_proj", + "model.language_model.layers.33.mlp.shared_experts.up_proj", + "model.language_model.layers.33.self_attn.b_proj", + "model.language_model.layers.33.self_attn.f_a_proj", + "model.language_model.layers.33.self_attn.f_b_proj", + "model.language_model.layers.33.self_attn.g_a_proj", + "model.language_model.layers.33.self_attn.g_b_proj", + "model.language_model.layers.33.self_attn.k_proj", + "model.language_model.layers.33.self_attn.o_proj", + "model.language_model.layers.33.self_attn.q_proj", + "model.language_model.layers.33.self_attn.v_proj", + "model.language_model.layers.34.mlp.shared_experts.down_proj", + "model.language_model.layers.34.mlp.shared_experts.gate_proj", + "model.language_model.layers.34.mlp.shared_experts.up_proj", + "model.language_model.layers.34.self_attn.b_proj", + "model.language_model.layers.34.self_attn.f_a_proj", + "model.language_model.layers.34.self_attn.f_b_proj", + "model.language_model.layers.34.self_attn.g_a_proj", + "model.language_model.layers.34.self_attn.g_b_proj", + "model.language_model.layers.34.self_attn.k_proj", + "model.language_model.layers.34.self_attn.o_proj", + "model.language_model.layers.34.self_attn.q_proj", + "model.language_model.layers.34.self_attn.v_proj", + "model.language_model.layers.35.mlp.shared_experts.down_proj", + "model.language_model.layers.35.mlp.shared_experts.gate_proj", + "model.language_model.layers.35.mlp.shared_experts.up_proj", + "model.language_model.layers.35.self_attn.indexer.weights_proj", + "model.language_model.layers.35.self_attn.indexer.wk", + "model.language_model.layers.35.self_attn.indexer.wq_b", + "model.language_model.layers.35.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.35.self_attn.kv_b_proj", + "model.language_model.layers.35.self_attn.o_proj", + "model.language_model.layers.35.self_attn.q_a_proj", + "model.language_model.layers.35.self_attn.q_b_proj", + "model.language_model.layers.36.mlp.shared_experts.down_proj", + "model.language_model.layers.36.mlp.shared_experts.gate_proj", + "model.language_model.layers.36.mlp.shared_experts.up_proj", + "model.language_model.layers.36.self_attn.b_proj", + "model.language_model.layers.36.self_attn.f_a_proj", + "model.language_model.layers.36.self_attn.f_b_proj", + "model.language_model.layers.36.self_attn.g_a_proj", + "model.language_model.layers.36.self_attn.g_b_proj", + "model.language_model.layers.36.self_attn.k_proj", + "model.language_model.layers.36.self_attn.o_proj", + "model.language_model.layers.36.self_attn.q_proj", + "model.language_model.layers.36.self_attn.v_proj", + "model.language_model.layers.37.mlp.shared_experts.down_proj", + "model.language_model.layers.37.mlp.shared_experts.gate_proj", + "model.language_model.layers.37.mlp.shared_experts.up_proj", + "model.language_model.layers.37.self_attn.b_proj", + "model.language_model.layers.37.self_attn.f_a_proj", + "model.language_model.layers.37.self_attn.f_b_proj", + "model.language_model.layers.37.self_attn.g_a_proj", + "model.language_model.layers.37.self_attn.g_b_proj", + "model.language_model.layers.37.self_attn.k_proj", + "model.language_model.layers.37.self_attn.o_proj", + "model.language_model.layers.37.self_attn.q_proj", + "model.language_model.layers.37.self_attn.v_proj", + "model.language_model.layers.38.mlp.shared_experts.down_proj", + "model.language_model.layers.38.mlp.shared_experts.gate_proj", + "model.language_model.layers.38.mlp.shared_experts.up_proj", + "model.language_model.layers.38.self_attn.b_proj", + "model.language_model.layers.38.self_attn.f_a_proj", + "model.language_model.layers.38.self_attn.f_b_proj", + "model.language_model.layers.38.self_attn.g_a_proj", + "model.language_model.layers.38.self_attn.g_b_proj", + "model.language_model.layers.38.self_attn.k_proj", + "model.language_model.layers.38.self_attn.o_proj", + "model.language_model.layers.38.self_attn.q_proj", + "model.language_model.layers.38.self_attn.v_proj", + "model.language_model.layers.39.mlp.shared_experts.down_proj", + "model.language_model.layers.39.mlp.shared_experts.gate_proj", + "model.language_model.layers.39.mlp.shared_experts.up_proj", + "model.language_model.layers.39.self_attn.indexer.weights_proj", + "model.language_model.layers.39.self_attn.indexer.wk", + "model.language_model.layers.39.self_attn.indexer.wq_b", + "model.language_model.layers.39.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.39.self_attn.kv_b_proj", + "model.language_model.layers.39.self_attn.o_proj", + "model.language_model.layers.39.self_attn.q_a_proj", + "model.language_model.layers.39.self_attn.q_b_proj", + "model.language_model.layers.4.mlp.shared_experts.down_proj", + "model.language_model.layers.4.mlp.shared_experts.gate_proj", + "model.language_model.layers.4.mlp.shared_experts.up_proj", + "model.language_model.layers.4.self_attn.b_proj", + "model.language_model.layers.4.self_attn.f_a_proj", + "model.language_model.layers.4.self_attn.f_b_proj", + "model.language_model.layers.4.self_attn.g_a_proj", + "model.language_model.layers.4.self_attn.g_b_proj", + "model.language_model.layers.4.self_attn.k_proj", + "model.language_model.layers.4.self_attn.o_proj", + "model.language_model.layers.4.self_attn.q_proj", + "model.language_model.layers.4.self_attn.v_proj", + "model.language_model.layers.40.mlp.shared_experts.down_proj", + "model.language_model.layers.40.mlp.shared_experts.gate_proj", + "model.language_model.layers.40.mlp.shared_experts.up_proj", + "model.language_model.layers.40.self_attn.b_proj", + "model.language_model.layers.40.self_attn.f_a_proj", + "model.language_model.layers.40.self_attn.f_b_proj", + "model.language_model.layers.40.self_attn.g_a_proj", + "model.language_model.layers.40.self_attn.g_b_proj", + "model.language_model.layers.40.self_attn.k_proj", + "model.language_model.layers.40.self_attn.o_proj", + "model.language_model.layers.40.self_attn.q_proj", + "model.language_model.layers.40.self_attn.v_proj", + "model.language_model.layers.41.mlp.shared_experts.down_proj", + "model.language_model.layers.41.mlp.shared_experts.gate_proj", + "model.language_model.layers.41.mlp.shared_experts.up_proj", + "model.language_model.layers.41.self_attn.b_proj", + "model.language_model.layers.41.self_attn.f_a_proj", + "model.language_model.layers.41.self_attn.f_b_proj", + "model.language_model.layers.41.self_attn.g_a_proj", + "model.language_model.layers.41.self_attn.g_b_proj", + "model.language_model.layers.41.self_attn.k_proj", + "model.language_model.layers.41.self_attn.o_proj", + "model.language_model.layers.41.self_attn.q_proj", + "model.language_model.layers.41.self_attn.v_proj", + "model.language_model.layers.42.mlp.shared_experts.down_proj", + "model.language_model.layers.42.mlp.shared_experts.gate_proj", + "model.language_model.layers.42.mlp.shared_experts.up_proj", + "model.language_model.layers.42.self_attn.b_proj", + "model.language_model.layers.42.self_attn.f_a_proj", + "model.language_model.layers.42.self_attn.f_b_proj", + "model.language_model.layers.42.self_attn.g_a_proj", + "model.language_model.layers.42.self_attn.g_b_proj", + "model.language_model.layers.42.self_attn.k_proj", + "model.language_model.layers.42.self_attn.o_proj", + "model.language_model.layers.42.self_attn.q_proj", + "model.language_model.layers.42.self_attn.v_proj", + "model.language_model.layers.43.mlp.shared_experts.down_proj", + "model.language_model.layers.43.mlp.shared_experts.gate_proj", + "model.language_model.layers.43.mlp.shared_experts.up_proj", + "model.language_model.layers.43.self_attn.indexer.weights_proj", + "model.language_model.layers.43.self_attn.indexer.wk", + "model.language_model.layers.43.self_attn.indexer.wq_b", + "model.language_model.layers.43.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.43.self_attn.kv_b_proj", + "model.language_model.layers.43.self_attn.o_proj", + "model.language_model.layers.43.self_attn.q_a_proj", + "model.language_model.layers.43.self_attn.q_b_proj", + "model.language_model.layers.44.mlp.shared_experts.down_proj", + "model.language_model.layers.44.mlp.shared_experts.gate_proj", + "model.language_model.layers.44.mlp.shared_experts.up_proj", + "model.language_model.layers.44.self_attn.b_proj", + "model.language_model.layers.44.self_attn.f_a_proj", + "model.language_model.layers.44.self_attn.f_b_proj", + "model.language_model.layers.44.self_attn.g_a_proj", + "model.language_model.layers.44.self_attn.g_b_proj", + "model.language_model.layers.44.self_attn.k_proj", + "model.language_model.layers.44.self_attn.o_proj", + "model.language_model.layers.44.self_attn.q_proj", + "model.language_model.layers.44.self_attn.v_proj", + "model.language_model.layers.45.mlp.shared_experts.down_proj", + "model.language_model.layers.45.mlp.shared_experts.gate_proj", + "model.language_model.layers.45.mlp.shared_experts.up_proj", + "model.language_model.layers.45.self_attn.indexer.weights_proj", + "model.language_model.layers.45.self_attn.indexer.wk", + "model.language_model.layers.45.self_attn.indexer.wq_b", + "model.language_model.layers.45.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.45.self_attn.kv_b_proj", + "model.language_model.layers.45.self_attn.o_proj", + "model.language_model.layers.45.self_attn.q_a_proj", + "model.language_model.layers.45.self_attn.q_b_proj", + "model.language_model.layers.5.mlp.shared_experts.down_proj", + "model.language_model.layers.5.mlp.shared_experts.gate_proj", + "model.language_model.layers.5.mlp.shared_experts.up_proj", + "model.language_model.layers.5.self_attn.b_proj", + "model.language_model.layers.5.self_attn.f_a_proj", + "model.language_model.layers.5.self_attn.f_b_proj", + "model.language_model.layers.5.self_attn.g_a_proj", + "model.language_model.layers.5.self_attn.g_b_proj", + "model.language_model.layers.5.self_attn.k_proj", + "model.language_model.layers.5.self_attn.o_proj", + "model.language_model.layers.5.self_attn.q_proj", + "model.language_model.layers.5.self_attn.v_proj", + "model.language_model.layers.6.mlp.shared_experts.down_proj", + "model.language_model.layers.6.mlp.shared_experts.gate_proj", + "model.language_model.layers.6.mlp.shared_experts.up_proj", + "model.language_model.layers.6.self_attn.b_proj", + "model.language_model.layers.6.self_attn.f_a_proj", + "model.language_model.layers.6.self_attn.f_b_proj", + "model.language_model.layers.6.self_attn.g_a_proj", + "model.language_model.layers.6.self_attn.g_b_proj", + "model.language_model.layers.6.self_attn.k_proj", + "model.language_model.layers.6.self_attn.o_proj", + "model.language_model.layers.6.self_attn.q_proj", + "model.language_model.layers.6.self_attn.v_proj", + "model.language_model.layers.7.mlp.shared_experts.down_proj", + "model.language_model.layers.7.mlp.shared_experts.gate_proj", + "model.language_model.layers.7.mlp.shared_experts.up_proj", + "model.language_model.layers.7.self_attn.indexer.weights_proj", + "model.language_model.layers.7.self_attn.indexer.wk", + "model.language_model.layers.7.self_attn.indexer.wq_b", + "model.language_model.layers.7.self_attn.kv_a_proj_with_mqa", + "model.language_model.layers.7.self_attn.kv_b_proj", + "model.language_model.layers.7.self_attn.o_proj", + "model.language_model.layers.7.self_attn.q_a_proj", + "model.language_model.layers.7.self_attn.q_b_proj", + "model.language_model.layers.8.mlp.shared_experts.down_proj", + "model.language_model.layers.8.mlp.shared_experts.gate_proj", + "model.language_model.layers.8.mlp.shared_experts.up_proj", + "model.language_model.layers.8.self_attn.b_proj", + "model.language_model.layers.8.self_attn.f_a_proj", + "model.language_model.layers.8.self_attn.f_b_proj", + "model.language_model.layers.8.self_attn.g_a_proj", + "model.language_model.layers.8.self_attn.g_b_proj", + "model.language_model.layers.8.self_attn.k_proj", + "model.language_model.layers.8.self_attn.o_proj", + "model.language_model.layers.8.self_attn.q_proj", + "model.language_model.layers.8.self_attn.v_proj", + "model.language_model.layers.9.mlp.shared_experts.down_proj", + "model.language_model.layers.9.mlp.shared_experts.gate_proj", + "model.language_model.layers.9.mlp.shared_experts.up_proj", + "model.language_model.layers.9.self_attn.b_proj", + "model.language_model.layers.9.self_attn.f_a_proj", + "model.language_model.layers.9.self_attn.f_b_proj", + "model.language_model.layers.9.self_attn.g_a_proj", + "model.language_model.layers.9.self_attn.g_b_proj", + "model.language_model.layers.9.self_attn.k_proj", + "model.language_model.layers.9.self_attn.o_proj", + "model.language_model.layers.9.self_attn.q_proj", + "model.language_model.layers.9.self_attn.v_proj" + ], + "nvfp4_mtp_targets": [ + "model.language_model.layers.45.mlp.experts" + ], + "mtp_activation_format": "BF16 (W4A16 reference group has no input_activations quantization)" +} diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 2f20bba9f695..a55bf4f6d076 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -8,7 +8,7 @@ from vllm.config.cache import CacheConfig from vllm.config.scheduler import SchedulerConfig -from vllm.config.speculative import SpeculativeConfig +from vllm.config.speculative import SpeculativeConfig, SpeculativeMethod from vllm.config.utils import get_hash_factors, hash_factors, normalize_value # Helpers @@ -237,7 +237,9 @@ def test_scheduler_config_hash_includes_max_num_seqs(): def test_speculative_config_hash_includes_graph_shape(): - def config(method: str, num_speculative_tokens: int) -> SpeculativeConfig: + def config( + method: SpeculativeMethod, num_speculative_tokens: int + ) -> SpeculativeConfig: return SpeculativeConfig( method=method, num_speculative_tokens=num_speculative_tokens, diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index c4e881ffd94e..9ab0f3b1139e 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -6,7 +6,7 @@ import sys import types from unittest.mock import Mock -from weakref import WeakKeyDictionary, ref +from weakref import ReferenceType, WeakKeyDictionary, ref import pytest import torch @@ -170,7 +170,7 @@ class _DeferredOnlineQuantAttention(_ReloadableAttentionLayer): def __init__(self): torch.nn.Module.__init__(self) - self.source_refs = [] + self.source_refs: list[ReferenceType[torch.Tensor]] = [] self.quant_method = _SourceLifetimeQuantMethod(self.source_refs) def tracking_weight_loader(param, loaded_weight): @@ -772,6 +772,30 @@ def __init__(self, quant_method): self.register_parameter("bias", bias) +@pytest.mark.parametrize("load_order", [("w13", "w2"), ("w2", "w13")]) +def test_online_processing_releases_replaced_full_precision_weights(load_order): + """Loader temporaries must not retain BF16 storage during kernel packing.""" + + class ReplacingMethod(_RecordingQuantMethod): + def process_weights_after_loading(self, layer): + references = [ref(layer.w13), ref(layer.w2)] + layer.w13 = torch.nn.Parameter(torch.empty(1), requires_grad=False) + layer.w2 = torch.nn.Parameter(torch.empty(1), requires_grad=False) + assert all(reference() is None for reference in references) + + layer = torch.nn.Module() + layer.quant_method = ReplacingMethod() + for name in ("w13", "w2"): + parameter = torch.nn.Parameter(torch.empty(4, 4, device="meta")) + parameter.weight_loader = default_weight_loader + layer.register_parameter(name, parameter) + initialize_online_processing(layer) + checkpoint_parameters = dict(layer.named_parameters()) + for name in load_order: + parameter = checkpoint_parameters[name] + parameter.weight_loader(parameter, torch.ones(4, 4)) + + def test_online_processing_waits_for_late_registered_bias(): # Regression test: `bias` is skipped by the meta device paths, but it is # still loaded by a weight loader. Excluding it from the processing trigger diff --git a/tests/model_executor/test_eagle_quantization.py b/tests/model_executor/test_eagle_quantization.py index a9622b3261d3..05738bb201bc 100644 --- a/tests/model_executor/test_eagle_quantization.py +++ b/tests/model_executor/test_eagle_quantization.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest import torch from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig -from vllm.model_executor.models.utils import get_draft_quant_config +from vllm.model_executor.models.utils import ( + get_draft_quant_config, + with_draft_quantization, +) from vllm.platforms import current_platform DEVICE_TYPE = current_platform.device_type @@ -54,6 +58,32 @@ def test_get_draft_quant_config_without_draft_model(): assert result is None +@pytest.mark.parametrize("quantized_draft", [False, True]) +def test_draft_quantization_preserves_target_runtime_and_unquantized_drafts( + quantized_draft, +): + """A draft must not inherit target quantization or rebuild target planning.""" + target_quantization = object() + draft_quantization = object() if quantized_draft else None + target = SimpleNamespace( + model_config=object(), + compilation_config=object(), + quant_config=target_quantization, + use_replayssm=True, + ) + with patch( + "vllm.model_executor.models.utils.get_draft_quant_config", + return_value=draft_quantization, + ): + draft = with_draft_quantization(target) + assert draft is not target + assert draft.quant_config is draft_quantization + assert target.quant_config is target_quantization + assert draft.model_config is target.model_config + assert draft.compilation_config is target.compilation_config + assert draft.use_replayssm + + @torch.inference_mode() @pytest.mark.parametrize("device", DEVICES) def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) -> None: diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 829e606ddfac..9fa9af1a0090 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -383,6 +383,7 @@ def test_glm53_selector_lazily_caches_fp32_head_projection() -> None: indexer = Glm5NextPooledIndexer.__new__(Glm5NextPooledIndexer) nn.Module.__init__(indexer) indexer.weights_proj = nn.Linear(hidden_size, 32, bias=False, dtype=torch.bfloat16) + indexer.weights_proj.quant_method = None indexer._weights_proj_fp32 = None hidden = torch.randn(3, hidden_size, dtype=torch.bfloat16) @@ -398,6 +399,32 @@ def test_glm53_selector_lazily_caches_fp32_head_projection() -> None: assert indexer._weights_proj_fp32.data_ptr() == pointer +def test_glm53_selector_dequantizes_mxfp8_weights_before_fp32_projection() -> None: + """Packed weights must include block scales; the source parameter is empty.""" + indexer = Glm5NextPooledIndexer.__new__(Glm5NextPooledIndexer) + nn.Module.__init__(indexer) + values = torch.arange(32 * 64).remainder(17).reshape(32, 64).float() + values = values.to(torch.float8_e4m3fn) + scales = torch.tensor([125, 130], dtype=torch.uint8).expand(32, 2).clone() + indexer.weights_proj = nn.Module() + indexer.weights_proj.weight = nn.Parameter(torch.empty(0), requires_grad=False) + indexer.weights_proj.b12x_mxfp8_packed_weight = SimpleNamespace( + out_features=32, + in_features=64, + weight=SimpleNamespace(values=values, scale_rows=scales.reshape(1, 32, 2)), + ) + indexer._weights_proj_fp32 = None + hidden = torch.randn(3, 64, dtype=torch.bfloat16) + dequantized = values.float() * torch.tensor([0.25, 8.0]).repeat_interleave(32) + expected = torch.nn.functional.linear(hidden.float(), dequantized) + actual = indexer._project_head_weights(hidden) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.dtype == torch.float32 + pointer = indexer._weights_proj_fp32.data_ptr() + indexer._project_head_weights(hidden) + assert indexer._weights_proj_fp32.data_ptr() == pointer + + def test_glm53_mla_spec_scales_fp8_index_tail_with_manager_block() -> None: spec = MLAAttentionSpec( block_size=256, diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index fbd5cbc56392..ea364a5cd259 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -124,11 +124,14 @@ def on_gfx942() -> bool: ) +@pytest.mark.parametrize("use_a16", [False, True]) def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( monkeypatch, + use_a16, ) -> None: method = object.__new__(Nvfp4OnlineMoEMethod) method.moe = SimpleNamespace(is_act_and_mul=True) + method.use_a16 = use_a16 method.nvfp4_backend = object() method.experts_cls = object method.moe_quant_config = None @@ -169,6 +172,14 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( make_kernel.assert_called_once() get_quant_config.assert_called_once() assert process_weights.call_count == 2 + assert convert_weights.call_args.kwargs["use_a16"] is use_a16 + assert make_kernel.call_args.kwargs["per_token_activation"] is not use_a16 + + +def test_online_nvfp4_a16_rejects_fp16_activations() -> None: + config = cast(Any, SimpleNamespace(in_dtype=torch.float16)) + with pytest.raises(ValueError, match="requires BF16 activations"): + Nvfp4OnlineMoEMethod(moe=config, use_a16=True) def _fully_quantized_quark_config() -> QuarkConfig: @@ -908,6 +919,65 @@ def test_online_quantization_records_global_config( } +@pytest.mark.parametrize( + "target", ["re:.*attn.*", "model.*.q_proj", "model.layers.[01].q_proj"] +) +def test_strict_online_targets_reject_patterns(target) -> None: + with pytest.raises(ValueError, match="exact names"): + QuantizationConfigArgs(targets={target: "mxfp8"}, strict_targets=True) + + +def test_strict_online_targets_cover_every_fused_projection( + default_vllm_config, + dist_init, +) -> None: + """A fused QKV counts all three source targets; unlisted layers stay BF16.""" + default_vllm_config.model_config = ModelConfig() + prefix = "model.layers.0.self_attn" + targets = {f"{prefix}.{p}_proj": "fp8_per_tensor" for p in ("q", "k", "v")} + config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets=targets, strict_targets=True) + ) + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + layer = ColumnParallelLinear(32, 32, bias=False, disable_tp=True) + with pytest.raises(ValueError, match="every strict target"): + config.validate_target_coverage() + layer.quant_method = config.get_quant_method(layer, f"{prefix}.qkv_proj") + config.validate_target_coverage() + assert set(config.quantized_targets) == set(targets) + assert isinstance( + config.get_quant_method(layer, f"{prefix}.o_proj"), UnquantizedLinearMethod + ) + assert set(config.quantized_layers) == {f"{prefix}.qkv_proj"} + model = torch.nn.Module() + model.add_module("qkv", layer) + config.validate_target_coverage(model) + with pytest.raises(ValueError, match="not present in the constructed model"): + config.validate_target_coverage(torch.nn.Module()) + head = torch.nn.Module() + head.quant_method = Fp8PerTensorOnlineLinearMethod() + model.add_module("head", head) + with pytest.raises(ValueError, match="outside strict targets"): + config.validate_target_coverage(model) + + +def test_strict_online_targets_reject_missing_fused_projection( + default_vllm_config, + dist_init, +) -> None: + default_vllm_config.model_config = ModelConfig() + config = OnlineQuantizationConfig( + QuantizationConfigArgs( + targets={"model.q_proj": "fp8_per_tensor"}, + strict_targets=True, + ) + ) + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + layer = ColumnParallelLinear(32, 32, bias=False, disable_tp=True) + with pytest.raises(ValueError, match="unmatched shards"): + config.get_quant_method(layer, "model.qkv_proj") + + def test_online_quantization_targets_ignore_collision() -> None: """A targets/ignore collision is reported when the layer is dispatched.""" config = OnlineQuantizationConfig( @@ -968,9 +1038,11 @@ def record_info(message: str, *args: object) -> None: log_online_quantization(default_vllm_config) assert logged_messages == [ - "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: " - "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: " - "re:.*qkv_proj.*, mxfp4)" + ( + "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: " + "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: " + "re:.*qkv_proj.*, mxfp4)" + ) ] @@ -1113,11 +1185,14 @@ def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None if ( not ( current_platform.is_cuda() - and current_platform.is_device_capability_family(100) + and ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) ) or current_platform.is_xpu() ): - pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.") + pytest.skip("NVFP4 weight quantization needs a Blackwell GPU.") elif not is_quant_method_supported("fp8"): pytest.skip("FP8 is not supported on this GPU type.") @@ -1174,9 +1249,13 @@ def quantize(w2_in, moe_tp_size): @pytest.mark.skipif( not ( - current_platform.is_cuda() and current_platform.is_device_capability_family(100) + current_platform.is_cuda() + and ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) ), - reason="NVFP4 weight quantization needs a Blackwell (SM100) GPU.", + reason="NVFP4 weight quantization needs a Blackwell GPU.", ) def test_online_nvfp4_quantizes_original_expert_weights() -> None: torch.manual_seed(0) @@ -1207,6 +1286,32 @@ def test_online_nvfp4_quantizes_original_expert_weights() -> None: ) +def test_online_moe_uses_caller_workspace() -> None: + """Online methods must preserve the runner's routed/shared scratch split.""" + method = Nvfp4OnlineMoEMethod.__new__(Nvfp4OnlineMoEMethod) + workspace = tuple(torch.empty(2, 4) for _ in range(3)) + x = torch.ones(2, 4) + + def apply_kernel(inputs, *args, workspace, **kwargs): + torch.add(inputs, 1, out=workspace[-1]) + return workspace[-1] + + method.moe_kernel = SimpleNamespace(is_monolithic=False, apply=apply_kernel) + layer = SimpleNamespace( + w13_weight=None, + w2_weight=None, + activation=None, + global_num_experts=2, + expert_map=None, + apply_router_weight_on_input=False, + ) + result = method.apply_with_workspace( + layer, x, torch.ones(2, 1), torch.zeros(2, 1), None, None, workspace + ) + assert result is workspace[-1] + torch.testing.assert_close(result, x + 1) + + @pytest.mark.skipif( not is_quant_method_supported("fp8"), reason="FP8 is not supported on this GPU type.", diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 0bf203ac4516..9f10d2232b0a 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -69,6 +69,24 @@ def test_args_moe_string_resolves_via_online_shorthand(): assert args.moe == QuantSpec(weight=kFp8Static128BlockSym) +def test_nvfp4_a16_spec_preserves_bfloat16_activation_selection(): + from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static + + args = QuantizationConfigArgs(targets={"model.layers.45.mlp.experts": "nvfp4_a16"}) + expected = QuantSpec(weight=kNvfp4Static, activation_dtype="bfloat16") + config = OnlineQuantizationConfig(args) + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.online.nvfp4 import ( + Nvfp4OnlineA16MoEMethod, + ) + + result = config.resolve_quant_method_cls( + Mock(spec=RoutedExperts), "model.layers.45.mlp.experts" + ) + assert result[-1] is Nvfp4OnlineA16MoEMethod + assert str(expected).endswith("/bfloat16") + + def test_args_string_shorthand_missing_slot_raises(): # int8_per_channel_weight_only sets only `moe`; using it on `linear` # has no defined spec and should raise rather than silently no-op. diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 7d32a1730837..15566ef81bc0 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Annotated, Any +from typing import Annotated, Any, Literal import regex as re from pydantic import ( @@ -83,6 +83,11 @@ class QuantSpec: activation: QuantKeyField = None """Activation quantization key, or a name from QUANT_KEY_NAMES.""" + activation_dtype: Literal["bfloat16"] | None = None + """Keep activations in this dtype instead of the method's quantized format. + Supported by online NVFP4 MoE with BF16 activations. + """ + def __str__(self) -> str: def quant_key_str(quant_key: QuantKey | None) -> str: if quant_key is None: @@ -96,7 +101,10 @@ def quant_key_str(quant_key: QuantKey | None) -> str: str(quant_key), ) - return quant_key_str(self.weight) + name = quant_key_str(self.weight) + return ( + name if self.activation_dtype is None else f"{name}/{self.activation_dtype}" + ) @config @@ -125,6 +133,12 @@ class QuantizationConfigArgs: Mutually exclusive with `linear` and `moe`. """ + strict_targets: bool = False + """Require exact target names and complete coverage at model construction. + Use separate target sets for the main and draft models. A missing target, + including a target outside the local pipeline partition, is an error. + """ + @field_validator("linear", "moe", mode="before") @classmethod def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: @@ -170,6 +184,14 @@ def _validate_targets(cls, v: Any) -> Any: @model_validator(mode="after") def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs": + if self.strict_targets: + if not self.targets: + raise ValueError("strict_targets requires a nonempty targets mapping") + if any( + name.startswith("re:") or any(char in name for char in "*?[") + for name in self.targets + ): + raise ValueError("strict_targets requires exact names without patterns") if self.targets is None: return self if self.linear is not None or self.moe is not None: @@ -216,6 +238,9 @@ def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs": "nvfp4_per_token": QuantizationConfigArgs( moe=QuantSpec(weight=kNvfp4Static), ), + "nvfp4_a16": QuantizationConfigArgs( + moe=QuantSpec(weight=kNvfp4Static, activation_dtype="bfloat16"), + ), } @@ -275,4 +300,5 @@ def resolve_quantization_config( moe=quantization_config.moe or base.moe, ignore=quantization_config.ignore or base.ignore, targets=quantization_config.targets or base.targets, + strict_targets=quantization_config.strict_targets, ) diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index a16a33eaaa33..4d3930872ea1 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -15,7 +15,6 @@ ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoEMethodBase, RoutedExperts, ) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( @@ -45,6 +44,7 @@ from vllm.model_executor.layers.quantization.online.int8 import ( Int8OnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.moe_base import OnlineMoEMethodBase from vllm.model_executor.layers.quantization.online.mxfp4 import ( Mxfp4OnlineLinearMethod, Mxfp4OnlineMoEMethod, @@ -54,6 +54,7 @@ Mxfp8OnlineMoEMethod, ) from vllm.model_executor.layers.quantization.online.nvfp4 import ( + Nvfp4OnlineA16MoEMethod, Nvfp4OnlineMoEMethod, ) from vllm.model_executor.layers.quantization.utils.config_utils import ( @@ -155,6 +156,40 @@ def __init__( self.args = args self.ignored_layers: list[str] = args.ignore self.quantized_layers: dict[str, tuple[str, str, str | None]] = {} + self.quantized_targets: dict[str, str] = {} + + def validate_target_coverage(self, model: torch.nn.Module | None = None) -> None: + """Reject incomplete exact target sets before checkpoint loading.""" + if not self.args.strict_targets: + return + missing = set(self.args.targets or ()) - self.quantized_targets.keys() + if missing: + raise ValueError( + "Online quantization did not resolve every strict target: " + + ", ".join(sorted(missing)) + ) + if model is not None: + present = set() + for name, layer in model.named_modules(): + method = getattr(layer, "quant_method", None) + if isinstance(method, (OnlineLinearBase, OnlineMoEMethodBase)): + prefix = getattr(method, "_online_target_prefix", None) + if prefix not in self.quantized_layers: + raise ValueError( + f"Online quantizer at {name} is outside strict targets; " + "disable independent head or model quantization flags" + ) + present.add(prefix) + absent = { + target + for target, prefix in self.quantized_targets.items() + if prefix not in present + } + if absent: + raise ValueError( + "Strict online targets are not present in the constructed model: " + + ", ".join(sorted(absent)) + ) @property def quantized_layer_summaries(self) -> list[str]: @@ -230,6 +265,17 @@ def _get_method_cls( """ if spec is None or spec.weight is None: return None + if spec.activation_dtype is not None: + if ( + table is not _ONLINE_MOE_METHODS + or spec.weight != kNvfp4Static + or spec.activation is not None + ): + raise ValueError( + "activation_dtype is supported only for online NVFP4 MoE " + "without an activation quantization override" + ) + return Nvfp4OnlineA16MoEMethod cls = table.get(spec.weight) if cls is None: raise ValueError( @@ -374,17 +420,38 @@ def get_quant_method( resolved = self.resolve_quant_method_cls(layer, prefix) if resolved is not None: source, quant_key_str, target_pattern, _, quant_method_cls = resolved + if self.args.strict_targets: + assert self.args.targets is not None + matches = find_matching_patterns( + prefix, + self.args.targets, + self.packed_modules_mapping, + use_fnmatch=False, + ) + for shard_matches in matches: + for target in shard_matches: + previous = self.quantized_targets.setdefault(target, prefix) + if previous != prefix: + raise ValueError( + f"Strict target {target} resolved to both " + f"{previous} and {prefix}" + ) self.quantized_layers[prefix] = ( source.value, quant_key_str, target_pattern, ) if isinstance(layer, RoutedExperts): - assert issubclass(quant_method_cls, FusedMoEMethodBase) - return quant_method_cls(moe=layer.moe_config) - - assert issubclass(quant_method_cls, OnlineLinearBase) - return quant_method_cls() + assert issubclass(quant_method_cls, OnlineMoEMethodBase) + method: OnlineMoEMethodBase | OnlineLinearBase = quant_method_cls( + moe=layer.moe_config + ) + else: + assert issubclass(quant_method_cls, OnlineLinearBase) + method = quant_method_cls() + if self.args.strict_targets: + method._online_target_prefix = prefix + return method if isinstance(layer, LinearBase): return UnquantizedLinearMethod() diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 0234702c8bab..20eb47b76217 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -117,6 +117,7 @@ class OnlineLinearBase(LinearMethodBase): weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True + _online_target_prefix: str | None = None def __init__(self): self.out_dtype = torch.get_default_dtype() diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 866c88b3d649..9822d035c792 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -23,6 +23,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): """ uses_meta_device: bool = True + _online_target_prefix: str | None = None def create_weights( self, @@ -166,6 +167,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, + workspace: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None @@ -181,4 +183,5 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, shared_experts=shared_experts, shared_experts_input=shared_experts_input, + workspace=workspace, ) diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index e09acee29fb1..6c97a3791403 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -4,7 +4,7 @@ import torch from torch.nn import Module -from vllm._custom_ops import scaled_fp4_quant +from vllm._custom_ops import create_fp4_output_tensors, scaled_fp4_quant from vllm.model_executor.kernels.linear.nvfp4.b12x import ( B12xNvFp4LinearKernel, run_b12x_nvfp4_serialized_linear, @@ -127,56 +127,61 @@ def _quantize_moe_weight_to_nvfp4( # Keep the original BF16/FP16 values as the quantizer input. Folding each # expert's FP32 global scale into the weight would add a BF16/FP16 rounding # before the group-16 scale and E2M1 values are selected. - weight = weight.contiguous() - quantized_experts = [ - scaled_fp4_quant( - expert_weight, - expert_scale, - is_sf_swizzled_layout=False, - ) - for expert_weight, expert_scale in zip( - weight, - global_scale, - strict=True, + num_experts, rows, _ = weight.shape + qweight, block_scale = create_fp4_output_tensors( + num_experts * rows, k, weight.device, is_sf_swizzled_layout=False + ) + qweight = qweight.view(num_experts, rows, k // 2) + block_scale = block_scale.view(num_experts, rows, k // 16) + # Write directly into the destination to avoid a second full packed tensor + # while the BF16 expert weights are still resident during loading. + for index in range(num_experts): + torch.ops._C.scaled_fp4_quant.out( + weight[index].contiguous(), + global_scale[index], + False, + output=qweight[index], + output_scale=block_scale[index], ) - ] - qweight = torch.stack([quantized for quantized, _ in quantized_experts]) - block_scale = torch.stack([block_scale for _, block_scale in quantized_experts]) return ( qweight, - block_scale, + block_scale.view(torch.float8_e4m3fn), weight_scale_2, ) class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): - """Online NVFP4 MoE quantization with per-token activation scales. + """Online NVFP4 MoE with per-token FP4 or unquantized BF16 activations. - Quantizes fp16/bf16 expert weights to NVFP4 at load time; the FlashInfer - TRTLLM kernel computes per-token activation scales at runtime. Blackwell - (SM100) only. + Per-token FP4 activations require Blackwell SM100. The BF16 variant selects + a W4A16 backend through the MoE configuration. """ def __init__( self, *, moe: FusedMoEConfig, + use_a16: bool = False, ): - if not current_platform.is_device_capability_family(100): + if use_a16 and moe.in_dtype != torch.bfloat16: + raise ValueError("nvfp4_a16 online MoE requires BF16 activations") + if not use_a16 and not current_platform.is_device_capability_family(100): raise ValueError( "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU." ) super().__init__(moe) + self.use_a16 = use_a16 self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=kNvfp4Dynamic, + activation_key=None if use_a16 else kNvfp4Dynamic, ) def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return + self._zero_padding(layer) self._quantize_weights(layer) self._setup_kernel(layer) @@ -187,19 +192,18 @@ def _quantize_weights(self, layer: Module) -> None: w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4( layer.w13_weight, moe_tp_size ) - w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4( - layer.w2_weight, moe_tp_size - ) - replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) + w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4( + layer.w2_weight, moe_tp_size + ) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w2_weight_scale", w2_scale) replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) - # Neutral (1.0) activation global scales: the kernel derives per-token - # scales at runtime, so the output scalars reduce to the weight scales. + # Neutral activation scales: the FP4 path derives per-token scales at + # runtime, while W4A16 consumes BF16 values without activation scaling. ones = torch.ones(layer.num_experts, device=w13.device, dtype=torch.float32) replace_parameter(layer, "w13_input_scale", ones) replace_parameter(layer, "w2_input_scale", ones.clone()) @@ -226,6 +230,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: w2_scale_2=layer.w2_weight_scale_2, a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, + use_a16=self.use_a16, ) replace_parameter(layer, "w13_weight", w13) @@ -246,7 +251,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), - per_token_activation=True, + per_token_activation=not self.use_a16, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -262,4 +267,12 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), layer=layer, + use_a16=self.use_a16, ) + + +class Nvfp4OnlineA16MoEMethod(Nvfp4OnlineMoEMethod): + """Quantize BF16 expert weights at load time and retain BF16 activations.""" + + def __init__(self, *, moe: FusedMoEConfig): + super().__init__(moe=moe, use_a16=True) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 6f356464585a..5b4a3d961801 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -48,7 +48,7 @@ def create_model( model_config=model_config, prefix=prefix, ) - log_online_quantization(vllm_config) + log_online_quantization(vllm_config, model) log_model_inspection(model) return model @@ -103,7 +103,9 @@ def log_model_inspection(model: nn.Module) -> None: logger.info("vLLM model structure:\n%s", format_model_inspection(model)) -def log_online_quantization(vllm_config: VllmConfig) -> None: +def log_online_quantization( + vllm_config: VllmConfig, model: nn.Module | None = None +) -> None: """Log the online-quantized layer count and types, when applicable.""" from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, @@ -118,6 +120,8 @@ def log_online_quantization(vllm_config: VllmConfig) -> None: if not isinstance(quant_config, OnlineQuantizationConfig): return + quant_config.validate_target_coverage(model) + logger.info( "Quantized %d layers of types: %s", len(quant_config.quantized_layers), diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 0aa1dc406ff6..e7be7a726942 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -309,6 +309,9 @@ def online_process_loader(*args, **kwargs): # Process and copy when all weights are loaded if info.load_numel >= info.load_numel_total: # type: ignore[operator] + # Quantizers replace BF16 parameters before packing kernels. Do not + # retain the final direct-load destination through that conversion. + del bound_args _layerwise_process(layer, info) LOADING_LAYERS.discard(layer) @@ -456,6 +459,9 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): # formerly held by checkpoint sources for its temporary workspace. info.loaded_weights.clear() loaded_args = None + # Loop locals must not keep a replaced full-precision parameter resident + # while its quantized replacement is prepared for the execution backend. + param = None _zero_online_processing_unloaded(layer) # Process weights (quantization, repacking, etc.) quant_method = getattr(layer, "quant_method", None) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 6cd21763a095..e8646e0763c8 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -971,6 +971,21 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non ) +def with_draft_quantization(vllm_config: VllmConfig) -> VllmConfig: + """Select draft weights without revalidating target execution planning. + + Runtime options such as ReplaySSM belong to the target architecture. A + shallow copy preserves that planning and shared compilation state. Assigning + quant_config after construction also preserves an unquantized draft's None + instead of reinitializing it from the target checkpoint in __post_init__. + """ + import copy + + draft_config = copy.copy(vllm_config) + draft_config.quant_config = get_draft_quant_config(vllm_config) + return draft_config + + def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int: """ Extract the layer index from the module name. diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 78859c0f2556..71264e19c73f 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -18,6 +18,9 @@ from vllm.model_executor.layers.layernorm import LayerNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) from vllm.models.deepseek_v4.nvidia.b12x_indexer import ( B12xC4SparseIndexer, ) @@ -236,7 +239,7 @@ def __init__( hidden_size, _INDEX_HEADS, bias=False, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.weights_proj", ) self.k_norm = LayerNorm(_INDEX_HEAD_DIM, eps=1e-6) @@ -490,7 +493,11 @@ def get_b12x_physical_selection( def _project_head_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: if self._weights_proj_fp32 is None: - self._weights_proj_fp32 = self.weights_proj.weight.detach().float() + # Honor quantized checkpoint and online weights while retaining the + # selector's FP32 projection arithmetic. + self._weights_proj_fp32 = get_and_maybe_dequant_weights( + self.weights_proj, out_dtype=torch.float32 + ).detach() return F.linear(hidden_states.float(), self._weights_proj_fp32) @staticmethod diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 1088311e234e..a65e109c4a62 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1325,8 +1325,9 @@ def _get_model(self) -> nn.Module: need to customize model loading. """ from vllm.compilation.backends import set_model_tag + from vllm.model_executor.models.utils import with_draft_quantization - draft_vllm_config = self._create_draft_vllm_config() + draft_vllm_config = with_draft_quantization(self._create_draft_vllm_config()) with set_model_tag("eagle_head"): model = get_model( vllm_config=draft_vllm_config, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 89ed39d14f2b..135975af8236 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -7,7 +7,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.model_loader import get_model -from vllm.model_executor.models.utils import PPMissingLayer +from vllm.model_executor.models.utils import PPMissingLayer, with_draft_quantization from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -138,6 +138,7 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod draft_load_config = get_pp_safe_draft_load_config(vllm_config.load_config) if draft_load_config is not vllm_config.load_config: vllm_config = replace(vllm_config, load_config=draft_load_config) + vllm_config = with_draft_quantization(vllm_config) with set_model_tag("eagle_head"): eagle_model = get_model( vllm_config=vllm_config, model_config=draft_model_config From ae168049b57247eef1a753496e192c060fffc0f3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 3 Oct 2026 16:57:00 +0000 Subject: [PATCH 344/365] Link GLM online weight release metadata to PR 965 Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- .lil/changes/vllm-glm-spark-online-weights.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-glm-spark-online-weights.json b/.lil/changes/vllm-glm-spark-online-weights.json index 788d782004d5..0a18e3446528 100644 --- a/.lil/changes/vllm-glm-spark-online-weights.json +++ b/.lil/changes/vllm-glm-spark-online-weights.json @@ -15,6 +15,8 @@ "authors": [ "voipmonitor" ], - "pull_requests": [], + "pull_requests": [ + 965 + ], "requires": [] } From 64303cb5323c8fbf616c2515dc36351fa08a0512 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 02:15:31 +0000 Subject: [PATCH 345/365] 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 d56ce3013728..60201b734bed 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -3,6 +3,7 @@ """b12x modular tensor-parallel fused MoE backend.""" import functools +import os import weakref from collections.abc import Sequence from dataclasses import dataclass @@ -229,6 +230,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): @@ -423,6 +472,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 @@ -481,7 +560,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 @@ -964,6 +1048,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 9e60da6cddfe78c818aa15413bad97e124421ad2 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 02:20:55 +0000 Subject: [PATCH 346/365] 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 1b6c2a3bde1f0da79731930237e385a5745b68e5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:06:27 +0000 Subject: [PATCH 347/365] release: the vLLM A4 prefill change requires the b12x feature fragment Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-b12x-w4a16-a4-prefill.json | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json index 4a82b4504c18..90e0c570a1a6 100644 --- a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json +++ b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json @@ -3,14 +3,20 @@ "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).", "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"], + "authors": [ + "voipmonitor" + ], "pull_requests": [], - "requires": [] + "requires": [ + "b12x-w4a16-a4-prefill" + ] } From d2e74c036de3ca37799ddf181b894ce306416428 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:20:28 +0000 Subject: [PATCH 348/365] release: link the A4 prefill fragment to #970 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-b12x-w4a16-a4-prefill.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json index 90e0c570a1a6..09962b3906a0 100644 --- a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json +++ b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json @@ -15,7 +15,9 @@ "authors": [ "voipmonitor" ], - "pull_requests": [], + "pull_requests": [ + 970 + ], "requires": [ "b12x-w4a16-a4-prefill" ] From 8fa22e45602290e39b64a1564482c4e4588eb317 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:51:39 +0000 Subject: [PATCH 349/365] 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 60201b734bed..490a9590d689 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 functools import os import weakref @@ -271,10 +272,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 @@ -483,9 +482,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( @@ -1059,7 +1056,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 9e4d0f6beac15ea60105525e005dd044e3c9d23e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:52:08 +0000 Subject: [PATCH 350/365] release: A4 prefill fragment with the measured GLM numbers Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-b12x-w4a16-a4-prefill.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json index 09962b3906a0..b623221aa416 100644 --- a/.lil/changes/vllm-b12x-w4a16-a4-prefill.json +++ b/.lil/changes/vllm-b12x-w4a16-a4-prefill.json @@ -8,7 +8,7 @@ ], "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." ], From a628c6ad0d36e525ea359f30c6c5e15032388528 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 02:31:17 +0000 Subject: [PATCH 351/365] Serve lossless MXFP4-CSF DeepSeek-V4-Flash and Vision-Exp checkpoints Add the deepseek_v4_flash MXFP4-CSF family (43 target layers of 256 routed experts, hidden 4096, MoE intermediate 2048, DeepSeek-native tensor names, TP1/2/4/8) and dispatch it to DeepseekV4Mxfp4CsfConfig. The config derives from the class CUDA registers as deepseek_v4_fp8 and parses the retained source block-FP8 fields, so attention, dense, shared-expert, vision and MTP/DSpark draft weights keep their native methods. Target routed experts read compressed scales with MXFP8 activations unless VLLM_B12X_MOE_FP4_FORCE_A16=1 requests BF16. Expert parallelism is rejected: MegaMoE experts bypass quantization methods and would never read the compressed experts. The loader family table records each family's compressed layer range and rejects reads outside it. The FP4-CSF guide documents the DeepSeek MXFP8 activation default and the DeepSeek-V4-Flash serving contract. Validation: tests/quantization/test_mxfp4_csf.py and test_nvfp4_csf.py pass on CPU (74 tests, 13 new); ruff, mypy hook, typos pass. No GPU serving run yet. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-mxfp4-csf-ds4-flash.json | 22 ++ docs/features/quantization/fp4_csf.md | 44 ++- tests/quantization/test_mxfp4_csf.py | 312 ++++++++++++++++++ vllm/config/load.py | 2 +- .../layers/quantization/mxfp4_csf.py | 7 +- .../model_loader/mxfp4_csf_loader.py | 46 ++- vllm/models/deepseek_v4/mxfp4_csf.py | 105 ++++++ vllm/models/deepseek_v4_1/mxfp4_csf.py | 2 +- 8 files changed, 515 insertions(+), 25 deletions(-) create mode 100644 .lil/changes/vllm-mxfp4-csf-ds4-flash.json create mode 100644 vllm/models/deepseek_v4/mxfp4_csf.py diff --git a/.lil/changes/vllm-mxfp4-csf-ds4-flash.json b/.lil/changes/vllm-mxfp4-csf-ds4-flash.json new file mode 100644 index 000000000000..432453d914cb --- /dev/null +++ b/.lil/changes/vllm-mxfp4-csf-ds4-flash.json @@ -0,0 +1,22 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-mxfp4-csf-ds4-flash", + "category": "feature", + "summary": "Load lossless MXFP4-CSF checkpoints of DeepSeek-V4-Flash and DeepSeek-V4-Flash-Vision-Exp", + "models": [ + "DeepSeek-V4-Flash", + "DeepSeek-V4-Flash-Vision-Exp" + ], + "compatibility": "Opt-in: serve a deepseek_v4_flash MXFP4-CSF serving directory with --quantization mxfp4_csf --load-format mxfp4_csf. Original checkpoints and their launch flags are unchanged.", + "details": [ + "The deepseek_v4_flash family reads the compressed scales of 256 routed MXFP4 experts in each of the 43 target layers by DeepSeek's native tensor names, at TP1, TP2, TP4 or TP8.", + "Routed experts use MXFP8 activations like the native B12X MXFP4 path unless VLLM_B12X_MOE_FP4_FORCE_A16=1 requests BF16.", + "Attention, dense and shared-expert weights keep the source block-FP8 configuration. MTP/DSpark draft experts, the vision tower, the aligner and the image embeddings load unchanged.", + "Expert-parallel execution, including MegaMoE, is rejected because those experts would bypass the compressed-scale reader." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [], + "requires": [] +} diff --git a/docs/features/quantization/fp4_csf.md b/docs/features/quantization/fp4_csf.md index b7449ca4708f..75ccc6241da7 100644 --- a/docs/features/quantization/fp4_csf.md +++ b/docs/features/quantization/fp4_csf.md @@ -5,16 +5,19 @@ original scale bytes, calibration values and other tensors are preserved. The B12X backend keeps scales compressed between layers and reconstructs the needed scales into shared GPU scratch before expert computation. -Use `--quantization mxfp4_csf --load-format mxfp4_csf` for supported Kimi-K3 -or DeepSeek-V4.1-Flash MXFP4 containers. Use -`--quantization nvfp4_csf --load-format nvfp4_csf` for GLM-5.3-Flash or -Qwen3.8-Flash-Next NVFP4 containers. Automatic loading does not select these readers; -set the load format explicitly. +Use `--quantization mxfp4_csf --load-format mxfp4_csf` for supported Kimi-K3, +DeepSeek-V4.1-Flash, DeepSeek-V4-Flash or DeepSeek-V4-Flash-Vision-Exp MXFP4 +containers. Use `--quantization nvfp4_csf --load-format nvfp4_csf` for +GLM-5.3-Flash or Qwen3.8-Flash-Next NVFP4 containers. Automatic loading does +not select these readers; set the load format explicitly. MXFP4 uses a one-bit unsigned offset from a row's base byte. NVFP4 uses a four-bit offset. Both preserve out-of-interval scale bytes as exceptions. -They retain their source arithmetic: MXFP4-CSF uses BF16 expert activations; -NVFP4-CSF uses native calibrated FP4 expert activations. +Expert activations follow each model's native B12X policy. DeepSeek-V4.1-Flash +and DeepSeek-V4-Flash MXFP4-CSF quantize expert activations to MXFP8 (W4A8) +unless `VLLM_B12X_MOE_FP4_FORCE_A16=1` requests BF16 activations. Kimi-K3 +MXFP4-CSF uses BF16 expert activations. NVFP4-CSF uses native calibrated FP4 +expert activations. ## GLM-5.3-Flash on two 96 GiB Blackwell GPUs @@ -76,6 +79,33 @@ The example reserves 2 GiB/GPU for KV cache. Qwen's intermediate dimension supports TP1 and TP2 in this reader; TP4 does not satisfy the 64-channel local alignment requirement. +## DeepSeek-V4-Flash on two 96 GiB Blackwell GPUs + +This also applies to DeepSeek-V4-Flash-Vision-Exp. The serving directory +contains the source tokenizer metadata and a `config.json` whose +`quantization_config` keeps the source block-FP8 fields, including +`weight_block_size` equal to `[128, 128]`, and sets `quant_method` to +`mxfp4_csf`, `format_version` to 1 and an absolute `checkpoint_root`. + +The `deepseek_v4_flash` container compresses the routed-expert scales of the +43 target layers. Attention, dense and shared-expert weights retain block +FP8. MTP/DSpark draft layers (`mtp.*`), the vision tower, the aligner and the +image embeddings retain their source tensors and precision. A `config.json` +with a vision tower selects the multimodal model. + +```bash +vllm serve /models/DeepSeek-V4-Flash-MXFP4-CSF/serve \ + --quantization mxfp4_csf --load-format mxfp4_csf \ + --tensor-parallel-size 2 --moe-backend b12x --attention-backend B12X \ + --block-size 256 --max-num-seqs 8 --max-num-batched-tokens 4096 \ + --speculative-config '{"method":"dspark","num_speculative_tokens":5}' \ + --tokenizer-mode deepseek_v4 --reasoning-parser deepseek_v4 \ + --tool-call-parser deepseek_v4 --host 0.0.0.0 --port 8000 +``` + +The reader supports TP1, TP2, TP4 and TP8 with pipeline parallel size 1, +without expert/data parallelism or ubatching. + ## Compatibility Install matching vLLM and B12X CSF support. vLLM owns the checkpoint manifest, diff --git a/tests/quantization/test_mxfp4_csf.py b/tests/quantization/test_mxfp4_csf.py index dedcc69c9b6a..2f2c5f1691b5 100644 --- a/tests/quantization/test_mxfp4_csf.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -618,3 +618,315 @@ def test_expert_inventory_must_match_declared_count(count): w13_scale_scratch=None, w2_scale_scratch=None, ) + + +# DeepSeek-V4-Flash (0731 and Vision-Exp) serving config: the source block-FP8 +# fields plus the MXFP4-CSF method and container location. +DS4_FLASH_SERVING = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "mxfp4_csf", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + "format_version": 1, + "checkpoint_root": "/csf", +} + + +@pytest.fixture +def ds4_flash_moe(): + return FusedMoEConfig( + num_experts=256, + num_local_experts=256, + num_logical_experts=256, + experts_per_token=6, + hidden_dim=4096, + intermediate_size=1024, + in_dtype=torch.bfloat16, + device="cpu", + activation=MoEActivation.SILU, + swiglu_limit=10.0, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + ) + + +def ds4_flash_vllm_config( + *, pp=1, ubatching=False, load_format="mxfp4_csf", ep=False, expert_dtype="fp4" +): + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(num_hidden_layers=43, expert_dtype=expert_dtype) + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=pp, + use_ubatching=ubatching, + enable_expert_parallel=ep, + ), + load_config=SimpleNamespace(load_format=load_format), + ) + + +def test_ds4_flash_reader_requests_native_target_expert_names(monkeypatch): + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint + + requested = [] + + class Reader: + def __init__(self, root, source_names, codec): + assert codec == "mxfp4" + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def matrix(self, weight, scale): + requested.append((weight, scale)) + return weight + + monkeypatch.setattr( + checkpoint, + "checkpoint_contract", + lambda _: {"family": "deepseek_v4_flash", "source_names": {}}, + ) + monkeypatch.setattr(checkpoint, "CsfTensorReader", Reader) + monkeypatch.setattr( + checkpoint, "_load_mxfp4_csf_weights", lambda experts, **_: list(experts) + ) + experts = checkpoint.read_mxfp4_csf_layer( + "/csf", + 42, + num_experts=256, + hidden_size=4096, + intermediate_size=2048, + tp_rank=1, + tp_size=2, + device="cpu", + w13_scale_scratch=None, + w2_scale_scratch=None, + ) + assert len(experts) == len(requested) // 3 == 256 + prefix = "layers.42.ffn.experts.255" + assert requested[-3:] == [ + (f"{prefix}.{p}.weight", f"{prefix}.{p}.scale") for p in ("w1", "w3", "w2") + ] + assert experts[-1] == tuple(weight for weight, _ in requested[-3:]) + + +@pytest.mark.parametrize( + "layer,geometry,tp_size,error", + [ + (43, (256, 4096, 2048), 2, "outside the compressed expert inventory"), + (0, (384, 5120, 2304), 2, "geometry"), + (0, (256, 4096, 2048), 16, "supports TP"), + ], +) +def test_ds4_flash_reader_rejects_draft_layers_geometry_and_tp( + monkeypatch, layer, geometry, tp_size, error +): + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint + + monkeypatch.setattr( + checkpoint, + "checkpoint_contract", + lambda _: {"family": "deepseek_v4_flash", "source_names": {}}, + ) + num_experts, hidden_size, intermediate_size = geometry + with pytest.raises(ValueError, match=error): + checkpoint.read_mxfp4_csf_layer( + "/csf", + layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_rank=0, + tp_size=tp_size, + device="cpu", + w13_scale_scratch=None, + w2_scale_scratch=None, + ) + + +def test_mxfp4_csf_dispatches_ds4_flash_with_source_dense_config(monkeypatch): + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint + from vllm.models.deepseek_v4.mxfp4_csf import DeepseekV4Mxfp4CsfConfig + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config + + monkeypatch.setattr( + checkpoint, "checkpoint_contract", lambda _: {"family": "deepseek_v4_flash"} + ) + owner = get_quantization_config("mxfp4_csf").from_config(DS4_FLASH_SERVING) + source = { + key: value + for key, value in DS4_FLASH_SERVING.items() + if key not in ("format_version", "checkpoint_root") + } + native = DeepseekV4FP8Config.from_config({**source, "quant_method": "fp8"}) + assert isinstance(owner, DeepseekV4Mxfp4CsfConfig) + assert owner.get_name() == "mxfp4_csf" and owner.get_min_capability() == 120 + assert owner.checkpoint_root == "/csf" and owner.scale_scratch is None + for field in ( + "is_checkpoint_fp8_serialized", + "activation_scheme", + "ignored_layers", + "weight_block_size", + "is_scale_e8m0", + ): + assert getattr(owner, field) == getattr(native, field) + # The serving config must select the CSF reader, not the native FP8 path. + hf_config = SimpleNamespace(model_type="deepseek_v4") + assert ( + DeepseekV4FP8Config.override_quantization_method( + DS4_FLASH_SERVING, "mxfp4_csf", hf_config=hf_config + ) + is None + ) + for method in (Mxfp4CsfConfig, DeepseekV4Mxfp4CsfConfig): + assert ( + method.override_quantization_method( + DS4_FLASH_SERVING, "mxfp4_csf", hf_config=hf_config + ) + == "mxfp4_csf" + ) + with pytest.raises(ValueError, match=r"\[128, 128\]"): + DeepseekV4Mxfp4CsfConfig.from_config( + {**DS4_FLASH_SERVING, "weight_block_size": [32, 32]} + ) + with pytest.raises(ValueError, match="format_version"): + DeepseekV4Mxfp4CsfConfig.from_config({**DS4_FLASH_SERVING, "format_version": 2}) + + +@pytest.mark.parametrize("force_a16", [False, True]) +def test_ds4_flash_target_experts_use_csf_while_dense_and_drafts_stay_native( + monkeypatch, ds4_flash_moe, force_a16 +): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.models.deepseek_v4 import mxfp4_csf + + owner = mxfp4_csf.DeepseekV4Mxfp4CsfConfig.from_config(DS4_FLASH_SERVING) + config = ds4_flash_vllm_config() + monkeypatch.setattr(mxfp4_csf.envs, "VLLM_B12X_MOE_FP4_FORCE_A16", force_a16) + monkeypatch.setattr(mxfp4_csf, "get_current_vllm_config", lambda: config) + monkeypatch.setattr(mxfp4_csf, "get_current_vllm_config_or_none", lambda: config) + native = object() + monkeypatch.setattr( + mxfp4_csf.DeepseekV4FP8Config, "get_quant_method", lambda *a, **k: native + ) + layer = Mock(spec=RoutedExperts) + layer.moe_config = ds4_flash_moe + for prefix in ( + "model.layers.0.ffn.experts", + "model.layers.42.ffn.experts", + "language_model.model.layers.42.ffn.experts", + ): + method = owner.get_quant_method(layer, prefix) + assert isinstance(method, Mxfp4CsfMoEMethod) and method.owner is owner + assert method.activation_mode == ("a16" if force_a16 else "a8") + assert method.get_fused_moe_quant_config(layer).quant_dtype == ( + None if force_a16 else "mxfp8" + ) + # MTP and DSpark draft layers are numbered after the 43 target layers. + for prefix in ("model.layers.43.ffn.experts", "model.layers.45.ffn.experts"): + assert owner.get_quant_method(layer, prefix) is native + assert owner.get_quant_method(torch.nn.Module(), "model.layers.0.attn.wq_a") is ( + native + ) + + +@pytest.mark.parametrize( + "mode,error,match", + [ + ("pipeline", NotImplementedError, "PP1 without ubatching"), + ("ubatching", NotImplementedError, "PP1 without ubatching"), + ("wrong_loader", ValueError, "load-format"), + ("fp8_experts", ValueError, "MXFP4 routed experts"), + ("expert_parallel", NotImplementedError, "expert parallelism"), + ], +) +def test_ds4_flash_csf_rejects_unsupported_execution( + monkeypatch, ds4_flash_moe, mode, error, match +): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.models import deepseek_v41 + from vllm.models.deepseek_v4 import mxfp4_csf + + owner = mxfp4_csf.DeepseekV4Mxfp4CsfConfig.from_config(DS4_FLASH_SERVING) + config = ds4_flash_vllm_config( + pp=2 if mode == "pipeline" else 1, + ubatching=mode == "ubatching", + load_format="safetensors" if mode == "wrong_loader" else "mxfp4_csf", + ep=mode == "expert_parallel", + expert_dtype="fp8" if mode == "fp8_experts" else "fp4", + ) + for module in (mxfp4_csf, deepseek_v41.quant_config): + monkeypatch.setattr(module, "get_current_vllm_config", lambda: config) + monkeypatch.setattr(mxfp4_csf, "get_current_vllm_config_or_none", lambda: config) + layer = Mock(spec=RoutedExperts) + layer.moe_config = ds4_flash_moe + with pytest.raises(error, match=match): + owner.get_quant_method(layer, "model.layers.42.ffn.experts") + if mode == "expert_parallel": + # MegaMoE experts never reach a quantization method. + with pytest.raises(error, match=match): + owner.get_quant_method(torch.nn.Module(), "model.layers.0.attn.wq_a") + + +def test_ds4_flash_loader_keeps_vision_draft_and_dense_tensors(tmp_path): + from vllm.model_executor.model_loader.mxfp4_csf_loader import CODEC, SCHEMA + + tensor_dir = tmp_path / "tensors" + tensor_dir.mkdir() + name = "model-00001-of-00001.safetensors" + + def byte(): + return torch.ones(8, dtype=torch.uint8) + + compressed = { + "layers.0.ffn.experts.0.w1.weight": byte(), + "layers.0.ffn.experts.0.w1.scale.mxfp4_csf_fixed": byte(), + "layers.0.ffn.experts.0.w1.scale.mxfp4_csf_exceptions": torch.zeros( + 1, dtype=torch.uint32 + ), + "layers.42.ffn.experts.255.w2.weight": byte(), + } + draft = { + "mtp.0.ffn.experts.0.w1.weight": torch.full((8,), 42, dtype=torch.uint8), + "mtp.0.ffn.experts.0.w1.scale": byte(), + "mtp.2.markov_head.markov_w1.weight": torch.ones(8, dtype=torch.bfloat16), + } + retained = { + **draft, + "layers.0.attn.wq_a.scale": byte(), + "layers.0.ffn.gate.bias_vl": torch.ones(8), + "layers.0.ffn.gate.tid2eid": torch.ones(8, dtype=torch.int32), + "layers.0.hc_attn_fn": torch.ones(8), + "vision.blocks.0.attn.wqkv.weight": torch.ones(8, dtype=torch.bfloat16), + "aligner.w1.weight": torch.ones(8, dtype=torch.bfloat16), + "image_start": torch.ones(8), + } + tensors = {**compressed, **retained} + save_file(tensors, tensor_dir / name) + common = {"schema": SCHEMA, "codec": CODEC, "family": "deepseek_v4_flash"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": name}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: name for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace( + num_hidden_layers=43, + quantization_config={**DS4_FLASH_SERVING, "checkpoint_root": str(tmp_path)}, + ) + ) + loader = Mxfp4CsfModelLoader(LoadConfig(load_format="mxfp4_csf")) + result = dict(loader.get_all_weights(config, SimpleNamespace())) + assert sorted(result) == sorted(retained) + for key, value in retained.items(): + assert torch.equal(result[key], value) + dspark = SimpleNamespace(checkpoint_weight_name_prefixes=("mtp.",)) + assert sorted(dict(loader.get_all_weights(config, dspark))) == sorted(draft) diff --git a/vllm/config/load.py b/vllm/config/load.py index 8bb6760736ac..573353e18b19 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -35,7 +35,7 @@ class LoadConfig: back to the pytorch bin format if safetensors format is not available. - "pt" will load the weights in the pytorch bin format. - "safetensors" will load the weights in the safetensors format. - - "mxfp4_csf" reads losslessly compressed Kimi/DS4.1 + - "mxfp4_csf" reads losslessly compressed Kimi/DS4.1/DS4-Flash MXFP4 expert scales and retains file-backed loading for Engram tables. - "nvfp4_csf" reads losslessly compressed GLM-5.3-Flash or Qwen3.8-Flash-Next NVFP4 expert scales and retains native activation diff --git a/vllm/model_executor/layers/quantization/mxfp4_csf.py b/vllm/model_executor/layers/quantization/mxfp4_csf.py index c0e2bf95eade..de0306934d21 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_csf.py +++ b/vllm/model_executor/layers/quantization/mxfp4_csf.py @@ -8,7 +8,7 @@ class Mxfp4CsfConfig(KimiMxfp4CsfConfig): - """Select Kimi or DeepSeek dense precision from the checkpoint contract.""" + """Select Kimi, DeepSeek-V4.1 or DeepSeek-V4 dense precision by family.""" @classmethod def get_name(cls): @@ -32,6 +32,11 @@ def from_config(cls, config): family = checkpoint_contract(config["checkpoint_root"])["family"] if family == "deepseek_v41": return DeepseekV41Mxfp4CsfConfig.from_config(config) + if family == "deepseek_v4_flash": + # Deferred: the DeepSeek-V4 package imports its platform model. + from vllm.models.deepseek_v4.mxfp4_csf import DeepseekV4Mxfp4CsfConfig + + return DeepseekV4Mxfp4CsfConfig.from_config(config) if family == "kimi_k3": return super().from_config(config) raise ValueError(f"Unsupported MXFP4-CSF model family: {family}") diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py index e060dbec4b2c..fbd3d4a316c6 100644 --- a/vllm/model_executor/model_loader/mxfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -31,9 +31,11 @@ SCHEMA = "lil-mxfp4-csf-checkpoint/1" CODEC = "row-base-offset1-u24-exceptions/1" +# Routed experts, hidden size, MoE intermediate size and compressed layers. FAMILIES = { - "deepseek_v41": (384, 5120, 2304), - "kimi_k3": (896, 3584, 3072), + "deepseek_v41": (384, 5120, 2304, range(40)), + "deepseek_v4_flash": (256, 4096, 2048, range(43)), + "kimi_k3": (896, 3584, 3072, range(1, 93)), } @@ -43,6 +45,24 @@ def checkpoint_contract(root: str) -> dict: return read_csf_contract(root, schema=SCHEMA, codec=CODEC, families=FAMILIES) +def expert_tensor_names(family, layer_index, expert): + """Source weight and scale names of an expert's gate, up and down matrices.""" + if family == "kimi_k3": + prefix = ( + f"language_model.model.layers.{layer_index}." + f"block_sparse_moe.experts.{expert}" + ) + weight, scale = "weight_packed", "weight_scale" + else: + # DeepSeek-V4.1-Flash and DeepSeek-V4-Flash, including its vision + # variant, keep DeepSeek's native routed-expert names. + prefix = f"layers.{layer_index}.ffn.experts.{expert}" + weight, scale = "weight", "scale" + return tuple( + (f"{prefix}.{p}.{weight}", f"{prefix}.{p}.{scale}") for p in ("w1", "w3", "w2") + ) + + def read_mxfp4_csf_layer( root, layer_index, @@ -59,9 +79,12 @@ def read_mxfp4_csf_layer( """Read rank-local gate/up/down tensors and unprepared compressed scales.""" contract = checkpoint_contract(str(Path(root).resolve())) family = contract["family"] - if (num_experts, hidden_size, intermediate_size) != FAMILIES[family]: + e, h, n, layers = FAMILIES[family] + if (num_experts, hidden_size, intermediate_size) != (e, h, n): raise ValueError("MXFP4-CSF expert geometry differs from the checkpoint family") - supported_tp = (1, 2, 4, 8) if family == "deepseek_v41" else (1, 2, 4, 8, 12, 16) + if layer_index not in layers: + raise ValueError("MXFP4-CSF layer is outside the compressed expert inventory") + supported_tp = (1, 2, 4, 8, 12, 16) if family == "kimi_k3" else (1, 2, 4, 8) if tp_size not in supported_tp or not 0 <= tp_rank < tp_size: raise ValueError( f"MXFP4-CSF {family} supports TP {supported_tp} with a valid rank" @@ -70,18 +93,11 @@ def read_mxfp4_csf_layer( def experts(): for expert in range(num_experts): - if family == "kimi_k3": - prefix = ( - f"language_model.model.layers.{layer_index}." - f"block_sparse_moe.experts.{expert}" - ) - weight, scale = "weight_packed", "weight_scale" - else: - prefix = f"layers.{layer_index}.ffn.experts.{expert}" - weight, scale = "weight", "scale" yield tuple( - reader.matrix(f"{prefix}.{p}.{weight}", f"{prefix}.{p}.{scale}") - for p in ("w1", "w3", "w2") + reader.matrix(weight, scale) + for weight, scale in expert_tensor_names( + family, layer_index, expert + ) ) return _load_mxfp4_csf_weights( diff --git a/vllm/models/deepseek_v4/mxfp4_csf.py b/vllm/models/deepseek_v4/mxfp4_csf.py new file mode 100644 index 000000000000..d1f121df88ec --- /dev/null +++ b/vllm/models/deepseek_v4/mxfp4_csf.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Losslessly compressed DeepSeek-V4-Flash routed experts with native dense tensors. + +Serves DeepSeek-V4-Flash and DeepSeek-V4-Flash-Vision-Exp. Only the target +model's routed experts read compressed scales; dense, attention, shared-expert, +vision and MTP/DSpark draft tensors keep the source configuration. +""" + +from typing import cast + +import regex as re +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config, get_current_vllm_config_or_none +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.models.deepseek_v4_1.mxfp4_csf import Mxfp4CsfMoEMethod +from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config + + +class DeepseekV4Mxfp4CsfConfig(DeepseekV4FP8Config): + """MXFP4-CSF target experts; native block-FP8 dense and draft weights. + + The base is the class CUDA registers as ``deepseek_v4_fp8``, which serves + the uncompressed DeepSeek-V4-Flash checkpoints. + """ + + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "mxfp4_csf" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if ( + user_quant in (None, cls.get_name()) + and hf_quant_cfg is not None + and hf_quant_cfg.get("quant_method") == cls.get_name() + and getattr(hf_config, "model_type", None) == "deepseek_v4" + ): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("MXFP4-CSF requires format_version=1 and checkpoint_root") + if config.get("weight_block_size") != [128, 128]: + raise ValueError( + "DeepSeek-V4-Flash MXFP4-CSF requires the source block-FP8 " + "[128, 128] quantization config" + ) + # The serving config keeps the source fields and replaces the method. + source = { + key: value + for key, value in config.items() + if key not in ("format_version", "checkpoint_root") + } + result = cast( + "DeepseekV4Mxfp4CsfConfig", + super().from_config({**source, "quant_method": "fp8"}), + ) + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + current = get_current_vllm_config_or_none() + if current is not None and current.parallel_config.enable_expert_parallel: + # MegaMoE experts bypass quantization methods and would never read + # the compressed experts. + raise NotImplementedError( + "MXFP4-CSF DeepSeek-V4-Flash supports TP without expert parallelism" + ) + if isinstance(layer, RoutedExperts): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) + if match is None: + raise ValueError("DeepSeek-V4 routed experts require a numbered layer") + config = get_current_vllm_config() + # MTP and DSpark draft layers follow the target layers. + if int(match.group(1)) < config.model_config.hf_config.num_hidden_layers: + if self.expert_dtype != "fp4" or self.moe_quant_algo: + raise ValueError("MXFP4-CSF requires MXFP4 routed experts") + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "MXFP4-CSF shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format not in ("mxfp4_csf",): + raise ValueError("MXFP4-CSF requires --load-format mxfp4_csf ") + return Mxfp4CsfMoEMethod( + layer.moe_config, + self, + activation_mode="a16" if envs.VLLM_B12X_MOE_FP4_FORCE_A16 else "a8", + ) + return super().get_quant_method(layer, prefix) diff --git a/vllm/models/deepseek_v4_1/mxfp4_csf.py b/vllm/models/deepseek_v4_1/mxfp4_csf.py index a156b664ef33..32ea19c7e479 100644 --- a/vllm/models/deepseek_v4_1/mxfp4_csf.py +++ b/vllm/models/deepseek_v4_1/mxfp4_csf.py @@ -113,7 +113,7 @@ def __init__(self, moe, owner, *, activation_mode="a16"): or parallel.use_all2all_kernels or parallel.enable_eplb ): - raise NotImplementedError("MXFP4-CSF DS4.1 supports TP without EP/DP") + raise NotImplementedError("MXFP4-CSF experts support TP without EP/DP") if ( moe.activation not in (MoEActivation.SILU, MoEActivation.SITU) or moe.in_dtype != torch.bfloat16 From 2338b9cb4d62b8a84afb0c3009dc6cb06b8da80e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 03:57:25 +0000 Subject: [PATCH 352/365] release: link the DS4-Flash MXFP4-CSF fragment to #971 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-mxfp4-csf-ds4-flash.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-mxfp4-csf-ds4-flash.json b/.lil/changes/vllm-mxfp4-csf-ds4-flash.json index 432453d914cb..1ae3ae26079a 100644 --- a/.lil/changes/vllm-mxfp4-csf-ds4-flash.json +++ b/.lil/changes/vllm-mxfp4-csf-ds4-flash.json @@ -17,6 +17,8 @@ "authors": [ "voipmonitor" ], - "pull_requests": [], + "pull_requests": [ + 971 + ], "requires": [] } From 44100d39a0c02fdcad541ff288bb02957dcd5cb5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 05:39:26 +0000 Subject: [PATCH 353/365] qwen4_exp: resolve the PLE storage dtype from an FP4-CSF checkpoint root An FP4-CSF serving directory holds only metadata; its config points at the checkpoint root whose tensors/ directory holds the weights, PLE shards included. The PLE dtype probe read the serving directory, found no shards, and left the dtype at its default, so the Qwen3.8-Flash-Next QAD CSF checkpoint (NVFP4 PLE table) failed to load with a PLE shape mismatch. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-qwen-ple-dtype-csf.json | 18 +++++++++++++ tests/models/qwen4_exp/test_b12x_config.py | 31 ++++++++++++++++++++++ vllm/models/qwen4_exp/config.py | 14 +++++++++- 3 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 .lil/changes/vllm-qwen-ple-dtype-csf.json diff --git a/.lil/changes/vllm-qwen-ple-dtype-csf.json b/.lil/changes/vllm-qwen-ple-dtype-csf.json new file mode 100644 index 000000000000..3c62e3d48990 --- /dev/null +++ b/.lil/changes/vllm-qwen-ple-dtype-csf.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-qwen-ple-dtype-csf", + "category": "fix", + "summary": "Qwen3.8-Flash-Next QAD FP4-CSF checkpoints start: the PLE storage dtype is read from the checkpoint the serving directory points at", + "models": [ + "Qwen3.8-Flash-Next" + ], + "compatibility": "Only FP4-CSF serving directories change; other checkpoints resolve the PLE dtype as before.", + "details": [ + "When the checkpoint config omits the PLE storage dtype, vLLM reads it from the safetensors headers of the model directory. An FP4-CSF serving directory holds metadata only, so the QAD checkpoint's NVFP4 PLE table was taken for BF16 and loading stopped with a PLE shard shape mismatch. The dtype is now read from the tensors under the CSF checkpoint root." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [], + "requires": [] +} diff --git a/tests/models/qwen4_exp/test_b12x_config.py b/tests/models/qwen4_exp/test_b12x_config.py index adf802d261d1..f44dc56853c6 100644 --- a/tests/models/qwen4_exp/test_b12x_config.py +++ b/tests/models/qwen4_exp/test_b12x_config.py @@ -116,6 +116,37 @@ def test_omitted_ple_dtype_is_resolved_from_checkpoint_headers( assert config.get_text_config().ple_embedding_dtype == expected +@pytest.mark.parametrize("method", ["nvfp4_csf", "mxfp4_csf"]) +def test_omitted_ple_dtype_is_resolved_from_the_csf_checkpoint_root( + tmp_path, method +) -> None: + """An FP4-CSF serving directory has no tensors; the PLE dtype comes from the + checkpoint root it points at.""" + root = tmp_path / "checkpoint" + (root / "tensors").mkdir(parents=True) + save_file( + { + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_0.weight": torch.zeros((2, 8), dtype=torch.uint8), + }, + root / "tensors" / "hybrid-main-00001.safetensors", + ) + serving = tmp_path / "serving" + serving.mkdir() + config_dict = Qwen3_8FlashNextConfig(**_TEXT_CONFIG, ple_layer_ids=[2]).to_dict() + config_dict.get("text_config", config_dict).pop("ple_embedding_dtype") + config_dict["quantization_config"] = { + "quant_method": method, + "format_version": 1, + "checkpoint_root": str(root), + } + (serving / "config.json").write_text(json.dumps(config_dict)) + + config = get_config(str(serving), trust_remote_code=False) + + assert config.get_text_config().ple_embedding_dtype == "nvfp4" + + def test_explicit_ple_dtype_does_not_probe_checkpoint(tmp_path, monkeypatch) -> None: config = Qwen4ExpConfig( **_TEXT_CONFIG, ple_layer_ids=[2], ple_embedding_dtype="bfloat16" diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py index 6d9c3f8e0a47..bf82c77987aa 100644 --- a/vllm/models/qwen4_exp/config.py +++ b/vllm/models/qwen4_exp/config.py @@ -38,7 +38,19 @@ def resolve_ple_embedding_dtype( from vllm.transformers_utils.config import get_safetensors_params_metadata - metadata = get_safetensors_params_metadata(str(model), revision=revision) + source = str(model) + quant_config = ( + config_dict.get("quantization_config") + or text_config_dict.get("quantization_config") + or {} + ) + if quant_config.get("quant_method") in ("nvfp4_csf", "mxfp4_csf") and ( + root := quant_config.get("checkpoint_root") + ): + # An FP4-CSF serving directory holds metadata only; the checkpoint's + # tensors, PLE shards included, live under its root. + source = str(Path(root) / "tensors") + metadata = get_safetensors_params_metadata(source, revision=revision) dtypes = { info["dtype"] for name, info in metadata.items() From 87b54c511571591cc2ed3df19bb2d36e65c87c35 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 05:39:50 +0000 Subject: [PATCH 354/365] release: link the PLE dtype fragment to #974 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-qwen-ple-dtype-csf.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-qwen-ple-dtype-csf.json b/.lil/changes/vllm-qwen-ple-dtype-csf.json index 3c62e3d48990..e2437f513ea1 100644 --- a/.lil/changes/vllm-qwen-ple-dtype-csf.json +++ b/.lil/changes/vllm-qwen-ple-dtype-csf.json @@ -13,6 +13,8 @@ "authors": [ "voipmonitor" ], - "pull_requests": [], + "pull_requests": [ + 974 + ], "requires": [] } From d7af3ad174040404bfd82b941abe9c72f4529d5a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:07:19 +0000 Subject: [PATCH 355/365] 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 --- 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 da6826cfbf66..85f72f3142d2 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 dbece926c942b894e2967c3f6961671870fb6aee Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:08:01 +0000 Subject: [PATCH 356/365] release: fragment for the zero-initialized NVFP4 MoE input scales (#976) Co-Authored-By: Claude Opus 5.5 --- .../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 df3d9833a4761b6e4132fed8350552e844be98be Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 13:53:16 +0000 Subject: [PATCH 357/365] 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 --- 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 490a9590d689..0fdf0714e4a1 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -236,21 +236,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 @@ -263,15 +265,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 @@ -1046,34 +1058,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 109d044369f4c680deda36ceefd830169efe08a5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 14:43:54 +0000 Subject: [PATCH 358/365] release: fragment for the row-type A4 prefill (#977) Co-Authored-By: Claude Opus 5.5 --- .../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" + ] +} From 53c7474b8641e96f937f782ad419938119785e32 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 15:48:00 +0000 Subject: [PATCH 359/365] NVFP4-CSF: expand the next layer's scales during its attention Layers of an NVFP4-CSF model share one scale scratch, and every eager prefill call expanded its experts' scales into it serially before its MoE (4% of a GLM-5.3 glm53-tp2 A4 prefill, with nothing overlapping). After a layer's MoE the method now expands the following layer's scales on a side stream, ordered after that MoE, and the following layer waits for it and binds with scales_expanded=True (b12x expand_scales). Only eager steps that expand - A4 prefill rows or A16 calls above the stage-read limit - prefetch; CUDA graph capture, decode-only steps and other forward passes are untouched, and any pending expansion is waited for before the scratch is used. VLLM_B12X_CSF_SCALE_PREFETCH=0 turns it off. Co-Authored-By: Claude Opus 5.5 --- tests/quantization/test_nvfp4_csf.py | 83 ++++++++++++ vllm/envs.py | 6 + vllm/model_executor/layers/fused_moe/b12x.py | 7 + .../layers/quantization/nvfp4_csf.py | 122 ++++++++++++++++-- 4 files changed, 204 insertions(+), 14 deletions(-) diff --git a/tests/quantization/test_nvfp4_csf.py b/tests/quantization/test_nvfp4_csf.py index f988b64d69e4..e2d34d53fb79 100644 --- a/tests/quantization/test_nvfp4_csf.py +++ b/tests/quantization/test_nvfp4_csf.py @@ -281,3 +281,86 @@ def test_missing_or_invalid_calibration_is_rejected(bad): w13_scale_scratch=None, w2_scale_scratch=None, ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA streams") +def test_scale_prefetch_expands_the_next_layer_for_this_forward_only(monkeypatch): + """Layer N expands layer N+1's scales on the side stream after its own MoE; + layer N+1 skips its expansion in the same forward pass only.""" + pytest.importorskip("b12x") + import b12x.moe.fused_moe as fused_moe + + import vllm.forward_context as forward_context + + expanded = [] + monkeypatch.setattr( + fused_moe, + "expand_scales", + lambda prepared: expanded.append((prepared, torch.cuda.current_stream())), + ) + monkeypatch.setenv("B12X_W4A16_A4_PREFILL_MIN_TOKENS", "1536") + + def step(starts, gdn_prefills=None): + metadata = { + "attn": SimpleNamespace( + query_start_loc=torch.tensor(starts, dtype=torch.int32), + num_actual_tokens=starts[-1], + ) + } + if gdn_prefills is not None: + metadata["kda"] = 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 torch.zeros(starts[-1], 8) + + owner = SimpleNamespace( + scale_layers={}, scale_stream=torch.cuda.Stream(), scale_prefetch=None + ) + calls = [] + + def method(index): + m = object.__new__(Nvfp4CsfMoEMethod) + m.owner, m.layer_index, m.a4_prefill = owner, index, True + m.backend = SimpleNamespace(scales_expanded=False) + m.prepared = f"layer-{index}" + m.moe_done, m.scales_ready = torch.cuda.Event(), torch.cuda.Event() + m.moe_kernel = SimpleNamespace( + apply=lambda **_: calls.append((index, m.backend.scales_expanded)) + ) + owner.scale_layers[index] = m + return m + + first, second = method(3), method(4) + layer = SimpleNamespace( + w13_weight=None, + w2_weight=None, + activation=MoEActivation.SILU, + global_num_experts=8, + expert_map=None, + apply_router_weight_on_input=False, + ) + + def forward(x): + for m in (first, second): + m.apply(layer, x, None, None, None, None) + + forward(step([0, 4, 300])) # decode rows, then prefill rows + assert calls == [(3, False), (4, True)] + assert [p for p, _ in expanded] == ["layer-4"] + assert expanded[0][1] == owner.scale_stream + assert owner.scale_prefetch is None + # A prefetch left from another forward pass is waited for, never reused. + owner.scale_prefetch = (4, -1, second.scales_ready) + calls.clear() + second.apply(layer, step([0, 300]), None, None, None, None) + assert calls == [(4, False)] + # Decode-only steps (a GDN layer counts no prefills) expand nothing. + expanded.clear() + calls.clear() + forward(step([0, 4, 8], gdn_prefills=0)) + assert calls == [(3, False), (4, False)] and expanded == [] diff --git a/vllm/envs.py b/vllm/envs.py index f24a795b84b6..01ef6bd7c2dd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -206,6 +206,7 @@ VLLM_B12X_MOE_FP4_FORCE_A16: bool = False VLLM_B12X_BF16_GEMV: bool = False VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE: Literal["0", "1", "all", "w13", "w2"] = "0" + VLLM_B12X_CSF_SCALE_PREFETCH: bool = True VLLM_DEFAULT_MOE_BACKEND: str = "auto" VLLM_B12X_DENSE_ACTIVATION_MODE: Literal["auto", "a16", "quantized"] = "auto" VLLM_B12X_NVFP4_ACTIVATION_MODE: Literal["auto", "a16", "quantized"] | None = None @@ -1822,6 +1823,11 @@ def _resolve_rust_cli_path() -> str | None: "0", ["0", "1", "all", "w13", "w2"], ), + # Expand the next NVFP4-CSF layer's expert scales on a side stream while its + # attention runs, instead of before its MoE (eager prefill calls only). + "VLLM_B12X_CSF_SCALE_PREFETCH": lambda: bool( + int(os.getenv("VLLM_B12X_CSF_SCALE_PREFETCH", "1")) + ), # Dense activation precision; recipe overrides take precedence. "VLLM_B12X_DENSE_ACTIVATION_MODE": env_with_choices( "VLLM_B12X_DENSE_ACTIVATION_MODE", "auto", ["auto", "a16", "quantized"] diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index 0fdf0714e4a1..da84a239396e 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -313,6 +313,10 @@ def _b12x_has_mxfp8_moe() -> bool: class B12xExperts(mk.FusedMoEExpertsModular): """Packed MoE experts backed by the b12x SM12x planned API.""" + # Set by the NVFP4-CSF method while the shared scale scratch holds this + # layer's scales, expanded ahead of the call (b12x expand_scales()). + scales_expanded = False + def __init__( self, moe_config: mk.FusedMoEConfig, @@ -1058,6 +1062,7 @@ def apply( raise ValueError("b12x MoE requires contiguous caller-owned workspace2") scratch = workspace2.view(-1).view(torch.uint8) tokens = int(hidden_states.shape[0]) + expanded = {"scales_expanded": True} if self.scales_expanded else {} if ( self._quant_mode == "w4a16" and _w4a16_a4_prefill_enabled() @@ -1095,6 +1100,7 @@ def apply( output=output[lo:hi], input_scales_static=True, a4_prefill=a4_prefill, + **expanded, ) _require_b12x_fused_moe().run(binding=binding) return @@ -1107,6 +1113,7 @@ def apply( topk_ids=topk_ids, output=output, input_scales_static=True, + **expanded, ) _require_b12x_fused_moe().run(binding=binding) diff --git a/vllm/model_executor/layers/quantization/nvfp4_csf.py b/vllm/model_executor/layers/quantization/nvfp4_csf.py index 55398c32102c..57d744c7b506 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_csf.py +++ b/vllm/model_executor/layers/quantization/nvfp4_csf.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Lossless NVFP4 expert scale storage with native ModelOpt arithmetic.""" +import itertools from dataclasses import replace import regex as re @@ -16,7 +17,12 @@ from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import modular_kernel as mk from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.b12x import ( + B12xExperts, + _is_current_stream_capturing, + _num_leading_decode_tokens, + _w4a16_a4_prefill_enabled, +) from vllm.model_executor.layers.fused_moe.config import ( nvfp4_moe_quant_config, nvfp4_w4a16_moe_quant_config, @@ -34,12 +40,40 @@ logger = init_logger(__name__) +_FORWARD_TOKENS = itertools.count() + + +def _forward_token() -> int | None: + """An id of the current forward pass, or None outside of one.""" + from vllm.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return None + context = get_forward_context() + token = getattr(context, "_b12x_csf_token", None) + if token is None: + token = next(_FORWARD_TOKENS) + context._b12x_csf_token = token + return token + + +def _stage_max_tokens() -> int: + """Calls up to this size read compressed scales per stage, without expansion.""" + try: + from b12x.moe.fused_moe._impl import W4A16_CSF_STAGE_MAX_TOKENS + except ImportError: + return 1536 + return W4A16_CSF_STAGE_MAX_TOKENS + class Nvfp4CsfConfig(ModelOptMixedPrecisionConfig): """Compressed main experts; retained ModelOpt formats for all other tensors.""" checkpoint_root: str scale_scratch: tuple[torch.Tensor, ...] | None + scale_layers: dict[int, "Nvfp4CsfMoEMethod"] + scale_stream: torch.cuda.Stream | None + scale_prefetch: tuple[int, int, torch.cuda.Event] | None @classmethod def get_name(cls): @@ -66,6 +100,11 @@ def from_config(cls, config): assert isinstance(result, cls) result.checkpoint_root = config["checkpoint_root"] result.scale_scratch = None + # Layers sharing scale_scratch by index, the side stream that expands + # the next layer's scales ahead of its MoE, and the pending expansion. + result.scale_layers = {} + result.scale_stream = None + result.scale_prefetch = None return result def get_quant_method(self, layer, prefix): @@ -240,6 +279,14 @@ def process_weights_after_loading(self, layer): self.moe_kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoDPEPModular(), backend ) + self.backend, self.prepared = backend, prepared + self.a4_prefill = bool(getattr(prepared._impl, "a4_prefill_scales", False)) + self.moe_done = torch.cuda.Event() + self.scales_ready = torch.cuda.Event() + if self.use_a16 and envs.VLLM_B12X_CSF_SCALE_PREFETCH: + self.owner.scale_layers[self.layer_index] = self + if self.owner.scale_stream is None: + self.owner.scale_stream = torch.cuda.Stream(device) logger.info( "NVFP4-CSF layer %d rank %d/%d: native NVFP4 %s, " "shared scale scratch %d bytes", @@ -261,17 +308,64 @@ def apply( workspace=None, ): assert self.moe_kernel is not None - return self.moe_kernel.apply( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts=shared_experts, - shared_experts_input=shared_experts_input, - workspace=workspace, + + def run(): + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + workspace=workspace, + ) + + owner = self.owner + if self.layer_index not in owner.scale_layers or _is_current_stream_capturing(): + return run() + token = _forward_token() + pending, owner.scale_prefetch = owner.scale_prefetch, None + stream = torch.cuda.current_stream() + if pending is not None: + # The shared scratch must never race an expansion still in flight. + stream.wait_event(pending[2]) + self.backend.scales_expanded = ( + token is not None + and pending is not None + and (pending[:2] == (self.layer_index, token)) ) + try: + result = run() + finally: + self.backend.scales_expanded = False + following = owner.scale_layers.get(self.layer_index + 1) + if token is not None and following is not None and self._expands(x): + # The next layer's MoE would expand its scales into the same scratch: + # do it now on a side stream, overlapping that layer's attention. + from b12x.moe import fused_moe + + self.moe_done.record(stream) + side = owner.scale_stream + side.wait_event(self.moe_done) + with torch.cuda.stream(side): + fused_moe.expand_scales(following.prepared) + following.scales_ready.record(side) + owner.scale_prefetch = ( + following.layer_index, + token, + following.scales_ready, + ) + return result + + def _expands(self, x: torch.Tensor) -> bool: + """Whether this step's calls expand scales: A4 prefill rows or a large + A16 call (smaller A16 calls read compressed scales per stage).""" + tokens = int(x.shape[0]) + if self.a4_prefill and _w4a16_a4_prefill_enabled(): + return _num_leading_decode_tokens(tokens) < tokens + return tokens > _stage_max_tokens() From fcad04e4ebe4918dffc002f85625c97a82f71ba5 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 16:43:53 +0000 Subject: [PATCH 360/365] release: fragment for the NVFP4-CSF scale prefetch (#978) Co-Authored-By: Claude Opus 5.5 --- .../vllm-nvfp4-csf-scale-prefetch.json | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 .lil/changes/vllm-nvfp4-csf-scale-prefetch.json diff --git a/.lil/changes/vllm-nvfp4-csf-scale-prefetch.json b/.lil/changes/vllm-nvfp4-csf-scale-prefetch.json new file mode 100644 index 000000000000..d27186f83d8d --- /dev/null +++ b/.lil/changes/vllm-nvfp4-csf-scale-prefetch.json @@ -0,0 +1,22 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-nvfp4-csf-scale-prefetch", + "category": "performance", + "summary": "GLM-5.3-Flash prefill on FP4-CSF checkpoints is 2-3% faster with A4 and 1-2% with A16: the next layer's scales expand during its attention", + "models": [ + "GLM-5.3-Flash" + ], + "compatibility": "On by default for NVFP4-CSF W4A16 layers; VLLM_B12X_CSF_SCALE_PREFETCH=0 turns it off. Decode is unchanged.", + "details": [ + "Layers share one scale scratch and each eager prefill call expanded its experts' scales serially before its MoE. After a layer's MoE, the next layer's scales now expand on a side stream, overlapping that layer's attention. glm53-tp2 on two RTX PRO 6000 Max-Q: 8K/32K prefill 10,254/10,620 instead of 10,032/10,312 tok/s with A4 and 7,763/8,058 instead of 7,672/7,925 with A16; needle-checksum A4 495/490/492, A16 498/495/497." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 978 + ], + "requires": [ + "b12x-nvfp4-csf-expand-scales" + ] +} From f2d127220f50c37c9e3ef50e754c84fcabfcd3ba Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 19:47:32 +0000 Subject: [PATCH 361/365] [Worker] Load weights without the split limit under expandable segments #41268 sets max_split_size_mb:20 while loading weights so cached blocks do not fragment. With expandable segments the allocator releases free pages instead of whole blocks, and the limit instead strands the pages that persistent weights share with freed loading temporaries: empty_cache() cannot unmap them, and the memory profile then charges them against the KV cache. Measured with torch.cuda.memory_reserved() - memory_allocated() after loading, per GPU: - GLM-5.3-Flash glm53-tp2 (NVFP4-CSF, stage-read scales): 0.59 -> 0.08 GiB; with per-layer expansion 0.42 -> 0.12 GiB. Both start and serve; the preset's fixed KV cache leaves that much more headroom. - DeepSeek-V4.1-Flash TP4 with b12x#479's inline scales before its own allocator fix: 1.01 -> 0.04 GiB, KV cache 11.96M -> 12.86M tokens, weight loading 143 -> 120 s. The classic allocator keeps the limit. Co-authored-by: Claude Opus 5.5 Signed-off-by: Martin Vit --- tests/v1/worker/test_gpu_worker.py | 29 +++++++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 6 ++++++ 2 files changed, 35 insertions(+) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index c17eec5e6cf1..961542d88707 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -197,6 +197,35 @@ def test_startup_plan_survives_profiling_config_rewrites(plan_env): assert second.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes +@pytest.mark.parametrize( + "conf,limited", + [ + ("expandable_segments:True", False), + ("expandable_segments:True,large_segment_size_mb:12", False), + ("", True), + ("expandable_segments:False", True), + ], +) +def test_weight_loading_split_limit_skips_expandable_segments( + monkeypatch, conf, limited +): + """Expandable segments load weights without the temporary split limit. + + With them, the limit strands pages that weights share with freed loading + temporaries; the classic allocator keeps it against cached-block + fragmentation. + """ + settings: list[str] = [] + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", conf) + monkeypatch.setattr(gpu_worker.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + gpu_worker.torch._C, "_accelerator_setAllocatorSettings", settings.append + ) + with gpu_worker.Worker._scoped_allocator_max_split(None, max_split_size_mb=20): + assert settings == (["max_split_size_mb:20"] if limited else []) + assert len(settings) == (2 if limited else 0) + + # Memory accounting of the profiling run (Worker.determine_available_memory). diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3c12397a1186..4557b2ec5c84 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -375,6 +375,12 @@ def _scoped_allocator_max_split(self, max_split_size_mb: int): return conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") + if re.search(r"expandable_segments:\s*True", conf): + # Expandable segments release free pages, not whole cached blocks. + # A split limit there strands the pages that weights share with + # freed loading temporaries (1 GiB per GPU on DeepSeek-V4.1 TP4). + yield + return match = re.search(r"max_split_size_mb:(\d+)", conf) original_value = match.group(1) if match else None From ffb0d84638c3a9778630869000e3a0f2bcb606c9 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 4 Oct 2026 19:48:03 +0000 Subject: [PATCH 362/365] release: fragment for loading without the split limit (#980) Co-authored-by: Claude Opus 5.5 Signed-off-by: Martin Vit --- .../vllm-load-split-expandable-segments.json | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .lil/changes/vllm-load-split-expandable-segments.json diff --git a/.lil/changes/vllm-load-split-expandable-segments.json b/.lil/changes/vllm-load-split-expandable-segments.json new file mode 100644 index 000000000000..0d5c2c6f235b --- /dev/null +++ b/.lil/changes/vllm-load-split-expandable-segments.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-load-split-expandable-segments", + "category": "performance", + "summary": "Weight loading no longer strands GPU memory when expandable segments are on: 0.5 GiB more free per GPU on glm53-tp2", + "models": [ + "GLM-5.3-Flash", + "DeepSeek-V4.1-Flash" + ], + "compatibility": "No user action required. Applies when PYTORCH_CUDA_ALLOC_CONF enables expandable_segments, as the shipped profiles do; the classic allocator still loads with max_split_size_mb:20.", + "details": [ + "vLLM loaded weights with max_split_size_mb:20 (upstream #41268). With expandable segments that limit left pages shared by weights and freed loading temporaries mapped, which empty_cache cannot release and the memory profile counts against the KV cache. Measured after loading, per GPU: glm53-tp2 0.59 -> 0.08 GiB (per-layer scale expansion 0.42 -> 0.12 GiB); DeepSeek-V4.1-Flash TP4 with inline CSF scales before b12x's own fix 1.01 -> 0.04 GiB, KV cache 11.96M -> 12.86M tokens. glm53-tp2 keeps its fixed 8 GiB KV cache, so the memory is headroom." + ], + "authors": [ + "voipmonitor" + ], + "pull_requests": [ + 980 + ], + "requires": [] +} From 65c16ddea68f1e8a85476a210869019adcfc20cc Mon Sep 17 00:00:00 2001 From: Christopher Maher Date: Sun, 4 Oct 2026 20:10:13 -0700 Subject: [PATCH 363/365] MXFP4-CSF: load DeepSeek-V4.1-Flash checkpoints at TP3 The deepseek_v41 reader admitted TP 1, 2, 4 and 8 only. At TP3 each rank owns 768 of the 2304 routed-expert channels, which is 32-channel aligned and falls on whole 16-row scale slabs, so the slicer and the b12x W4A8 CSF geometry already hold. Add 3 for deepseek_v41 only; the deepseek_v4_flash (2048 channels) and kimi_k3 lists are unchanged. Tests: the TP byte-slicing test also runs at TP3, and a new test pins which families admit TP3. Signed-off-by: Christopher Maher --- .lil/changes/vllm-mxfp4-csf-ds41-tp3.json | 15 +++++ tests/quantization/test_mxfp4_csf.py | 56 +++++++++++++++++-- .../model_loader/mxfp4_csf_loader.py | 6 +- 3 files changed, 72 insertions(+), 5 deletions(-) create mode 100644 .lil/changes/vllm-mxfp4-csf-ds41-tp3.json diff --git a/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json b/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json new file mode 100644 index 000000000000..8bec5137a6e4 --- /dev/null +++ b/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-mxfp4-csf-ds41-tp3", + "category": "feature", + "summary": "Load DeepSeek-V4.1-Flash MXFP4-CSF checkpoints at TP3", + "models": ["DeepSeek-V4.1-Flash"], + "compatibility": "No user action required. TP1, TP2, TP4 and TP8 are unchanged; DeepSeek-V4-Flash and Kimi-K3 keep their existing TP sizes.", + "details": [ + "The deepseek_v41 MXFP4-CSF reader admits --tensor-parallel-size 3. Each rank owns 768 of the 2304 routed-expert channels, which keeps the 32-channel and 16-row scale-slab alignment the slicer requires.", + "The checkpoint keeps its source config; the existing TP3 attention head padding applies at runtime." + ], + "authors": ["Defilan"], + "pull_requests": [], + "requires": [] +} diff --git a/tests/quantization/test_mxfp4_csf.py b/tests/quantization/test_mxfp4_csf.py index 2f2c5f1691b5..5ac7e42491b6 100644 --- a/tests/quantization/test_mxfp4_csf.py +++ b/tests/quantization/test_mxfp4_csf.py @@ -530,15 +530,16 @@ def test_rejects_unaligned_rows_and_bad_padding(): slice_scale_plane(fixed, exceptions, 64, 9, (0, 64), (0, 9)) -@pytest.mark.parametrize("rank", [0, 1]) -def test_tensor_sources_preserve_projection_order_and_tp_bytes(rank): +@pytest.mark.parametrize("tp_size,rank", [(2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]) +def test_tensor_sources_preserve_projection_order_and_tp_bytes(tp_size, rank): """An arbitrary geometry loads without any model identity or checkpoint path.""" from vllm.model_executor.model_loader.mxfp4_csf_loader import ( _load_mxfp4_csf_weights as load_mxfp4_csf_weights, ) device = "cpu" - experts, hidden, intermediate, local = 3, 256, 192, 96 + experts, hidden, local = 3, 256, 96 + intermediate = local * tp_size sources, scales = [], [] for expert in range(experts): pairs = [ @@ -561,7 +562,7 @@ def test_tensor_sources_preserve_projection_order_and_tp_bytes(rank): hidden_size=hidden, intermediate_size=intermediate, tp_rank=rank, - tp_size=2, + tp_size=tp_size, device=device, w13_scale_scratch=scratch13, w2_scale_scratch=scratch2, @@ -749,6 +750,53 @@ def test_ds4_flash_reader_rejects_draft_layers_geometry_and_tp( ) +@pytest.mark.parametrize( + "family,geometry,tp_size,admitted", + [ + ("deepseek_v41", (384, 5120, 2304), 3, True), + ("deepseek_v41", (384, 5120, 2304), 6, False), + ("deepseek_v4_flash", (256, 4096, 2048), 3, False), + ("kimi_k3", (896, 3584, 3072), 3, False), + ], +) +def test_reader_admits_tp3_for_deepseek_v41_only( + monkeypatch, family, geometry, tp_size, admitted +): + from vllm.model_executor.model_loader import mxfp4_csf_loader as checkpoint + + class PastTpGuard(Exception): + pass + + def reader(*args, **kwargs): + raise PastTpGuard + + monkeypatch.setattr( + checkpoint, + "checkpoint_contract", + lambda _: {"family": family, "source_names": {}}, + ) + monkeypatch.setattr(checkpoint, "CsfTensorReader", reader) + num_experts, hidden_size, intermediate_size = geometry + expected = ( + pytest.raises(PastTpGuard) + if admitted + else pytest.raises(ValueError, match="supports TP") + ) + with expected: + checkpoint.read_mxfp4_csf_layer( + "/csf", + 1, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_rank=tp_size - 1, + tp_size=tp_size, + device="cpu", + w13_scale_scratch=None, + w2_scale_scratch=None, + ) + + def test_mxfp4_csf_dispatches_ds4_flash_with_source_dense_config(monkeypatch): from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.layers.quantization.mxfp4_csf import Mxfp4CsfConfig diff --git a/vllm/model_executor/model_loader/mxfp4_csf_loader.py b/vllm/model_executor/model_loader/mxfp4_csf_loader.py index fbd3d4a316c6..dd3f27bcf27e 100644 --- a/vllm/model_executor/model_loader/mxfp4_csf_loader.py +++ b/vllm/model_executor/model_loader/mxfp4_csf_loader.py @@ -84,7 +84,11 @@ def read_mxfp4_csf_layer( raise ValueError("MXFP4-CSF expert geometry differs from the checkpoint family") if layer_index not in layers: raise ValueError("MXFP4-CSF layer is outside the compressed expert inventory") - supported_tp = (1, 2, 4, 8, 12, 16) if family == "kimi_k3" else (1, 2, 4, 8) + supported_tp = { + "kimi_k3": (1, 2, 4, 8, 12, 16), + # 2304 / 3 = 768 local channels: 32-aligned, whole 16-row scale slabs. + "deepseek_v41": (1, 2, 3, 4, 8), + }.get(family, (1, 2, 4, 8)) if tp_size not in supported_tp or not 0 <= tp_rank < tp_size: raise ValueError( f"MXFP4-CSF {family} supports TP {supported_tp} with a valid rank" From 3b488b1481b80c1f7afe830102de8324fb8cecd2 Mon Sep 17 00:00:00 2001 From: Christopher Maher Date: Sun, 4 Oct 2026 21:42:35 -0700 Subject: [PATCH 364/365] lil: link the release-note fragment to #982 Signed-off-by: Christopher Maher --- .lil/changes/vllm-mxfp4-csf-ds41-tp3.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json b/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json index 8bec5137a6e4..0948cd8b62bc 100644 --- a/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json +++ b/.lil/changes/vllm-mxfp4-csf-ds41-tp3.json @@ -10,6 +10,6 @@ "The checkpoint keeps its source config; the existing TP3 attention head padding applies at runtime." ], "authors": ["Defilan"], - "pull_requests": [], + "pull_requests": [982], "requires": [] } From 7f37b5961e0e8b6ad623b6ab2479ff650cb9919b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 5 Oct 2026 07:22:56 +0000 Subject: [PATCH 365/365] [b12x] Clarify A4 prefill logs for FP4-CSF and draft MoE layers FP4-CSF MoE layers logged 'native NVFP4 A16' even when their A4 prefill path was prepared, and the MTP draft layer, which has no calibrated input scales, warned 'A4 prefill disabled for it'. Users read that as hybrid A4/A16 having fallen back to full A16. CSF layers now log 'W4A16 decode, A4 prefill'. The draft-layer message names the layer, explains that it keeps A16, and is INFO. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-983.json | 14 ++++++++++++++ vllm/model_executor/layers/fused_moe/b12x.py | 19 ++++++++++++------- .../layers/quantization/nvfp4_csf.py | 8 +++++++- 3 files changed, 33 insertions(+), 8 deletions(-) create mode 100644 .lil/changes/vllm-983.json diff --git a/.lil/changes/vllm-983.json b/.lil/changes/vllm-983.json new file mode 100644 index 000000000000..343baabe27c5 --- /dev/null +++ b/.lil/changes/vllm-983.json @@ -0,0 +1,14 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-983", + "category": "fix", + "summary": "Startup logs show which GLM-5.3-Flash MoE layers prefill with A4; the MTP draft layer no longer warns as if A4 prefill were off", + "models": ["GLM-5.3-Flash", "Qwen3.8-Flash-Next"], + "compatibility": "Log messages only; serving is unchanged.", + "details": [ + "With PREFILL_ACTIVATIONS=a4 every FP4-CSF MoE layer used to log 'native NVFP4 A16', and the one layer without calibrated input scales (the MTP draft layer) logged the WARNING 'layer has no valid activation scales; A4 prefill disabled for it (A16 only)'. That read as if hybrid A4/A16 had fallen back to full A16, but the target layers did prefill with A4.", + "FP4-CSF layers now log 'W4A16 decode, A4 prefill' when the A4 path is prepared. The draft layer's message names the layer, explains that it stays A16 because it has no calibrated scales, and is logged at INFO." + ], + "pull_requests": [983], + "requires": [] +} diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index da84a239396e..9960f9c47fe3 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -421,6 +421,7 @@ def _prepare_experts( w2: torch.Tensor, activation: MoEActivation, params_dtype: torch.dtype, + layer_name: str | None = None, ) -> Any: """Build the b12x prepared-experts representation for these weights. @@ -502,18 +503,21 @@ def _prepare_experts( 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"), + "b12x W4A16 %s: prefill rows run A4, decode rows W4A16 " + "(shared input global scale %.4g, %d intermediate scales)", + layer_name or "MoE layer", float(a1_gscale[0]), int(a2_gscale.numel()), ) else: - # No calibrated activation scales (e.g. a draft layer): stay A16. + # No calibrated activation scales (an MTP draft layer): its + # prefill stays A16; layers with scales keep A4 prefill. a1_gscale = a2_gscale = None - logger.warning( - "b12x W4A16: layer has no valid activation scales; A4 prefill " - "disabled for it (A16 only)" + logger.info( + "b12x W4A16 %s: no calibrated activation scales (expected " + "for an MTP draft layer), so its prefill stays A16; every " + "layer with scales still prefills with A4", + layer_name or "MoE layer", ) else: a1_gscale = unit_scale @@ -620,6 +624,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2=layer.w2_weight, activation=layer.activation, params_dtype=self.moe_config.in_dtype, + layer_name=getattr(layer, "layer_name", None), ) prepared = self._reuse_prepared_storage(layer, prepared) if prepared.plan._impl.discards_source_parameters: diff --git a/vllm/model_executor/layers/quantization/nvfp4_csf.py b/vllm/model_executor/layers/quantization/nvfp4_csf.py index 57d744c7b506..477a5f0e9297 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_csf.py +++ b/vllm/model_executor/layers/quantization/nvfp4_csf.py @@ -287,13 +287,19 @@ def process_weights_after_loading(self, layer): self.owner.scale_layers[self.layer_index] = self if self.owner.scale_stream is None: self.owner.scale_stream = torch.cuda.Stream(device) + if not self.use_a16: + activations = "A4" + elif self.a4_prefill: + activations = "W4A16 decode, A4 prefill" + else: + activations = "A16" logger.info( "NVFP4-CSF layer %d rank %d/%d: native NVFP4 %s, " "shared scale scratch %d bytes", self.layer_index, rank, tp, - "A16" if self.use_a16 else "A4", + activations, sum(t.numel() for t in scratch), )