feat(glm5next): reduce speculative state memory with atomic LMCache recovery

.lil/changes/vllm-821.json

@@ -0,0 +1,21 @@
+{
+  "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.",
+    "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."
+  ],
+  "pull_requests": [821],
+  "requires": ["b12x-408", "b12x-409"],
+  "evidence": [
+    "https://github.com/local-inference-lab/rtx6kpro/blob/master/benchmarks/glm53-kda-recovery-lmcache.md"
+  ]
+}

tests/models/glm5next/test_kda_recoverssm.py

@@ -0,0 +1,242 @@
+# 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 import CacheConfig, VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.config.mamba import MambaBackendEnum
+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 _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(
+            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)

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,
+        )

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"