diff --git a/tests/e2e/nightly/single_node/ops/singlecard_ops/test_glm5next_conv_state.py b/tests/e2e/nightly/single_node/ops/singlecard_ops/test_glm5next_conv_state.py new file mode 100644 index 000000000000..0b9f403b4358 --- /dev/null +++ b/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) diff --git a/tests/ut/models/test_glm5next_kda_contracts.py b/tests/ut/models/test_glm5next_kda_contracts.py new file mode 100644 index 000000000000..123fb5ea112c --- /dev/null +++ b/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() diff --git a/tests/ut/models/test_glm5next_kda_dispatch.py b/tests/ut/models/test_glm5next_kda_dispatch.py new file mode 100644 index 000000000000..2780f927da1d --- /dev/null +++ b/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) diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py index eb20e8061fb4..1c9624be1e58 100644 --- a/tests/ut/ops/test_kimi_kda.py +++ b/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) diff --git a/tests/ut/worker/test_glm5next_pooled_cache.py b/tests/ut/worker/test_glm5next_pooled_cache.py index 54a697182524..f3c583cd842a 100644 --- a/tests/ut/worker/test_glm5next_pooled_cache.py +++ b/tests/ut/worker/test_glm5next_pooled_cache.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """CPU tests for GLM-Next model-runner pooled cache views.""" +from itertools import permutations from types import SimpleNamespace from unittest.mock import patch +import pytest import torch from vllm.v1.core.single_type_kv_cache_manager import ( register_all_kvcache_specs, @@ -14,6 +16,7 @@ from vllm_ascend.core.kv_cache_interface import ( AscendIndexerKPoolStateSpec, AscendMLAAttentionSpec, + requires_padded_page_layout, ) from vllm_ascend.models.glm5next.cache_config import ( get_glm5_next_kv_cache_config, @@ -206,9 +209,10 @@ def test_glm5_next_runner_allocates_contiguous_slot_backings(): page_size = descriptors[name].size // plan.num_blocks assert cache.stride(0) * cache.element_size() == page_size assert cache.data_ptr() == raw_caches[name].data_ptr() - assert all(cache.is_contiguous() for cache in caches[MAMBA]) + for cache in caches[MAMBA]: + assert cache.stride(0) * cache.element_size() == descriptors[MAMBA].size // plan.num_blocks - mamba_second_offset = caches[MAMBA][0].numel() * caches[MAMBA][0].element_size() + mamba_second_offset = caches[MAMBA][0][0].numel() * caches[MAMBA][0].element_size() assert caches[MAMBA][1].data_ptr() - raw_caches[MAMBA].data_ptr() == mamba_second_offset mamba_payload_size = sum(cache.numel() * cache.element_size() for cache in caches[MAMBA]) assert mamba_payload_size < descriptors[MAMBA].size @@ -267,3 +271,72 @@ def test_glm5_next_initialize_passes_all_pooled_views_to_cache_binding(): runner.kv_caches, 1, ) + + +@pytest.mark.parametrize("offset", [0, 64]) +def test_glm_shared_mla_mamba_pages_preserve_other_block_ids(offset): + config, _, plan = _make_plan(num_blocks=4) + runner = _make_runner(config) + raw_tensors = runner._allocate_kv_cache_tensors(plan) + assert raw_tensors[MAIN] is raw_tensors[MAMBA] + page_bytes = raw_tensors[MAIN].numel() // plan.num_blocks + backing = torch.zeros(offset + raw_tensors[MAIN].numel() + 64, dtype=torch.int8) + raw = backing[offset : offset + plan.num_blocks * page_bytes] + raw_tensors[MAIN] = raw_tensors[MAMBA] = raw + caches = runner._reshape_kv_cache_tensors(plan, raw_tensors) + latent, _ = caches[MAIN] + conv, ssm = caches[MAMBA] + for state in (conv, ssm): + assert state.untyped_storage().data_ptr() == raw.untyped_storage().data_ptr() + assert state.stride(0) * state.element_size() == page_bytes + assert ssm.data_ptr() - conv.data_ptr() == conv[0].numel() * conv.element_size() + state_bytes = sum(state[0].numel() * state.element_size() for state in (conv, ssm)) + for mla_id, state_id in permutations(range(plan.num_blocks), 2): + raw.zero_() + latent[mla_id].fill_(7) + conv[state_id].fill_(11) + ssm[state_id].fill_(13) + torch.testing.assert_close(latent[mla_id], torch.full_like(latent[mla_id], 7)) + latent[mla_id].fill_(17) + torch.testing.assert_close(conv[state_id], torch.full_like(conv[state_id], 11)) + torch.testing.assert_close(ssm[state_id], torch.full_like(ssm[state_id], 13)) + assert torch.count_nonzero(raw.view(plan.num_blocks, page_bytes)[state_id, state_bytes:]) == 0 + assert torch.count_nonzero(backing[:offset]) == 0 + assert torch.count_nonzero(backing[offset + raw.numel() :]) == 0 + + +def test_padded_page_layout_detected_for_shared_state_pages(): + # The pooled layout pads the state caches to the page size of the + # block-stride addressed MLA/indexer caches they share physical pages with. + # Probe the specs the runner itself derives from the KV cache config. + config, _, plan = _make_plan() + layer_specs = _make_runner(config)._get_layer_kv_cache_specs(plan) + assert requires_padded_page_layout(layer_specs.values()) + + +def test_padded_page_layout_rejected_without_state_caches(): + # A standalone MTP runner has the same attention specs but no recurrent + # state caches, so no state view needs the padded page stride. + specs = {name: spec for name, spec in _make_specs().items() if not isinstance(spec, MambaSpec)} + assert not requires_padded_page_layout(specs.values()) + + +def test_padded_page_layout_rejected_for_packed_hybrid_pool(): + # Other hybrid models pad Mamba pages to the attention page size + # (``cache_config.mamba_page_size_padded``) but keep the packed contiguous + # state layout, so they must not take the shared padded page path. + specs = [ + AscendMLAAttentionSpec( + block_size=8, + num_kv_heads=1, + head_size=4, + dtype=torch.bfloat16, + ), + MambaSpec( + block_size=8, + shapes=((2, 2), (1, 2, 2)), + dtypes=(torch.bfloat16, torch.float32), + page_size_padded=64, + ), + ] + assert not requires_padded_page_layout(specs) diff --git a/vllm_ascend/core/kv_cache_interface.py b/vllm_ascend/core/kv_cache_interface.py index 8516e7ad4b87..501b9e784a3c 100644 --- a/vllm_ascend/core/kv_cache_interface.py +++ b/vllm_ascend/core/kv_cache_interface.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +from collections.abc import Iterable from dataclasses import dataclass, replace import torch @@ -13,6 +14,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheSpec, + MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, UniformTypeKVCacheSpecs, @@ -46,6 +48,30 @@ def get_storage_block_size(kv_cache_spec: KVCacheSpec) -> int: return getattr(kv_cache_spec, "storage_block_size", kv_cache_spec.block_size) +def requires_padded_page_layout(kv_cache_specs: Iterable[KVCacheSpec]) -> bool: + """Whether state caches advance by a padded page size shared with attention. + + Ascend mixed pools can bind one physical page to both an attention cache and a + recurrent state cache. The attention side of such a pool is addressed by an explicit + physical block stride (``indexes_kv_by_block_stride``) and the state caches are + padded to that same page size, so every view over the shared page has to advance by + the padded page size: otherwise a state update for one scheduler block ID would land + in the pages owned by another block ID. + + The capability belongs to the pool rather than to a single spec, because hybrid + models also pad Mamba pages through ``cache_config.mamba_page_size_padded`` while + keeping the packed contiguous state layout, so padding alone does not require this + layout. + """ + specs = list(kv_cache_specs) + state_specs = [spec for spec in specs if isinstance(spec, MambaSpec)] + if not state_specs: + return False + if not any(getattr(spec, "page_size_padded", None) is not None for spec in state_specs): + return False + return any(getattr(spec, "indexes_kv_by_block_stride", False) for spec in specs) + + @dataclass(frozen=True, kw_only=True) class AscendMLAAttentionSpec(MLAAttentionSpec): """MLA cache spec with Ascend-specific layout metadata. diff --git a/vllm_ascend/models/glm5next/kda.py b/vllm_ascend/models/glm5next/kda.py index b00d483a13d0..883f804bf6b7 100644 --- a/vllm_ascend/models/glm5next/kda.py +++ b/vllm_ascend/models/glm5next/kda.py @@ -20,11 +20,7 @@ is_conv_state_dim_first, ) from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader -from vllm.model_executor.utils import ( - maybe_disable_graph_partition, - set_weight_attrs, -) -from vllm.platforms import current_platform +from vllm.model_executor.utils import set_weight_attrs # FusedRMSNormGated is a CustomOp, so the Ascend implementation is picked up # through the OOT registration in `vllm_ascend.utils` rather than by import. @@ -32,18 +28,9 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm_ascend.models.glm5next.config import Glm5NextConfig -from vllm_ascend.models.glm5next.ops.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from vllm_ascend.models.glm5next.ops.state_ops import ( - gather_initial_states, - scatter_states, -) -from vllm_ascend.ops.triton.kda.kda import ( - chunk_kda_with_fused_gate, - fused_recurrent_kda, -) +from vllm_ascend.models.glm5next.ops.causal_conv1d import causal_conv1d +from vllm_ascend.models.glm5next.ops.kda import KDA_MAX_RECURRENT_TOKENS, chunk_kda, recurrent_kda +from vllm_ascend.ops.gdn_attn_builder import AscendGDNAttentionBackend class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear): @@ -113,16 +100,6 @@ def weight_loader_v2( param.tp_rank = param_tp_rank -@torch.compile( - dynamic=True, - backend=current_platform.simple_compile_backend, - options=maybe_disable_graph_partition(current_platform.simple_compile_backend), -) -def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: - """Fuse the fp32 cast + sigmoid into one Inductor kernel.""" - return x.float().sigmoid() - - class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -139,7 +116,7 @@ def get_state_shape( self, ) -> tuple[tuple[int, ...], tuple[int, ...]]: # conv_state width must include num_spec so the spec-decode conv update - # (causal_conv1d_update with num_accepted_tokens + max_query_len) can + # (AscendC causal-conv with num_accepted_tokens) can # slide the window across the draft-verify tokens without reading past # the allocated width. Matches qwen_gdn_linear_attn.get_state_shape. return MambaStateShapeCalculator.kda_state_shape( @@ -164,6 +141,17 @@ def __init__( finally: vllm_config.quant_config = saved_quant_config + if config.linear_head_dim != 128 or vllm_config.model_config.dtype != torch.bfloat16: + raise ValueError("GLM AscendC KDA requires BF16 activations and head_dim=128.") + num_spec = vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config else 0 + if num_spec + 1 > KDA_MAX_RECURRENT_TOKENS: + raise ValueError("GLM AscendC KDA supports at most seven speculative tokens.") + if not 2 <= config.linear_conv_kernel_dim <= 4: + raise ValueError("GLM AscendC causal-conv requires a kernel width in [2, 4].") + if num_spec and config.linear_conv_kernel_dim != 4: + raise ValueError("GLM AscendC causal-conv requires kernel width=4 for MTP.") + if not -5 <= config.linear_lower_bound < 0: + raise ValueError("GLM AscendC KDA requires linear_lower_bound in [-5, 0).") self.head_dim = config.linear_head_dim self.num_heads = config.linear_num_heads self.conv_size = config.linear_conv_kernel_dim @@ -270,12 +258,14 @@ def _a_log_weight_loader(param, loaded_weight): # GLM-5.3-Flash uses a bounded sigmoid gate instead of the default # unbounded softplus gate. - self.kda_safe_gate = True self.kda_lower_bound = config.linear_lower_bound # Process-global conv-state layout, resolved once here instead of on # every _forward call (it reads an env-derived flag each time). self._conv_state_dim_first = is_conv_state_dim_first() + def get_attn_backend(self): + return AscendGDNAttentionBackend + def forward( self, hidden_states: torch.Tensor, @@ -297,7 +287,7 @@ def forward( # Beta stays raw (bf16) here: the recurrent kernel sigmoids it in fp32 # at load (SIGMOID_BETA), and only the chunked prefill path needs the # pre-computed fp32 sigmoid — computed lazily in _forward. Pure decode - # / spec-verify steps then skip the _cast_sigmoid kernel and its fp32 + # / spec-verify steps then skip the separate sigmoid and its fp32 # intermediate entirely. beta = beta_raw.unsqueeze(0) g1 = self.f_b_proj(f_a)[0] @@ -312,8 +302,7 @@ def forward( dtype=hidden_states.dtype, device=hidden_states.device, ) - # Call the decorated eager break directly so host-side prefill branches - # are not captured by PIECEWISE CUDA graphs. + # Keep the layer's dispatch outside piecewise graph compilation. self._forward( qkv_proj_states=qkv, g1=g1, @@ -336,15 +325,16 @@ def _forward( attn_metadata_raw = forward_context.attn_metadata if attn_metadata_raw is None: + core_attn_out.zero_() return assert isinstance(attn_metadata_raw, dict) attn_metadata_narrowed = attn_metadata_raw.get(self.prefix) if attn_metadata_narrowed is None: # Profile/warmup dummy runs may omit mamba-family metadata. + core_attn_out.zero_() return assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) - has_initial_state = attn_metadata_narrowed.has_initial_state non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc non_spec_state_indices_tensor = attn_metadata_narrowed.non_spec_state_indices_tensor # noqa: E501 num_actual_tokens = attn_metadata_narrowed.num_actual_tokens @@ -358,7 +348,6 @@ def _forward( num_spec_decodes = attn_metadata_narrowed.num_spec_decodes use_spec = spec_sequence_masks is not None and num_spec_decodes > 0 # Safe-gate checkpoints use the bounded sigmoid variant. - safe_gate = self.kda_safe_gate lower_bound = self.kda_lower_bound constant_caches = self.kv_cache @@ -367,16 +356,15 @@ def _forward( beta = beta[:, :num_actual_tokens] (conv_state, recurrent_state) = constant_caches - # conv_state must be (..., dim, width-1) for the conv kernels. - # DS layout stores it that way directly; SD layout needs a transpose. + # AscendC consumes [cache, state_len, dim]. Preserve the original storage. # Layout is process-global and resolved once at init (see __init__). - if not self._conv_state_dim_first: + if self._conv_state_dim_first: conv_state = conv_state.transpose(-1, -2) # One merged short-conv over q|k|v instead of three separate calls. The # 1D conv is independent per channel, so concatenating q/k/v along the - # channel dim and running a single causal_conv1d is bit-identical to - # three calls. The merged weight is q|k|v conv weights concatenated; + # channel dim preserves the independent q/k/v convolutions. + # The merged weight is q|k|v conv weights concatenated; # built once and cached (params are fixed after load). conv_state is # already stored as the merged q|k|v state, so it is used directly. if self._merged_conv_weight is None: @@ -384,12 +372,16 @@ def _forward( def _w(m): return m.weight.view(m.weight.size(0), m.weight.size(2)) - self._merged_conv_weight = torch.cat( - [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)], - dim=0, - ).contiguous() + self._merged_conv_weight = ( + torch.cat( + [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)], + dim=0, + ) + .transpose(0, 1) + .to(dtype=qkv_proj_states.dtype) + .contiguous() + ) conv_weights = self._merged_conv_weight - conv_bias = self.q_conv1d.bias # Split projections / gating into spec (draft-verify) and non-spec token # groups when speculative decoding is active. Spec tokens carry @@ -426,18 +418,15 @@ def _w(m): if use_spec: assert spec_state_indices_tensor is not None assert num_accepted_tokens is not None - conv_idx = spec_state_indices_tensor[:, 0][:num_spec_decodes] - conv_mql = spec_state_indices_tensor.size(-1) - qkv_spec = causal_conv1d_update( + conv_meta = attn_metadata_narrowed.spec_decode_metadata.spec_causal_conv1d + qkv_spec = causal_conv1d( qkv_spec, - conv_state, conv_weights, - conv_bias, - activation="silu", - conv_state_indices=conv_idx, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=spec_query_start_loc, - max_query_len=conv_mql, + conv_state, + conv_meta.query_start_loc, + conv_meta.cache_indices, + run_mode=1, + num_accepted_tokens=conv_meta.num_accepted_tokens, ) q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1) @@ -445,151 +434,89 @@ def _w(m): q_ns = k_ns = v_ns = None if attn_metadata_narrowed.num_prefills > 0: assert qkv_ns is not None - qkv_ns = causal_conv1d_fn( - qkv_ns.transpose(0, 1), + conv_meta = attn_metadata_narrowed.non_spec_prefill_metadata.causal_conv1d + qkv_ns = causal_conv1d( + qkv_ns, conv_weights, - conv_bias, - activation="silu", - conv_states=conv_state, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) + conv_state, + conv_meta.query_start_loc, + conv_meta.cache_indices, + run_mode=0, + initial_state_mode=conv_meta.initial_state_mode, + ) q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1) elif attn_metadata_narrowed.num_decodes > 0: assert non_spec_state_indices_tensor is not None - decode_conv_indices = non_spec_state_indices_tensor[: attn_metadata_narrowed.num_decodes] - qkv_ns = causal_conv1d_update( + conv_meta = attn_metadata_narrowed.non_spec_decode_metadata.causal_conv1d + qkv_ns = causal_conv1d( qkv_ns, - conv_state, conv_weights, - conv_bias, - activation="silu", - conv_state_indices=decode_conv_indices, + conv_state, + conv_meta.query_start_loc, + conv_meta.cache_indices, + run_mode=1, ) q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1) - def _rearr(x): + def rearrange(x): return x.reshape(1, -1, self.local_num_heads, self.head_dim) - # --- core attention: spec (draft-verify) path --- - core_attn_out_spec = None - # In a pure spec-verify step (no non-spec tokens) the recurrent kernel - # can write straight into the layer output buffer, skipping the - # fresh allocation + copy below. Mixed steps must scatter via - # spec_token_indx, so they keep the kernel-managed output. - spec_out = ( - core_attn_out[0, :num_actual_tokens].unsqueeze(0) - if non_spec_token_indx is None or non_spec_token_indx.numel() == 0 - else None - ) + core_attn_out.zero_() if use_spec: - assert spec_state_indices_tensor is not None - assert num_accepted_tokens is not None - assert spec_query_start_loc is not None - # Gate computed inside the recurrent kernel (COMPUTE_GATE) from - # raw g1 — replicates fused_kda_gate's arithmetic bit-for-bit and - # skips its launch + fp32 [n, H, D] intermediate per layer. - core_attn_out_spec, _ = fused_recurrent_kda( - q=_rearr(q_spec), - k=_rearr(k_spec), - v=_rearr(v_spec), - g=g1_spec, - beta=beta_spec, - initial_state=recurrent_state, - use_qk_l2norm_in_kernel=True, - cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1], - ssm_state_indices=spec_state_indices_tensor, - num_accepted_tokens=num_accepted_tokens, - out=spec_out, - sigmoid_beta=True, - a_log=self.A_log, - g_bias=self.dt_bias, - compute_gate=True, - lower_bound=lower_bound, + spec_output = recurrent_kda( + rearrange(q_spec), + rearrange(k_spec), + rearrange(v_spec), + g1_spec, + beta_spec, + recurrent_state, + spec_query_start_loc[: num_spec_decodes + 1], + spec_state_indices_tensor, + self.A_log, + self.dt_bias, + lower_bound, + num_accepted_tokens, ) + core_attn_out[0].index_copy_(0, spec_token_indx, spec_output[0]) - # --- core attention: non-spec path (prefill or plain decode) --- - core_attn_out_non_spec = None - # Only the plain-decode recurrent kernel can write straight into the - # layer output buffer; the chunked prefill kernel cannot, so this - # stays None there and the merge copy below runs as before. - ns_out = None - if attn_metadata_narrowed.num_prefills > 0: - assert q_ns is not None - 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) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=_rearr(q_ns), - k=_rearr(k_ns), - v=_rearr(v_ns), - raw_g=g1_ns, - # Chunk path wants the pre-sigmoided fp32 beta (its kernels - # don't sigmoid); beta_ns is raw bf16 from forward. - beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), - A_log=self.A_log, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) - # Init cache - scatter_states( + if q_ns is None: + return + q_ns, k_ns, v_ns = rearrange(q_ns), rearrange(k_ns), rearrange(v_ns) + metadata = attn_metadata_narrowed + decode_tokens = metadata.num_decode_tokens if metadata.num_prefills > 0 else q_ns.shape[1] + output = None + if metadata.num_decodes > 0: + output = recurrent_kda( + q_ns[:, :decode_tokens], + k_ns[:, :decode_tokens], + v_ns[:, :decode_tokens], + g1_ns[:, :decode_tokens], + beta_ns[:, :decode_tokens], recurrent_state, - last_recurrent_state, + non_spec_query_start_loc[: metadata.num_decodes + 1], non_spec_state_indices_tensor, + self.A_log, + self.dt_bias, + lower_bound, ) - elif attn_metadata_narrowed.num_decodes > 0: - assert non_spec_query_start_loc is not None - assert non_spec_state_indices_tensor is not None - # Plain decode step (no spec tokens): token order is dense, so the - # kernel can write straight into the layer output buffer. A mixed - # step scatters non-spec output via non_spec_token_indx instead. - # Gate computed in-kernel (COMPUTE_GATE), beta sigmoided in-kernel. - if not use_spec: - ns_out = spec_out - core_attn_out_non_spec, _ = fused_recurrent_kda( - q=_rearr(q_ns), - k=_rearr(k_ns), - v=_rearr(v_ns), - g=g1_ns, - beta=beta_ns, - initial_state=recurrent_state, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc[: attn_metadata_narrowed.num_decodes + 1], - ssm_state_indices=non_spec_state_indices_tensor, - out=ns_out, - sigmoid_beta=True, - a_log=self.A_log, - g_bias=self.dt_bias, - compute_gate=True, - lower_bound=lower_bound, - ) - - # --- merge spec / non-spec outputs back into token order --- - if use_spec and core_attn_out_non_spec is not None: - assert core_attn_out_spec is not None - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, + if metadata.num_prefills > 0: + prefill_output = chunk_kda( + q_ns[:, decode_tokens:], + k_ns[:, decode_tokens:], + v_ns[:, decode_tokens:], + g1_ns[:, decode_tokens:], + beta_ns[:, decode_tokens:], + recurrent_state, + metadata.prefill_state_indices, + metadata.prefill_has_initial_state, + metadata.non_spec_prefill_metadata.chunk, + self.A_log, + self.dt_bias, + lower_bound, ) - 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.squeeze(0) - elif use_spec: - assert core_attn_out_spec is not None - if spec_out is None: - core_attn_out[0, :num_actual_tokens] = core_attn_out_spec.squeeze(0) + output = prefill_output if output is None else torch.cat((output, prefill_output), dim=1) + assert output is not None + if use_spec: + core_attn_out[0].index_copy_(0, non_spec_token_indx, output[0]) else: - assert core_attn_out_non_spec is not None - if ns_out is None: - core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[0, :num_actual_tokens] + core_attn_out[0, : output.shape[1]].copy_(output[0]) diff --git a/vllm_ascend/models/glm5next/ops/causal_conv1d.py b/vllm_ascend/models/glm5next/ops/causal_conv1d.py index a30634afa091..0a5d7f4b1173 100644 --- a/vllm_ascend/models/glm5next/ops/causal_conv1d.py +++ b/vllm_ascend/models/glm5next/ops/causal_conv1d.py @@ -1,72 +1,106 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Ascend causal conv1d entry points for the GLM-5.3-Flash KDA layers. - -The upstream CUDA kernels reference ``tl.extra.cuda.gdc_wait``, which Ascend -Triton does not provide -- the AST visitor raises even when ``launch_pdl`` is -False. Both entry points are therefore routed to Ascend implementations, and the -kwargs the upstream signatures grew for CUDA-side cache management are dropped -here rather than at every call site. - -``causal_conv1d_update`` prefers the NPU Triton kernel, which has no host sync -and accepts the spec-decode arguments directly. The PyTorch fallback calls -``.item()`` per request, so ACL graph capture stalls at decode-FULL when the -kernel is unavailable; ``has_npu_triton_conv1d_update()`` lets the caller report -that up front instead of hanging during capture. -""" +"""AscendC short convolution for GLM prefill, decode and MTP verification.""" import torch +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm_ascend.ops.causal_conv1d import ( - causal_conv1d_fn as _torch_causal_conv1d_fn, -) -from vllm_ascend.ops.causal_conv1d import ( - causal_conv1d_update as _torch_causal_conv1d_update, -) - -try: - from vllm_ascend.ops.triton.mamba.causal_conv1d import ( # type: ignore[attr-defined] - causal_conv1d_update_npu as _npu_triton_conv1d_update, - ) - - _HAS_NPU_TRITON_CONV1D_UPDATE = True -except ImportError: - _npu_triton_conv1d_update = None - _HAS_NPU_TRITON_CONV1D_UPDATE = False - -# Cache-management and validation kwargs the upstream CUDA signatures accept but -# the Ascend implementations neither need nor understand. -_UNSUPPORTED_KWARGS = ( - "null_block_id", - "block_idx_first_scheduled_token", - "block_idx_last_scheduled_token", - "initial_state_idx", - "num_computed_tokens", - "block_size_to_align", - "validate_data", - "metadata", -) +CONV_STATE_COPY_BLOCK_SIZE = 256 -_UPDATE_FALLBACK_UNSUPPORTED_KWARGS = (*_UNSUPPORTED_KWARGS, "max_query_len", "out") +@triton.jit +def _copy_conv_state( + cache, + packed, + cache_indices, + starts, + packed_indices, + cache_stride, + index_stride, + num_slots, + STATE_LEN: tl.constexpr, + DIM: tl.constexpr, + STATE_STRIDE: tl.constexpr, + DIM_STRIDE: tl.constexpr, + WRITE_BACK: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + slot = tl.load(cache_indices + request * index_stride).to(tl.int64) + active = (slot >= 0) & (slot < num_slots) & (tl.load(starts + request + 1) > tl.load(starts + request)) + in_range = offsets < STATE_LEN * DIM + safe_slot = tl.where(active, slot, 0) + cache_offsets = safe_slot * cache_stride + offsets // DIM * STATE_STRIDE + offsets % DIM * DIM_STRIDE + packed_offsets = request * STATE_LEN * DIM + offsets + if WRITE_BACK: + values = tl.load(packed + packed_offsets, mask=in_range, other=0) + tl.store(cache + cache_offsets, values, mask=active & in_range) + else: + values = tl.load(cache + cache_offsets, mask=active & in_range, other=0) + tl.store(packed + packed_offsets, values, mask=in_range) + if tl.program_id(1) == 0: + tl.store(packed_indices + request, tl.where(active, request, -1)) -def has_npu_triton_conv1d_update() -> bool: - """Whether the host-sync-free NPU Triton update kernel is available.""" - return _HAS_NPU_TRITON_CONV1D_UPDATE - -def causal_conv1d_fn(*args, **kwargs) -> torch.Tensor: - for key in _UNSUPPORTED_KWARGS: - kwargs.pop(key, None) - return _torch_causal_conv1d_fn(*args, **kwargs) - - -def causal_conv1d_update(*args, **kwargs) -> torch.Tensor: - if _npu_triton_conv1d_update is not None: - for key in _UNSUPPORTED_KWARGS: - kwargs.pop(key, None) - return _npu_triton_conv1d_update(*args, **kwargs) - - for key in _UPDATE_FALLBACK_UNSUPPORTED_KWARGS: - kwargs.pop(key, None) - return _torch_causal_conv1d_update(*args, **kwargs) +def causal_conv1d( + x: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + *, + run_mode: int, + initial_state_mode: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, +) -> torch.Tensor: + """Consume GDN metadata and update the caller's [cache, state_len, dim] state.""" + # Padded requests can be skipped by the kernel; their output must stay zero. + output = torch.zeros_like(x) + if cache_indices.shape[0] == 0: + return output + kernel_state = conv_state + kernel_indices = cache_indices + # aclnnCausalConv1d materializes a non-contiguous state without writing its + # mutations back to the view. Stage only this batch's rows, retaining both + # page strides and DS layouts; never copy the entire persistent cache. + if not conv_state.is_contiguous(): + requests = cache_indices.shape[0] + state_len, dim = conv_state.shape[1:] + kernel_state = torch.empty((requests, state_len, dim), dtype=conv_state.dtype, device=conv_state.device) + kernel_indices = torch.empty(requests, dtype=torch.int32, device=cache_indices.device) + copy_grid = (requests, triton.cdiv(state_len * dim, CONV_STATE_COPY_BLOCK_SIZE)) + copy_args = ( + conv_state, + kernel_state, + cache_indices, + query_start_loc, + kernel_indices, + conv_state.stride(0), + cache_indices.stride(0), + conv_state.shape[0], + state_len, + dim, + conv_state.stride(1), + conv_state.stride(2), + ) + _copy_conv_state[copy_grid](*copy_args, WRITE_BACK=False, BLOCK=CONV_STATE_COPY_BLOCK_SIZE) + # Return the declared result so graph functionalization retains the call. + result = torch.ops._C_ascend.npu_causal_conv1d_custom( + output, + x, + weight, + conv_state=kernel_state, + bias_opt=None, + query_start_loc_opt=query_start_loc, + cache_indices_opt=kernel_indices, + initial_state_mode_opt=initial_state_mode, + num_accepted_tokens_opt=num_accepted_tokens, + activation_mode=1, + pad_slot_id=PAD_SLOT_ID, + run_mode=run_mode, + ) + if not conv_state.is_contiguous(): + _copy_conv_state[copy_grid](*copy_args, WRITE_BACK=True, BLOCK=CONV_STATE_COPY_BLOCK_SIZE) + return result diff --git a/vllm_ascend/models/glm5next/ops/kda.py b/vllm_ascend/models/glm5next/ops/kda.py new file mode 100644 index 000000000000..f8dd6d19680e --- /dev/null +++ b/vllm_ascend/models/glm5next/ops/kda.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM bounded-gate contracts for the AscendC KDA operators.""" + +import torch + +from vllm_ascend.models.glm5next.ops.state_ops import gather_initial_states, scatter_states +from vllm_ascend.ops.kda import run_chunk_kda, run_recurrent_kda + +KDA_MAX_RECURRENT_TOKENS = 8 + + +def recurrent_kda( + q, + k, + v, + raw_gate, + raw_beta, + state, + cu_seqlens, + state_indices, + a_log, + dt_bias, + lower_bound, + num_accepted_tokens=None, +): + """Update the selected VK state slots, including MTP rejection rollback.""" + num_seqs = cu_seqlens.numel() - 1 + state_indices = state_indices[:num_seqs] + if num_accepted_tokens is not None: + num_accepted_tokens = num_accepted_tokens[:num_seqs] + output = run_recurrent_kda( + q, + k, + v, + raw_gate, + raw_beta, + state, + cu_seqlens, + state_indices, + a_log, + dt_bias, + lower_bound=lower_bound, + beta_is_preprocessed=False, + num_accepted_tokens=num_accepted_tokens, + ) + valid = torch.arange(q.shape[1], device=q.device) < cu_seqlens[-1] + return output.masked_fill(~valid[None, :, None, None], 0) + + +def chunk_kda( + q, + k, + v, + raw_gate, + raw_beta, + state, + state_indices, + has_initial_state, + metadata, + a_log, + dt_bias, + lower_bound, +): + """Run prefill with CPU chunk descriptors prepared by the GDN builder.""" + if metadata.keep_meta is not None: + state_indices = state_indices[metadata.keep_meta] + has_initial_state = has_initial_state[metadata.keep_meta] + initial_state = gather_initial_states(state, state_indices, has_initial_state).float().contiguous() + cu_seqlens = metadata.cu_seqlens_host if metadata.cu_seqlens_kern is None else metadata.cu_seqlens_kern + output, final_state = run_chunk_kda( + q, + k, + v, + raw_gate, + raw_beta.float().sigmoid(), + initial_state, + cu_seqlens, + metadata.chunk_indices_chunk64_host, + a_log, + dt_bias, + lower_bound=lower_bound, + ) + scatter_states(state, final_state.to(state.dtype), state_indices) + return output diff --git a/vllm_ascend/ops/kda.py b/vllm_ascend/ops/kda.py new file mode 100644 index 000000000000..1c500139db72 --- /dev/null +++ b/vllm_ascend/ops/kda.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared AscendC KDA execution; callers own projections and cache updates.""" + +from collections.abc import Sequence + +import torch +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd + +KDA_CHUNK_SIZE = 64 + + +def run_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, + cu_seqlens: torch.Tensor, + state_indices: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + lower_bound: float | None, + beta_is_preprocessed: bool = True, + num_accepted_tokens: torch.Tensor | None = None, +) -> torch.Tensor: + # Recurrent KDA consumes independent Q/K/V token/head strides directly. + return torch.ops._C_ascend.recurrent_kda( + q, + k, + v, + raw_gate.contiguous(), + beta.contiguous(), + state, + cu_seqlens, + state_indices, + a_log.reshape(-1).contiguous(), + dt_bias.contiguous(), + num_accepted_tokens=num_accepted_tokens, + scale=q.shape[-1] ** -0.5, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=not beta_is_preprocessed, + allow_neg_eigval=False, + safe_gate=lower_bound is not None, + lower_bound=lower_bound if lower_bound is not None else -5.0, + ) + + +def run_chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor | Sequence[int], + chunk_indices: torch.Tensor | Sequence[int], + a_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Consume preprocessed beta and return output plus the final VK state.""" + output, final_state, *_ = torch.ops._C_ascend.chunk_kda_fwd( + l2norm_fwd(q.contiguous()), + l2norm_fwd(k.contiguous()), + v.contiguous(), + raw_gate.contiguous(), + beta.contiguous(), + q.shape[-1] ** -0.5, + KDA_CHUNK_SIZE, + layout="BSND", + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + safe_gate=lower_bound is not None, + lower_bound=lower_bound if lower_bound is not None else -5.0, + use_gate_in_kernel=True, + A_log=a_log.reshape(-1).contiguous(), + dt_bias=dt_bias.contiguous(), + disable_recompute=False, + return_intermediate_states=False, + state_v_first=True, + ) + return output, final_state diff --git a/vllm_ascend/ops/kimi_kda.py b/vllm_ascend/ops/kimi_kda.py index e5f5705bd38f..02c8fd3362e5 100644 --- a/vllm_ascend/ops/kimi_kda.py +++ b/vllm_ascend/ops/kimi_kda.py @@ -22,12 +22,12 @@ from vllm.models.kimi_k3.nvidia.kda import ( KimiK3DeltaAttention, ) -from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm_ascend.ops.gdn_attn_builder import AscendGDNAttentionBackend +from vllm_ascend.ops.kda import run_chunk_kda, run_recurrent_kda from vllm_ascend.ops.triton.fla.utils import clear_ssm_states from vllm_ascend.quantization.methods.w4a8.w4a8_mxfp4 import ( AscendW4A8MXFPDynamicLinearMethod, @@ -37,7 +37,6 @@ ) from vllm_ascend.utils import npu_stream_switch -_KDA_CHUNK_SIZE = 64 _PACKED_CONV_WEIGHT_NAME = "ascend_conv1d_weight" _F_PROJ_SHARD_ID = 1 _KDA_BFG_STREAM: torch.npu.Stream | None = None @@ -422,27 +421,19 @@ def _run_recurrent( *, num_accepted_tokens: torch.Tensor | None = None, ) -> torch.Tensor: - # Preserve the strided Q/K/V views split from the fused projection; - # recurrent KDA consumes their independent token/head strides. - return torch.ops._C_ascend.recurrent_kda( + return run_recurrent_kda( q, k, v, - raw_gate.contiguous(), - beta.contiguous(), + raw_gate, + beta, recurrent_state, cu_seqlens, state_indices, - self.A_log.reshape(-1).contiguous(), - self.dt_bias.contiguous(), + self.A_log, + self.dt_bias, + lower_bound=self.gate_lower_bound, num_accepted_tokens=num_accepted_tokens, - scale=self.head_dim**-0.5, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=False, - allow_neg_eigval=False, - safe_gate=self.gate_lower_bound is not None, - lower_bound=(self.gate_lower_bound if self.gate_lower_bound is not None else -5.0), ) def _run_prefill( @@ -472,32 +463,21 @@ def _run_prefill( initial_state_vk = recurrent_state[state_indices].contiguous() clear_ssm_states(initial_state_vk, has_initial_state) - q = l2norm_fwd(q.contiguous()) - k = l2norm_fwd(k.contiguous()) - result = torch.ops._C_ascend.chunk_kda_fwd( + output, final_state = run_chunk_kda( q, k, - v.contiguous(), - raw_gate.contiguous(), - beta.contiguous(), - self.head_dim**-0.5, - _KDA_CHUNK_SIZE, - layout="BSND", - initial_state=initial_state_vk, - output_final_state=True, - cu_seqlens=cu_seqlens, - chunk_indices=prebuilt_metadata.chunk_indices_chunk64_host, - safe_gate=self.gate_lower_bound is not None, - lower_bound=self.gate_lower_bound if self.gate_lower_bound is not None else -5.0, - use_gate_in_kernel=True, - A_log=self.A_log.reshape(-1).contiguous(), - dt_bias=self.dt_bias.contiguous(), - disable_recompute=False, - return_intermediate_states=False, - state_v_first=True, + v, + raw_gate, + beta, + initial_state_vk, + cu_seqlens, + prebuilt_metadata.chunk_indices_chunk64_host, + self.A_log, + self.dt_bias, + lower_bound=self.gate_lower_bound, ) - recurrent_state[state_indices] = result[1].to(recurrent_state.dtype) - return result[0] + recurrent_state[state_indices] = final_state.to(recurrent_state.dtype) + return output @eager_break_during_capture def _forward( diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index c1437a8d2c2a..c3dae5c92ed8 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -257,6 +257,7 @@ AscendSlidingWindowMLASpec, get_kv_cache_compression_ratio, get_storage_block_size, + requires_padded_page_layout, ) from vllm_ascend.core.profiling_chunk_predictor import ( _finish_profiling_chunk_timing, @@ -4522,10 +4523,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str # standardized descriptors, whose ``size`` is the size of one common # backing allocation rather than the size of an individual layer. use_legacy_shared_by_layout = vllm_version_is("0.28.0") - is_glm5_next = any( - getattr(spec, "model_version", None) == "glm5_next" - for spec in layer_kv_cache_spec.values() - ) + uses_padded_page_layout = requires_padded_page_layout(layer_kv_cache_spec.values()) is_dsv4_main = not use_legacy_shared_by_layout and any( getattr(spec, "model_version", None) == "deepseek_v4" for spec in layer_kv_cache_spec.values() @@ -4562,7 +4560,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str # the compressed indexer and its state cache). This differs from the # generic main layout, which treats descriptor layers as independent # regions within one common backing. - if is_glm5_next: + if uses_padded_page_layout: for descriptor in kv_cache_config.kv_cache_tensors: shared_layers = get_kv_cache_tensor_layers(descriptor) if not shared_layers: @@ -4649,7 +4647,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str if ( not use_legacy_shared_by_layout and not is_dsv4_main - and not is_glm5_next + and not uses_padded_page_layout and self.hybrid_with_attn_and_mamba and supports_shared_backing_with_kv_transfer and not self.use_sparse @@ -4956,7 +4954,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str def _adjust_kv_layout( self, raw_tensor: torch.Tensor, - kv_cache_shape_list: list[int], + kv_cache_shape_list: list[tuple[int, ...]], kv_cache_dtype_list: list[int], page_size_bytes: int, overlap_full_kv_cache: bool = False, @@ -5004,6 +5002,7 @@ def _reshape_kv_cache_tensors( """ kv_caches: dict[str, torch.Tensor] = {} layer_kv_cache_spec = self._get_layer_kv_cache_specs(kv_cache_config) + uses_padded_page_layout = requires_padded_page_layout(layer_kv_cache_spec.values()) for group in self._kv_cache_spec_attn_group_iterator(): attn_backend = group.backend current_kv_cache_spec = group.kv_cache_spec @@ -5371,6 +5370,19 @@ def _reshape_kv_cache_tensors( assert raw_tensor.numel() % current_kv_cache_spec.page_size_bytes == 0 num_blocks = raw_tensor.numel() // current_kv_cache_spec.page_size_bytes assert num_blocks >= kv_cache_config.num_blocks + if uses_padded_page_layout: + # Recurrent state shares its physical pages with the + # block-strided attention caches of the same pool. Both + # views must advance by the same padded page size so + # different scheduler block IDs cannot overwrite each + # other. + kv_caches[layer_name] = self._adjust_kv_layout( + raw_tensor, + [(num_blocks, *shape) for shape in current_kv_cache_spec.shapes], + current_kv_cache_spec.dtypes, + current_kv_cache_spec.page_size_bytes, + ) + continue # `num_blocks` is the number of blocks the model runner can use. # `kv_cache_config.num_blocks` is the number of blocks that