[Feature][Model] Integrate AscendC KDA and causal convolution for GLM-5.3-Flash

tests/e2e/nightly/single_node/ops/singlecard_ops/test_glm5next_conv_state.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch_npu  # noqa: F401
+from vllm.triton_utils import triton
+
+from vllm_ascend.models.glm5next.ops.causal_conv1d import _copy_conv_state
+
+
+@pytest.mark.parametrize("dim_first", [False, True])
+def test_conv_state_copy_masks_invalid_slots_and_preserves_shared_pages(dim_first):
+    slots, width, dim = 4, 6, 384
+    # Leave a gap after each physical state and preserve the alternate layout.
+    backing = torch.arange(slots * width * dim * 2, dtype=torch.float32, device="npu")
+    strides = (width * dim * 2, 1, width) if dim_first else (width * dim * 2, dim, 1)
+    cache = backing.as_strided((slots, width, dim), strides)
+    saved = backing.clone()
+    indices = torch.tensor([2, -1, slots, 1, 3], dtype=torch.int32, device="npu")
+    starts = torch.tensor([0, 1, 2, 3, 3, 4], dtype=torch.int32, device="npu")
+    packed = torch.empty(5, width, dim, device="npu")
+    packed_indices = torch.empty_like(indices)
+    args = (cache, packed, indices, starts, packed_indices, cache.stride(0), 1, slots, width, dim, *cache.stride()[1:])
+    grid = (5, triton.cdiv(width * dim, 256))
+    _copy_conv_state[grid](*args, WRITE_BACK=False, BLOCK=256)
+    torch.testing.assert_close(packed_indices.cpu(), torch.tensor([0, -1, -1, -1, 4], dtype=torch.int32))
+    torch.testing.assert_close(packed[0], cache[2])
+    torch.testing.assert_close(packed[4], cache[3])
+    assert torch.count_nonzero(packed[1:4]).item() == 0
+    packed.fill_(17)
+    _copy_conv_state[grid](*args, WRITE_BACK=True, BLOCK=256)
+    expected = saved.as_strided(cache.shape, strides)
+    expected[2].fill_(17)
+    expected[3].fill_(17)
+    torch.testing.assert_close(backing, saved)

