fix(kda): preserve gate alias contract during GLM decode

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

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,