diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py index 6c9e67b0a2cc..ec46ad85f38a 100644 --- a/tests/models/test_glm5next_model.py +++ b/tests/models/test_glm5next_model.py @@ -252,7 +252,7 @@ def test_glm5next_kda_splits_mixed_decode_prefill_batch(monkeypatch) -> None: layer.gate_lower_bound = -5.0 layer.A_log = torch.ones(1) layer.dt_bias = torch.ones(1) - layer._b12x_kda_binding = None + layer._b12x_kda_plan = None layer.conv1d = SimpleNamespace( weight=torch.ones(3, 1, 3), bias=torch.zeros(3), @@ -339,6 +339,61 @@ def test_glm5next_alone_opts_into_b12x_kda_decode() -> None: assert Glm5NextLinearAttention.b12x_kda_null_state_index == 0 +@pytest.mark.parametrize( + ("backend", "speculative", "uses_b12x"), + [ + ("auto", False, False), + ("auto", True, True), + ("b12x", False, True), + ("b12x", True, True), + ("triton", False, False), + ("triton", True, False), + ], +) +def test_glm5next_selects_configured_kda_decode_backend( + monkeypatch, + backend: str, + speculative: bool, + uses_b12x: bool, +) -> None: + monkeypatch.setattr( + KimiGatedDeltaNetAttention, + "__init__", + lambda self, config, vllm_config, prefix: None, + ) + vllm_config = SimpleNamespace( + additional_config={"glm53_kda_decode_backend": backend}, + speculative_config=object() if speculative else None, + ) + + layer = Glm5NextLinearAttention(object(), vllm_config) + + assert layer.enable_b12x_kda_decode is uses_b12x + + +def test_glm5next_defaults_to_hybrid_kda_decode(monkeypatch) -> None: + monkeypatch.setattr( + KimiGatedDeltaNetAttention, + "__init__", + lambda self, config, vllm_config, prefix: None, + ) + vllm_config = SimpleNamespace(additional_config={}, speculative_config=None) + + layer = Glm5NextLinearAttention(object(), vllm_config) + + assert layer._glm53_kda_decode_backend == "auto" + assert not layer.enable_b12x_kda_decode + + +def test_glm5next_rejects_unknown_kda_decode_backend() -> None: + vllm_config = SimpleNamespace( + additional_config={"glm53_kda_decode_backend": "unknown"} + ) + + with pytest.raises(ValueError, match="KDA decode backend"): + Glm5NextLinearAttention(object(), vllm_config) + + def test_glm5next_b12x_mhc_builds_first_layer_broadcast_fn() -> None: hidden_size = 4 hc_mult = 4 @@ -500,6 +555,122 @@ def plan(caps): assert captured_caps["null_state_index"] == 0 +def test_glm5next_b12x_kda_binds_spec_tensors_without_staging(monkeypatch) -> None: + scratch = torch.empty(32, dtype=torch.uint8) + bindings: list[dict[str, object]] = [] + runs: list[object] = [] + + class FakePlan: + @staticmethod + def shapes_and_dtypes(): + return (((32,), torch.uint8),) + + class FakeApi: + @staticmethod + def bind_kda(plan, **kwargs): + binding = {"plan": plan, **kwargs} + bindings.append(binding) + return binding + + @staticmethod + def run_kda(binding, **kwargs): + runs.append(binding) + + workspace = SimpleNamespace(get_simultaneous=lambda *specs: (scratch,)) + monkeypatch.setattr( + kimi_gdn_linear_attn, + "current_workspace_manager", + lambda: workspace, + ) + + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + torch.nn.Module.__init__(layer) + layer._b12x_kda_api = FakeApi() + layer._b12x_kda_plan = FakePlan() + layer._b12x_kda_max_tokens = 4 + layer._b12x_kda_max_seqs = 2 + layer._b12x_kda_state_index_columns = 2 + layer.local_num_heads = 1 + layer.head_dim = 2 + layer.gate_lower_bound = -5.0 + layer.A_log = torch.ones(1) + layer.dt_bias = torch.ones(2) + layer.o_norm = SimpleNamespace(weight=torch.ones(2), eps=1e-6) + layer.kv_cache = (torch.empty(0), torch.zeros(3, 1, 2, 2)) + layer._b12x_kda_mixed_qkv = torch.empty(4, 6) + layer._b12x_kda_raw_g = torch.empty(4, 1, 2) + layer._b12x_kda_raw_beta = torch.empty(4, 1) + layer._b12x_kda_z = torch.empty(4, 1, 2) + layer._b12x_kda_output = torch.zeros(4, 1, 2) + layer._b12x_kda_query_start_loc = torch.zeros(3, dtype=torch.int32) + layer._b12x_kda_num_accepted_tokens = torch.ones(2, dtype=torch.int32) + layer._b12x_kda_state_indices = torch.zeros(2, 2, 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) + + kwargs = dict( + mixed_qkv=torch.ones(1, 6), + raw_g=torch.ones(1, 1, 2), + raw_beta=torch.ones(1, 1), + z=torch.ones(1, 1, 2), + output=torch.empty(1, 1, 2), + state_indices=torch.tensor([[1]], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + num_accepted_tokens=None, + num_requests=1, + ) + layer._run_b12x_kda_decode_post_conv(**kwargs) + spec_kwargs = dict( + mixed_qkv=torch.ones(4, 6), + raw_g=torch.ones(4, 1, 2), + raw_beta=torch.ones(4, 1), + z=torch.ones(4, 1, 2), + output=torch.empty(4, 1, 2), + state_indices=torch.tensor([[1, 2], [2, 1]], dtype=torch.int32), + query_start_loc=torch.tensor([0, 2, 4], dtype=torch.int32), + num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32), + num_requests=2, + ) + layer._run_b12x_kda_decode_post_conv(**spec_kwargs) + + assert len(bindings) == 2 + assert bindings[0] is not bindings[1] + assert bindings[0]["scratch"] is scratch + assert bindings[1]["scratch"] is scratch + assert bindings[0]["mixed_qkv"] is layer._b12x_kda_mixed_qkv + assert bindings[0]["raw_g"] is layer._b12x_kda_raw_g + assert bindings[0]["raw_beta"] is layer._b12x_kda_raw_beta + assert bindings[0]["z"] is layer._b12x_kda_z + assert bindings[0]["output"] is layer._b12x_kda_output + assert bindings[0]["query_start_loc"] is layer._b12x_kda_query_start_loc + assert bindings[0]["state_indices"] is layer._b12x_kda_state_indices + assert bindings[0]["num_tokens"] is layer._b12x_kda_num_tokens + torch.testing.assert_close(bindings[0]["mixed_qkv"][0], kwargs["mixed_qkv"][0]) + assert bindings[1]["mixed_qkv"] is spec_kwargs["mixed_qkv"] + assert ( + bindings[1]["num_accepted_tokens"].data_ptr() + == spec_kwargs["num_accepted_tokens"].data_ptr() + ) + assert ( + bindings[1]["query_start_loc"].data_ptr() + == spec_kwargs["query_start_loc"].data_ptr() + ) + assert ( + bindings[1]["state_indices"].data_ptr() + == spec_kwargs["state_indices"].data_ptr() + ) + assert bindings[1]["num_tokens"] is layer._b12x_kda_num_tokens + assert bindings[1]["num_tokens"].item() == spec_kwargs["mixed_qkv"].shape[0] + assert ( + bindings[1]["num_tokens"].data_ptr() + != spec_kwargs["query_start_loc"][2:].data_ptr() + ) + assert bindings[1]["output"].data_ptr() == spec_kwargs["output"].data_ptr() + assert len(runs) == 2 + assert runs[0] is bindings[0] + assert runs[1] is bindings[1] + + def test_glm5next_sparse_mla_selects_b12x_backend(monkeypatch) -> None: captured: dict[str, object] = {} 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 a6760cb856e7..fb43fc68d6e5 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 @@ -26,6 +26,7 @@ from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.utils.b12x import get_b12x_gdn_decode from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.worker.workspace import current_workspace_manager from ...linear import ( ColumnParallelLinear, @@ -308,7 +309,6 @@ def __init__( self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") self._b12x_kda_api: Any | None = None self._b12x_kda_plan = None - self._b12x_kda_binding = None self._initialize_b12x_kda_decode(vllm_config) self.o_proj = RowParallelLinear( self.projection_size, @@ -412,12 +412,6 @@ def _initialize_b12x_kda_decode(self, vllm_config: VllmConfig) -> None: torch.zeros(1, dtype=torch.int32, device=device), persistent=False, ) - scratch_spec = provisional.scratch_specs()[0] - self.register_buffer( - "_b12x_kda_scratch", - torch.empty(scratch_spec.shape, dtype=scratch_spec.dtype, device=device), - persistent=False, - ) def _make_b12x_kda_plan(self, max_state_slots: int): api = self._b12x_kda_api @@ -450,27 +444,8 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: recurrent_state = self.kv_cache[1] plan = self._make_b12x_kda_plan(max_state_slots=recurrent_state.shape[0]) self._b12x_kda_plan = plan - self._b12x_kda_binding = api.bind_kda( - plan, - scratch=self._b12x_kda_scratch, - 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, - 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=recurrent_state, - query_start_loc=self._b12x_kda_query_start_loc, - num_accepted_tokens=self._b12x_kda_num_accepted_tokens, - 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, - ) def unbind_kv_cache(self) -> None: - self._b12x_kda_binding = None self._b12x_kda_plan = None super().unbind_kv_cache() @@ -486,7 +461,7 @@ def rearrange_mixed_qkv( def _can_use_b12x_kda_decode(self, m: GDNAttentionMetadata) -> bool: if ( - self._b12x_kda_binding is None + self._b12x_kda_plan is None or m.num_prefills != 0 or (m.num_decodes == 0 and m.num_spec_decodes == 0) ): @@ -518,9 +493,9 @@ def _run_b12x_kda_decode_post_conv( num_accepted_tokens: torch.Tensor | None, num_requests: int, ) -> None: - binding = self._b12x_kda_binding + plan = self._b12x_kda_plan api = self._b12x_kda_api - if binding is None or api is None: + if plan is None or api is None: raise RuntimeError("b12x KDA KV cache was not bound before inference") num_tokens = int(mixed_qkv.shape[0]) state_columns = int(state_indices.shape[1]) @@ -528,6 +503,7 @@ def _run_b12x_kda_decode_post_conv( num_tokens > self._b12x_kda_max_tokens or num_requests > self._b12x_kda_max_seqs or state_columns > self._b12x_kda_state_index_columns + or num_tokens > num_requests * state_columns ): raise ValueError( "b12x KDA capacity exceeded: " @@ -537,6 +513,49 @@ def _run_b12x_kda_decode_post_conv( f"{self._b12x_kda_state_index_columns}" ) + scratch_buffers = current_workspace_manager().get_simultaneous( + *plan.shapes_and_dtypes() + ) + if not scratch_buffers: + raise RuntimeError("b12x KDA plan did not expose caller scratch") + scratch: torch.Tensor | tuple[torch.Tensor, ...] + scratch = ( + scratch_buffers[0] if len(scratch_buffers) == 1 else tuple(scratch_buffers) + ) + + if ( + num_accepted_tokens is not None + and state_columns == self._b12x_kda_state_index_columns + ): + self._b12x_kda_num_seqs.fill_(num_requests) + self._b12x_kda_num_tokens.fill_(num_tokens) + query_start_loc = query_start_loc[: num_requests + 1] + 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=num_accepted_tokens[:num_requests], + state_indices=state_indices[:num_requests, :state_columns], + num_seqs=self._b12x_kda_num_seqs, + num_tokens=self._b12x_kda_num_tokens, + output=output[:num_tokens], + ) + api.run_kda( + binding, + lower_bound=self.gate_lower_bound, + eps=self.o_norm.eps, + scale=self.head_dim**-0.5, + ) + return + self._b12x_kda_mixed_qkv[:num_tokens].copy_(mixed_qkv) self._b12x_kda_raw_g[:num_tokens].copy_(raw_g) self._b12x_kda_raw_beta[:num_tokens].copy_(raw_beta) @@ -555,10 +574,26 @@ def _run_b12x_kda_decode_post_conv( state_indices[:num_requests] ) self._b12x_kda_num_seqs.fill_(num_requests) - self._b12x_kda_num_tokens.copy_( - query_start_loc[num_requests : num_requests + 1] - ) + self._b12x_kda_num_tokens.fill_(num_tokens) + binding = api.bind_kda( + plan, + scratch=scratch, + 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, + 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=self._b12x_kda_query_start_loc, + num_accepted_tokens=self._b12x_kda_num_accepted_tokens, + 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, + ) api.run_kda( binding, lower_bound=self.gate_lower_bound, diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 47f2823fc6c2..deb4c43527d2 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -2,8 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GLM-5.3 KDA modeling adapter.""" +from typing import Any + import torch +from vllm.config import VllmConfig from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( KimiGatedDeltaNetAttention, ) @@ -15,6 +18,33 @@ class Glm5NextLinearAttention(KimiGatedDeltaNetAttention): enable_b12x_kda_decode = True b12x_kda_null_state_index = 0 + def __init__( + self, + config: Any, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("glm53_kda_decode_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + if backend not in ("auto", "b12x", "triton"): + raise ValueError( + "GLM-5.3 KDA decode backend must be 'auto', 'b12x', or " + "'triton', " + f"got {backend!r}." + ) + self._glm53_kda_decode_backend = backend + # A recurrent cache uses one KDA implementation for its full lifetime. + # Keeping its BF16 rounding consistent prevents speculative acceptance + # from depending on whether a step entered through plain or spec decode. + self.enable_b12x_kda_decode = backend == "b12x" or ( + backend == "auto" and vllm_config.speculative_config is not None + ) + super().__init__(config, vllm_config, prefix) + def forward( # type: ignore[override] self, hidden_states: torch.Tensor,