tests/ut/models/test_glm5next_kda_contracts.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_ascend.models.glm5next.ops.kda as kda
+import vllm_ascend.ops.kda as kda_ops
+
+
+@pytest.mark.parametrize("accepted", [None, [1, 2, 1]])
+@pytest.mark.parametrize("qkv_padding", [0, 64])
+def test_recurrent_raw_gates_rollback_slots_and_padding(monkeypatch, accepted, qkv_padding):
+    q, k, v = (torch.ones(1, 4, 1, 128 + qkv_padding * i, dtype=torch.bfloat16)[..., :128] for i in (1, 2, 3))
+    gate = q * 2
+    beta = torch.zeros(1, 4, 1, dtype=torch.bfloat16)
+    state = torch.zeros(8, 1, 128, 128)
+    starts = torch.tensor([0, 1, 3], dtype=torch.int32)
+    slots = torch.tensor([[2, 3], [5, 6], [0, 0]], dtype=torch.int32)
+    accepted_tensor = None if accepted is None else torch.tensor(accepted, dtype=torch.int32)
+
+    def recurrent(q_arg, k_arg, v_arg, gate_arg, beta_arg, state_arg, cu, ids, a_log, bias, **kwargs):
+        # Preserve each view's strides and storage without materializing Q/K/V.
+        assert q_arg is q and k_arg is k and v_arg is v
+        assert kwargs["use_gate_in_kernel"] and kwargs["use_beta_sigmoid_in_kernel"]
+        assert kwargs["safe_gate"] and kwargs["lower_bound"] == -4.0
+        assert state_arg is state
+        torch.testing.assert_close(gate_arg, gate)
+        torch.testing.assert_close(beta_arg, beta)
+        torch.testing.assert_close(ids, slots[:2])
+        if accepted_tensor is not None:
+            torch.testing.assert_close(kwargs["num_accepted_tokens"], accepted_tensor[:2])
+        result = q_arg.clone()
+        result[:, 3:] = float("nan")
+        return result
+
+    monkeypatch.setattr(torch.ops._C_ascend, "recurrent_kda", recurrent, raising=False)
+    out = kda.recurrent_kda(
+        q, k, v, gate, beta, state, starts, slots, torch.zeros(1), torch.zeros(128), -4, accepted_tensor
+    )
+    torch.testing.assert_close(out[:, :3], q[:, :3])
+    assert torch.count_nonzero(out[:, 3:]) == 0
+
+
+@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("compact", [False, True])
+def test_chunk_uses_host_descriptors_and_preserves_vk_cache(monkeypatch, state_dtype, compact):
+    q = torch.ones(1, 3, 1, 128, dtype=torch.bfloat16)
+    state = torch.arange(4 * 128 * 128, dtype=torch.float32).reshape(4, 1, 128, 128).to(state_dtype)
+    saved = state.clone()
+    indices = torch.tensor([3, 1], dtype=torch.int32)
+    has_initial = torch.tensor([True, False])
+    keep = torch.tensor([0]) if compact else None
+    metadata = SimpleNamespace(
+        keep_meta=keep,
+        cu_seqlens_host=torch.tensor([0, 3] if compact else [0, 1, 3]),
+        cu_seqlens_kern=None,
+        chunk_indices_chunk64_host=torch.tensor([[0, 0]] if compact else [[0, 0], [1, 0]]),
+    )
+
+    def chunk(q_arg, k_arg, v, g, beta, scale, chunk_size, **kwargs):
+        assert kwargs["state_v_first"] and kwargs["use_gate_in_kernel"] and kwargs["safe_gate"]
+        assert kwargs["cu_seqlens"] is metadata.cu_seqlens_host
+        assert kwargs["chunk_indices"] is metadata.chunk_indices_chunk64_host
+        assert kwargs["initial_state"].dtype == torch.float32
+        torch.testing.assert_close(kwargs["initial_state"][0], saved[3].float())
+        if not compact:
+            assert torch.count_nonzero(kwargs["initial_state"][1]) == 0
+        torch.testing.assert_close(beta, torch.full_like(beta, 0.5))
+        return v, torch.full_like(kwargs["initial_state"], 17)
+
+    monkeypatch.setattr(torch.ops._C_ascend, "chunk_kda_fwd", chunk, raising=False)
+    monkeypatch.setattr(kda_ops, "l2norm_fwd", lambda x: x)
+    out = kda.chunk_kda(
+        q, q, q, q, torch.zeros(1, 3, 1), state, indices, has_initial, metadata, torch.zeros(1), torch.zeros(128), -4
+    )
+    torch.testing.assert_close(out, q)
+    assert (state[3] == 17).all()
+    torch.testing.assert_close(state[0], saved[0])
+    torch.testing.assert_close(state[2], saved[2])
+    if compact:
+        torch.testing.assert_close(state[1], saved[1])
+    else:
+        assert (state[1] == 17).all()

tests/ut/models/test_glm5next_kda_dispatch.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_ascend.models.glm5next.kda as model_kda
+
+
+@pytest.mark.parametrize("speculative", [False, True])
+@pytest.mark.parametrize("dim_first", [False, True])
+def test_decode_and_prefill_use_their_own_metadata_and_merge_outputs(monkeypatch, speculative, dim_first):
+    layer = model_kda.Glm5NextLinearAttention.__new__(model_kda.Glm5NextLinearAttention)
+    torch.nn.Module.__init__(layer)
+    layer.prefix = "layer"
+    layer.local_num_heads = 1
+    layer.head_dim = layer.local_projection_size = 128
+    layer.kda_lower_bound = -4.0
+    state_shape = (8, 384, 6) if dim_first else (8, 6, 384)
+    layer.kv_cache = (torch.zeros(state_shape, dtype=torch.bfloat16), torch.zeros(8, 1, 128, 128))
+    layer._conv_state_dim_first = dim_first
+    layer._merged_conv_weight = None
+    # Checkpoints retain FP32 convolution weights; AscendC needs the activation
+    # dtype and [width, q|k|v] ordering in the cached packed weight.
+    for index, name in enumerate(("q_conv1d", "k_conv1d", "v_conv1d"), start=1):
+        setattr(layer, name, SimpleNamespace(bias=None, weight=torch.full((128, 1, 4), float(index))))
+    layer.A_log = torch.zeros(1)
+    layer.dt_bias = torch.zeros(128)
+    tokens = 5 if speculative else 4
+    metadata = object.__new__(model_kda.GDNAttentionMetadata)
+    values = dict(
+        has_initial_state=torch.tensor([True, False]),
+        non_spec_query_start_loc=torch.tensor([0, 1, 4]),
+        non_spec_state_indices_tensor=torch.tensor([2, 5]),
+        num_actual_tokens=tokens,
+        spec_sequence_masks=None,
+        spec_query_start_loc=None,
+        spec_state_indices_tensor=None,
+        spec_token_indx=None,
+        non_spec_token_indx=None,
+        num_accepted_tokens=None,
+        num_spec_decodes=0,
+        num_prefills=1,
+        num_decodes=1,
+        num_decode_tokens=1,
+        prefill_state_indices=torch.tensor([5]),
+        prefill_has_initial_state=torch.tensor([False]),
+        non_spec_prefill_metadata=SimpleNamespace(chunk=object()),
+    )
+    if speculative:
+        values.update(
+            has_initial_state=torch.tensor([False]),
+            non_spec_query_start_loc=torch.tensor([0, 3]),
+            non_spec_state_indices_tensor=torch.tensor([5]),
+            spec_sequence_masks=torch.tensor([True, False]),
+            spec_query_start_loc=torch.tensor([0, 2]),
+            spec_state_indices_tensor=torch.tensor([[2, 3]]),
+            spec_token_indx=torch.tensor([0, 1]),
+            non_spec_token_indx=torch.tensor([2, 3, 4]),
+            num_accepted_tokens=torch.tensor([1]),
+            num_spec_decodes=1,
+            num_decodes=0,
+            num_decode_tokens=0,
+        )
+    for name, value in values.items():
+        setattr(metadata, name, value)
+    prefill_conv = SimpleNamespace(
+        query_start_loc=metadata.non_spec_query_start_loc,
+        cache_indices=metadata.non_spec_state_indices_tensor[:, None],
+        initial_state_mode=metadata.has_initial_state,
+    )
+    metadata.non_spec_prefill_metadata.causal_conv1d = prefill_conv
+    if speculative:
+        metadata.spec_decode_metadata = SimpleNamespace(
+            spec_causal_conv1d=SimpleNamespace(
+                query_start_loc=metadata.spec_query_start_loc,
+                cache_indices=metadata.spec_state_indices_tensor,
+                num_accepted_tokens=metadata.num_accepted_tokens,
+            )
+        )
+    monkeypatch.setattr(model_kda, "get_forward_context", lambda: SimpleNamespace(attn_metadata={"layer": metadata}))
+    conv_calls = []
+    conv_entry = model_kda.causal_conv1d
+
+    def cpu_conv_entry(x, weight, state, *args, **kwargs):
+        # This UT checks model dispatch and layout; the NPU tests exercise the
+        # non-contiguous state's device-side gather/scatter and original alias.
+        assert state.data_ptr() == layer.kv_cache[0].data_ptr()
+        assert state.shape == (8, 6, 384)
+        return conv_entry(x, weight, state.contiguous(), *args, **kwargs)
+
+    monkeypatch.setattr(model_kda, "causal_conv1d", cpu_conv_entry)
+
+    def conv(output, x, weight, **kwargs):
+        conv_calls.append(kwargs["run_mode"])
+        state = kwargs["conv_state"]
+        assert state.shape == (8, 6, 384)
+        assert weight.shape == (4, 384)
+        assert weight.dtype == x.dtype == state.dtype == torch.bfloat16
+        expected_weight = torch.arange(1, 4, dtype=torch.bfloat16).repeat_interleave(128).expand(4, 384)
+        torch.testing.assert_close(weight, expected_weight)
+        assert weight is layer._merged_conv_weight
+        assert torch.count_nonzero(output) == 0
+        if kwargs["run_mode"] == 0:
+            assert kwargs["cache_indices_opt"] is prefill_conv.cache_indices
+            assert kwargs["initial_state_mode_opt"] is prefill_conv.initial_state_mode
+            assert kwargs["num_accepted_tokens_opt"] is None
+        else:
+            spec_conv = metadata.spec_decode_metadata.spec_causal_conv1d
+            assert kwargs["cache_indices_opt"] is spec_conv.cache_indices
+            assert kwargs["num_accepted_tokens_opt"] is spec_conv.num_accepted_tokens
+            assert kwargs["initial_state_mode_opt"] is None
+        # A distinct return catches accidentally ignoring the declared op result.
+        return x.clone()
+
+    monkeypatch.setattr(torch.ops._C_ascend, "npu_causal_conv1d_custom", conv, raising=False)
+    calls = []
+
+    def recurrent(q, k, v, gate, beta, state, starts, indices, *args):
+        calls.append("recurrent")
+        assert q.shape[1] == (2 if speculative else 1)
+        torch.testing.assert_close(starts, torch.tensor([0, q.shape[1]]))
+        return q * 2
+
+    def prefill(q, k, v, gate, beta, state, indices, initial, chunk, *args):
+        calls.append("prefill")
+        assert q.shape[1] == 3
+        assert chunk is metadata.non_spec_prefill_metadata.chunk
+        torch.testing.assert_close(indices, metadata.prefill_state_indices)
+        return q * 3
+
+    monkeypatch.setattr(model_kda, "recurrent_kda", recurrent)
+    monkeypatch.setattr(model_kda, "chunk_kda", prefill)
+    qkv = torch.arange(1, tokens + 1, dtype=torch.bfloat16)[:, None].expand(tokens, 384).clone()
+    out = torch.full((1, tokens + 1, 1, 128), float("nan"), dtype=torch.bfloat16)
+    layer._forward(qkv, torch.zeros(1, tokens, 1, 128), torch.zeros(1, tokens, 1), out)
+    expected = torch.arange(1, tokens + 1, dtype=torch.bfloat16) * 3
+    expected[: 2 if speculative else 1] = torch.arange(1, 3 if speculative else 2) * 2
+    torch.testing.assert_close(out[0, :tokens, 0, 0], expected)
+    assert torch.count_nonzero(out[:, tokens:]) == 0
+    assert calls == ["recurrent", "prefill"]
+    assert conv_calls == ([1, 0] if speculative else [0])
+
+
+@pytest.mark.parametrize(("width", "num_spec"), [(1, 0), (5, 0), (3, 3)])
+def test_unsupported_conv_width_is_rejected_before_execution(monkeypatch, width, num_spec):
+    monkeypatch.setattr(
+        model_kda.GatedDeltaNetAttention, "__init__", lambda self, *a, **kw: torch.nn.Module.__init__(self)
+    )
+    config = SimpleNamespace(linear_head_dim=128, linear_conv_kernel_dim=width)
+    vllm_config = SimpleNamespace(
+        quant_config=None,
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        speculative_config=SimpleNamespace(num_speculative_tokens=num_spec),
+    )
+    with pytest.raises(ValueError, match="causal-conv requires"):
+        model_kda.Glm5NextLinearAttention(config, vllm_config)

tests/ut/ops/test_kimi_kda.py

@@ -368,11 +368,12 @@ def test_fused_qkv_keeps_non_mxfp_quantization_in_linear_apply():
     adapter.apply.assert_called_once_with(attention.in_proj_qkvgfab, hidden_states, bias=None)
 
 
-def test_prefill_fuses_raw_gate_and_updates_v_first_state():
+@pytest.mark.parametrize("lower_bound", [None, -4.0])
+def test_prefill_fuses_raw_gate_and_updates_v_first_state(lower_bound):
     attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention)
     nn.Module.__init__(attention)
     attention.head_dim = 2
-    attention.gate_lower_bound = None
+    attention.gate_lower_bound = lower_bound
     attention.A_log = nn.Parameter(torch.randn(1))
     attention.dt_bias = nn.Parameter(torch.randn(2))
 
@@ -395,7 +396,7 @@ def test_prefill_fuses_raw_gate_and_updates_v_first_state():
 
     with (
         patch("vllm_ascend.ops.kimi_kda.clear_ssm_states"),
-        patch("vllm_ascend.ops.kimi_kda.l2norm_fwd", side_effect=lambda x: x),
+        patch("vllm_ascend.ops.kda.l2norm_fwd", side_effect=lambda x: x),
         patch.object(
             torch.ops._C_ascend,
             "chunk_kda_fwd",
@@ -419,10 +420,40 @@ def test_prefill_fuses_raw_gate_and_updates_v_first_state():
     assert chunk_kda_fwd.call_args.args[3] is raw_gate
     assert chunk_kda_fwd.call_args.kwargs["use_gate_in_kernel"] is True
     assert chunk_kda_fwd.call_args.kwargs["state_v_first"] is True
-    assert chunk_kda_fwd.call_args.kwargs["safe_gate"] is False
+    assert chunk_kda_fwd.call_args.args[4] is beta
+    assert chunk_kda_fwd.call_args.kwargs["safe_gate"] is (lower_bound is not None)
+    assert chunk_kda_fwd.call_args.kwargs["lower_bound"] == (lower_bound if lower_bound is not None else -5.0)
     torch.testing.assert_close(recurrent_state[state_indices], final_state)
 
 
+@pytest.mark.parametrize("lower_bound", [None, -4.0])
+@pytest.mark.parametrize("qkv_padding", [0, 64])
+def test_recurrent_preserves_preprocessed_beta_and_mtp_state(lower_bound, qkv_padding):
+    attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention)
+    nn.Module.__init__(attention)
+    attention.gate_lower_bound = lower_bound
+    attention.A_log, attention.dt_bias = torch.zeros(1), torch.zeros(128)
+    q, k, v = (torch.ones(1, 4, 1, 128 + qkv_padding * i, dtype=torch.bfloat16)[..., :128] for i in (1, 2, 3))
+    beta = torch.full((1, 4, 1), 0.25)
+    state = torch.zeros(8, 1, 128, 128)
+    starts = torch.tensor([0, 4], dtype=torch.int32)
+    slots = torch.tensor([[2, 3, 4, 5]], dtype=torch.int32)
+    accepted = torch.tensor([2], dtype=torch.int32)
+    with patch.object(torch.ops._C_ascend, "recurrent_kda", return_value=q, create=True) as recurrent:
+        output = attention._run_recurrent(q, k, v, q, beta, state, starts, slots, num_accepted_tokens=accepted)
+    assert output is q
+    assert recurrent.call_args.args[0] is q
+    assert recurrent.call_args.args[1] is k
+    assert recurrent.call_args.args[2] is v
+    assert recurrent.call_args.args[4] is beta
+    assert recurrent.call_args.args[5] is state
+    assert recurrent.call_args.args[7] is slots
+    assert recurrent.call_args.kwargs["num_accepted_tokens"] is accepted
+    assert recurrent.call_args.kwargs["use_beta_sigmoid_in_kernel"] is False
+    assert recurrent.call_args.kwargs["safe_gate"] is (lower_bound is not None)
+    assert recurrent.call_args.kwargs["lower_bound"] == (lower_bound if lower_bound is not None else -5.0)
+
+
 def test_kda_empty_forward_context_clears_preallocated_output():
     attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention)
     core_attn_out = torch.full((1, 4, 2, 3), torch.nan)