From 0bff0eb46d020b89ec19e1436d960a788b5f8807 Mon Sep 17 00:00:00 2001 From: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:27:12 +0800 Subject: [PATCH 1/3] feat(glm): integrate pooled indexing through the shared SFA attention backend Connect Triton pooled indexing to the shared SFA path. Keep the indexer backend and its hardware/context-parallel capability checks in attention, preserving physical cache pages, compressor state and causal-tail visibility. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> --- .../test_glm5next_kpool_model_backend.py | 239 ++++++++++++++++++ tests/ut/models/test_glm5next_kv_cache.py | 65 +++-- tests/ut/models/test_glm5next_sfa_routing.py | 86 +++++++ tests/ut/test_platform.py | 13 +- vllm_ascend/attention/indexer_kpool.py | 89 ++----- .../attention/indexer_kpool_backend.py | 187 ++++++++++++++ vllm_ascend/models/glm5next/attention.py | 89 +------ .../models/glm5next/ops/kpool_compress.py | 75 ------ .../glm5next/sparse_attn_indexer_kpool.py | 192 ++++++++------ vllm_ascend/platform.py | 9 - vllm_ascend/utils.py | 4 +- 11 files changed, 719 insertions(+), 329 deletions(-) create mode 100644 tests/ut/models/test_glm5next_kpool_model_backend.py create mode 100644 tests/ut/models/test_glm5next_sfa_routing.py create mode 100644 vllm_ascend/attention/indexer_kpool_backend.py delete mode 100644 vllm_ascend/models/glm5next/ops/kpool_compress.py diff --git a/tests/ut/models/test_glm5next_kpool_model_backend.py b/tests/ut/models/test_glm5next_kpool_model_backend.py new file mode 100644 index 000000000000..b46632fa26ae --- /dev/null +++ b/tests/ut/models/test_glm5next_kpool_model_backend.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from torch import nn + +from vllm_ascend.attention.indexer_kpool import ( + AscendIndexerKPoolMetadata, + AscendIndexerKPoolStateMetadata, +) + +# The operator PR owns these kernels. Each orchestration test below replaces +# the kernel call with an asserting CPU implementation. +with patch.dict( + "sys.modules", + { + "vllm_ascend.ops.triton.glm5_next_kpool_state_compress": MagicMock(), + "vllm_ascend.ops.triton.glm5_next_lightning_indexer": MagicMock(), + }, +): + import vllm_ascend.attention.indexer_kpool_backend as backend_module + import vllm_ascend.models.glm5next.sparse_attn_indexer_kpool as kpool_module + from vllm_ascend.attention.indexer_kpool_backend import Glm5NextKPoolIndexerBackend + from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import SparseAttnIndexerKpool, append_causal_tail + + +@pytest.mark.parametrize("pool_size", [1, 4]) +def test_causal_tail_follows_valid_history_without_holes(pool_size: int) -> None: + topk = 8 + positions = torch.tensor([0, 1, 2, 3, 4, 6, 7, 8, 10, 15, 18]) + indices = torch.full((positions.numel(), topk + pool_size - 1), -1, dtype=torch.int32) + expected_rows = [] + for row, position in enumerate(positions.tolist()): + tail_start = (position + 1) // pool_size * pool_size + # Reverse the selected pools to ensure their ranking is preserved. + groups = list(reversed(range(min(tail_start // pool_size, topk // pool_size)))) + history = [group * pool_size + offset for group in groups for offset in range(pool_size)] + indices[row, : len(history)] = torch.tensor(history, dtype=torch.int32) + tail = list(range(tail_start, position + 1)) + if tail: + indices[row, topk : topk + len(tail)] = torch.tensor(tail, dtype=torch.int32) + expected_rows.append(history + tail + [-1] * (indices.shape[1] - len(history) - len(tail))) + + pointer = indices.data_ptr() + append_causal_tail(indices, positions, topk, pool_size) + + assert indices.data_ptr() == pointer + torch.testing.assert_close(indices, torch.tensor(expected_rows, dtype=torch.int32)) + + +def _indexer_metadata(num_tokens: int = 8) -> AscendIndexerKPoolMetadata: + return AscendIndexerKPoolMetadata( + block_table=torch.tensor([[0], [1]], dtype=torch.int32), + slot_mapping=torch.tensor( + [-1, -1, -1, 0, -1, -1, -1, 2, -1, -1], + dtype=torch.int64, + )[:num_tokens], + seq_lens=torch.tensor([1, 1], dtype=torch.int32), + seq_lens_cpu=torch.tensor([1, 1], dtype=torch.int32), + positions=torch.tensor([0, 1, 2, 3, 0, 1, 2, 3, 0, 0])[:num_tokens], + block_size=2, + compress_ratio=4, + cum_query_lens=torch.tensor([4, 8], dtype=torch.int32), + raw_seq_lens=torch.tensor([4, 4], dtype=torch.int32), + num_actual_tokens=8, + ) + + +def _state_metadata() -> AscendIndexerKPoolStateMetadata: + return AscendIndexerKPoolStateMetadata( + block_table=torch.tensor([[0], [1]], dtype=torch.int32), + slot_mapping=torch.full((8,), -1, dtype=torch.int64), + block_size=4, + cache_role="indexer_state", + ) + + +@pytest.mark.parametrize("compute_topk", [False, True]) +def test_triton_indexer_updates_both_caches_and_masks_padding(monkeypatch, compute_topk): + metadata = _indexer_metadata(num_tokens=10) + state_metadata = _state_metadata() + state_metadata.slot_mapping = torch.arange(10) + indexer_cache = torch.zeros(2, 2, 1, 2, dtype=torch.bfloat16) + state_cache = torch.zeros(2, 4, 1, 4) + + def compress(state, cache, k, gate, ape, positions, query_ends, seq_lens, state_slots, table, indexer_slots, pool): + assert k.dtype == gate.dtype == state.dtype == torch.float32 + assert pool == 4 + torch.testing.assert_close(query_ends, metadata.cum_query_lens) + torch.testing.assert_close(seq_lens, metadata.raw_seq_lens) + assert table is state_metadata.block_table + torch.testing.assert_close(indexer_slots, metadata.slot_mapping) + state[0, 0].fill_(7) + cache[0, 0].fill_(11) + + select = MagicMock(return_value=torch.full((10, 1, 7), -1, dtype=torch.int32)) + monkeypatch.setattr(kpool_module, "glm5_next_kpool_state_compress_and_write_cache_triton", compress) + monkeypatch.setattr(kpool_module, "glm5_next_lightning_indexer_triton", select) + result = SparseAttnIndexerKpool(4, 2)( + torch.zeros(10, 2), + torch.zeros(10, 1, 2, dtype=torch.bfloat16), + torch.ones(10, 1, dtype=torch.bfloat16), + metadata.positions, + indexer_cache, + state_cache, + metadata, + state_metadata, + gate_score=torch.zeros(10, 2), + compress_ape=torch.zeros(4, 2), + index_kpool=4, + max_pool_seq_len=1, + compute_topk=compute_topk, + ) + torch.testing.assert_close(state_cache[0, 0], torch.full_like(state_cache[0, 0], 7)) + torch.testing.assert_close(indexer_cache[0, 0], torch.full_like(indexer_cache[0, 0], 11)) + if compute_topk: + assert result.shape == (10, 1, 7) + assert (result[8:] == -1).all() + select.assert_called_once() + else: + assert result is None + select.assert_not_called() + + +@pytest.mark.parametrize( + ("pcp_size", "dcp_size"), + [(2, 1), (1, 2)], +) +def test_kpool_backend_rejects_context_parallelism(pcp_size: int, dcp_size: int) -> None: + source = SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=dcp_size, + ) + ) + ) + + with pytest.raises(NotImplementedError, match="PCP or DCP"): + Glm5NextKPoolIndexerBackend(source, qk_rope_head_dim=0) + + +class _Projection(nn.Module): + def forward(self, q_c): + return q_c.repeat(1, 2), None + + +class _RecordingKPool(nn.Module): + def __init__(self): + super().__init__() + self.args = None + self.kwargs = None + + def forward(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + if kwargs["compute_topk"]: + return torch.zeros(args[0].shape[0], 1, 5, dtype=torch.int32) + return None + + +def test_backend_uses_normalized_q_c_and_separate_state_metadata( + monkeypatch, +) -> None: + backend = Glm5NextKPoolIndexerBackend.__new__(Glm5NextKPoolIndexerBackend) + nn.Module.__init__(backend) + backend.n_head = 2 + backend.head_dim = 2 + backend.topk_tokens = 2 + backend.index_kpool = 4 + backend.wq_b = _Projection() + backend.wk_weights_proj = nn.Linear(3, 4, bias=False) + backend.k_norm = nn.LayerNorm(2) + backend.index_kpool_compress_ape = nn.Parameter(torch.zeros(4, 2)) + backend.index_kpool_compress_gate = nn.Parameter(torch.zeros(2, 3)) + backend.k_cache = SimpleNamespace( + prefix="indexer.k_cache", + kv_cache=torch.zeros(2, 2, 1, 2, dtype=torch.bfloat16), + ) + backend.state_cache = SimpleNamespace( + prefix="indexer.state", + kv_cache=torch.zeros(2, 4, 1, 4, dtype=torch.float32), + ) + backend.topk_indices_buffer = None + backend.softmax_scale = 0.5 + backend._wk_weight_f32 = None + backend.indexer_op = _RecordingKPool() + state_metadata = _state_metadata() + monkeypatch.setattr( + backend_module, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={"indexer.state": state_metadata}, + cudagraph_runtime_mode=None, + virtual_engine=0, + ), + ) + + normalized_q_c = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + hidden = torch.ones(2, 3) + metadata = _indexer_metadata(num_tokens=2) + metadata.num_actual_tokens = 2 + metadata.cum_query_lens = torch.tensor([1, 2], dtype=torch.int32) + metadata.raw_seq_lens = torch.tensor([1, 1], dtype=torch.int32) + metadata.seq_lens = torch.tensor([0, 0], dtype=torch.int32) + metadata.positions = torch.tensor([0, 0]) + metadata.slot_mapping = torch.tensor([-1, -1]) + + result = backend.forward( + hidden, + normalized_q_c, + None, + None, + hidden + 3, + metadata, + compute_topk=True, + ) + + assert result is not None + assert backend.indexer_op.args is not None + torch.testing.assert_close(backend.indexer_op.args[1], normalized_q_c.repeat(1, 2).view(2, 2, 2)) + expected_k = torch.nn.functional.layer_norm( + torch.nn.functional.linear(hidden + 3, backend.wk_weights_proj.weight)[:, :2], + (2,), + backend.k_norm.weight, + backend.k_norm.bias, + backend.k_norm.eps, + ) + torch.testing.assert_close(backend.indexer_op.args[0], expected_k) + assert backend.indexer_op.args[0].dtype == torch.float32 + expected_weights = torch.nn.functional.linear(hidden, backend.wk_weights_proj.weight[2:]) * (0.5 * 2**-0.5) + torch.testing.assert_close(backend.indexer_op.args[2], expected_weights) + assert backend.indexer_op.args[7] is state_metadata + assert backend.indexer_op.kwargs is not None + assert backend.indexer_op.kwargs["compute_topk"] is True diff --git a/tests/ut/models/test_glm5next_kv_cache.py b/tests/ut/models/test_glm5next_kv_cache.py index aa13e365777d..8e1181080ca6 100644 --- a/tests/ut/models/test_glm5next_kv_cache.py +++ b/tests/ut/models/test_glm5next_kv_cache.py @@ -8,14 +8,12 @@ import pytest import torch from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager -from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm_ascend.attention.indexer_kpool import ( AscendIndexerKPoolBackend, AscendIndexerKPoolMetadataBuilder, AscendIndexerKPoolStateBackend, - select_indexer_block_size, ) from vllm_ascend.attention.utils import AscendCommonAttentionMetadata from vllm_ascend.core.kv_cache_interface import ( @@ -89,18 +87,55 @@ def test_invalid_pool_geometry_is_rejected(ratio): format_indexer_kpool_slot_mapping(torch.tensor([0]), torch.tensor([0]), 128, ratio) -@pytest.mark.parametrize( - ("storage_block_size", "expected"), - [(16, (16, 1)), (1024, (1024, 1)), (2048, (1024, 2)), (1536, (768, 2))], -) -def test_indexer_block_size_selection(storage_block_size, expected): - assert select_indexer_block_size(storage_block_size) == expected - - -@pytest.mark.parametrize("storage_block_size", [0, 7, 17]) -def test_invalid_indexer_block_size_is_rejected(storage_block_size): - with pytest.raises(ValueError): - select_indexer_block_size(storage_block_size) +@pytest.mark.parametrize("storage_block_size", [8, 24, 144, 1536, 2048]) +def test_indexer_metadata_addresses_complete_storage_pages(storage_block_size): + pool_size = 16 + logical_size = storage_block_size * pool_size + split = logical_size // 128 + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4, max_num_seqs=1), + model_config=SimpleNamespace(max_model_len=logical_size * 3), + ) + spec = AscendMLAAttentionSpec( + block_size=logical_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + **({"compress_ratio": pool_size} if vllm_version_is("0.28.0") else {"tokens_per_state": pool_size}), + model_version="glm5_next", + ) + builders = [ + AscendIndexerKPoolMetadataBuilder(spec, ["layer.indexer.k_cache"], config, torch.device("cpu")) + for _ in range(2) + ] + pages = torch.tensor([[7, 2, -1]], dtype=torch.int32) + expanded = (pages.unsqueeze(-1) * split + torch.arange(split)).reshape(1, -1).int() + expanded[:, -split:] = -1 + common = SimpleNamespace( + num_reqs=1, + num_input_tokens=4, + num_actual_tokens=3, + max_query_len=3, + query_start_loc=torch.tensor([0, 3], dtype=torch.int32), + seq_lens=torch.tensor([logical_size + pool_size], dtype=torch.int32), + _seq_lens_cpu=None, + seq_lens_cpu=None, + positions=torch.tensor([logical_size - 1, logical_size, logical_size + pool_size - 1, 0]), + slot_mapping=torch.tensor([8 * logical_size - 1, 2 * logical_size, 2 * logical_size + pool_size - 1, -1]), + block_table_tensor=expanded, + ) + first, draft = [builder.build(0, common) for builder in builders] + assert first.block_size == storage_block_size + torch.testing.assert_close(first.block_table, pages) + assert first.slot_mapping.tolist() == [8 * storage_block_size - 1, -1, 2 * storage_block_size, -1] + assert first.seq_lens.tolist() == [storage_block_size + 1] + address = first.block_table.data_ptr() + assert draft.block_table.data_ptr() != address + common.block_table_tensor[:, :split] = 3 * split + torch.arange(split) + refreshed = builders[0].build(0, common) + assert refreshed.block_table.data_ptr() == address + assert first.block_table.tolist() == [[3, 2, -1]] + assert draft.block_table.tolist() == [[7, 2, -1]] def test_model_cache_layers_publish_source_compatible_specs(): @@ -160,7 +195,7 @@ def test_indexer_metadata_preserves_raw_request_boundaries(): model_config=SimpleNamespace(max_model_len=512), ) builder = AscendIndexerKPoolMetadataBuilder( - MLAAttentionSpec( + AscendMLAAttentionSpec( block_size=256, num_kv_heads=1, head_size=128, diff --git a/tests/ut/models/test_glm5next_sfa_routing.py b/tests/ut/models/test_glm5next_sfa_routing.py new file mode 100644 index 000000000000..01c6cf1b5d5d --- /dev/null +++ b/tests/ut/models/test_glm5next_sfa_routing.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import vllm.config as vllm_config_module + +import vllm_ascend.platform as platform +from vllm_ascend.device.hardware_profile import AttentionBackendFamily, DeviceAdaptorFamily + +# Routing and visible-length checks do not execute the external kernels. +with patch.dict( + "sys.modules", + { + "vllm_ascend.ops.triton.glm5_next_kpool_state_compress": MagicMock(), + "vllm_ascend.ops.triton.glm5_next_lightning_indexer": MagicMock(), + }, +): + import vllm_ascend.attention.indexer_kpool_backend as backend_module + from vllm_ascend.attention.indexer_kpool_backend import Glm5NextKPoolIndexerBackend + + +@pytest.mark.parametrize("kpool", [False, True]) +@pytest.mark.parametrize("fp8_device", [False, True]) +def test_sparse_models_select_shared_sfa(monkeypatch, kpool, fp8_device): + config = SimpleNamespace(index_topk=2048) + if kpool: + config.index_kpool = 4 + monkeypatch.setattr( + vllm_config_module, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=SimpleNamespace(hf_text_config=config)), + ) + monkeypatch.setattr( + platform, + "get_current_hardware_profile", + lambda: SimpleNamespace( + attention_backend_family=AttentionBackendFamily.STANDARD, + device_adaptor_family=DeviceAdaptorFamily.FP8_OPTIMIZED if fp8_device else None, + ), + ) + monkeypatch.setattr(platform, "_validate_fa3_backend", lambda *args: False) + selector = SimpleNamespace(use_mla=True, use_sparse=True, use_pcp=False) + assert platform.NPUPlatform.get_attn_backend_cls(None, selector) == "vllm_ascend.attention.sfa_v1.AscendSFABackend" + + +@pytest.mark.parametrize("mode", ["310p", "pcp", "dcp"]) +def test_kpool_unsupported_routes_fail_explicitly(monkeypatch, mode): + monkeypatch.setattr( + backend_module, + "get_current_hardware_profile", + lambda: SimpleNamespace( + attention_backend_family=AttentionBackendFamily.COMPATIBILITY + if mode == "310p" + else AttentionBackendFamily.STANDARD + ), + ) + source = SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=2 if mode == "pcp" else 1, + decode_context_parallel_size=2 if mode == "dcp" else 1, + ) + ) + ) + with pytest.raises(NotImplementedError, match="requires Ascend|PCP or DCP"): + Glm5NextKPoolIndexerBackend(source, qk_rope_head_dim=0) + + +def _indexer(): + backend = object.__new__(Glm5NextKPoolIndexerBackend) + torch.nn.Module.__init__(backend) + backend.topk_tokens, backend.index_kpool = 2048, 4 + return backend + + +def test_indexer_owns_pool_boundary_lengths(): + indexer = _indexer() + assert indexer.topk_output_width == 2051 + torch.testing.assert_close( + indexer.get_topk_lengths(torch.tensor([0, 2, 3, 2047, 2048, 2050, 4095])), + torch.tensor([1, 3, 4, 2048, 2049, 2051, 2048]), + ) diff --git a/tests/ut/test_platform.py b/tests/ut/test_platform.py index 0cdfc7e64a09..7e2c27d887f1 100644 --- a/tests/ut/test_platform.py +++ b/tests/ut/test_platform.py @@ -4,6 +4,7 @@ import pytest import torch +from vllm.config import set_current_vllm_config from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.platforms import PlatformEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -1766,7 +1767,11 @@ def test_get_attn_backend_cls_supports_legacy_config_without_use_dcp(self): use_pcp=True, ) - result = self.platform.get_attn_backend_cls("ascend", attn_selector_config) + config = self.mock_vllm_config() + config.model_config.hf_text_config = SimpleNamespace(index_topk=2048) + config.model_config.hf_config = config.model_config.hf_text_config + with set_current_vllm_config(config): + result = self.platform.get_attn_backend_cls("ascend", attn_selector_config) self.assertEqual(result, "vllm_ascend.attention.sfa_v1.AscendSFABackend") @@ -1780,7 +1785,11 @@ def test_get_attn_backend_cls_selects_sfa_pcp_backend(self): use_sparse=True, use_pcp=True, ) - result = self.platform.get_attn_backend_cls("ascend", attn_selector_config) + config = self.mock_vllm_config() + config.model_config.hf_text_config = SimpleNamespace(index_topk=2048) + config.model_config.hf_config = config.model_config.hf_text_config + with set_current_vllm_config(config): + result = self.platform.get_attn_backend_cls("ascend", attn_selector_config) self.assertEqual(result, "vllm_ascend.attention.sfa_v1.AscendSFABackend") def test_get_attn_backend_cls_rejects_unsupported_pcp_backend(self): diff --git a/vllm_ascend/attention/indexer_kpool.py b/vllm_ascend/attention/indexer_kpool.py index 58cd84dd67a9..efe52e621668 100644 --- a/vllm_ascend/attention/indexer_kpool.py +++ b/vllm_ascend/attention/indexer_kpool.py @@ -26,39 +26,6 @@ ) GLM5_NEXT_SFA_KERNEL_BLOCK_SIZE = 128 -INDEXER_KPOOL_MAX_BLOCK_SIZE = 1024 -INDEXER_KPOOL_BLOCK_ALIGNMENT = 16 - - -def select_indexer_block_size(storage_block_size: int) -> tuple[int, int]: - """Select a CANN-compatible physical block size for the indexer cache. - - ``pool_key_indexer`` requires the ``PA_BBND`` block dimension to be a - multiple of 16 and no larger than 1024. GLM-Next's logical attention block - can be much larger than 1024 after compression, so split one storage block - into several CANN-compatible sub-blocks when necessary. - - Returns ``(block_size, blocks_per_logical_block)``. - """ - if storage_block_size <= 0: - raise ValueError(f"Indexer KPool storage block size must be positive, got {storage_block_size}.") - if storage_block_size <= INDEXER_KPOOL_MAX_BLOCK_SIZE and storage_block_size % INDEXER_KPOOL_BLOCK_ALIGNMENT == 0: - return storage_block_size, 1 - max_candidate = min(storage_block_size, INDEXER_KPOOL_MAX_BLOCK_SIZE) - max_candidate -= max_candidate % INDEXER_KPOOL_BLOCK_ALIGNMENT - for candidate in range( - max_candidate, - INDEXER_KPOOL_BLOCK_ALIGNMENT - 1, - -INDEXER_KPOOL_BLOCK_ALIGNMENT, - ): - if storage_block_size % candidate == 0: - return candidate, storage_block_size // candidate - raise ValueError( - "Indexer KPool storage block size " - f"{storage_block_size} cannot be split into a block size that is a " - f"multiple of {INDEXER_KPOOL_BLOCK_ALIGNMENT} and no larger than " - f"{INDEXER_KPOOL_MAX_BLOCK_SIZE}." - ) @dataclass @@ -68,7 +35,7 @@ class AscendIndexerKPoolMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor seq_lens: torch.Tensor - seq_lens_cpu: torch.Tensor + seq_lens_cpu: torch.Tensor | None positions: torch.Tensor block_size: int compress_ratio: int @@ -111,6 +78,8 @@ def __init__( super().__init__(kv_cache_spec, layer_names, vllm_config, device) self.logical_block_size = kv_cache_spec.block_size self.storage_block_size = get_storage_block_size(kv_cache_spec) + if self.storage_block_size <= 0: + raise ValueError(f"Indexer KPool storage block size must be positive, got {self.storage_block_size}.") self.compress_ratio = compress_ratio if self.logical_block_size % GLM5_NEXT_SFA_KERNEL_BLOCK_SIZE: raise ValueError( @@ -119,9 +88,6 @@ def __init__( f"kernel={GLM5_NEXT_SFA_KERNEL_BLOCK_SIZE}." ) self.kernel_blocks_per_logical_block = self.logical_block_size // GLM5_NEXT_SFA_KERNEL_BLOCK_SIZE - self.indexer_block_size, self.indexer_blocks_per_logical_block = select_indexer_block_size( - self.storage_block_size - ) scheduler_config = vllm_config.scheduler_config # ACLGraph replay keeps the addresses captured on the first run. The # derived compressed metadata therefore needs persistent storage that @@ -152,7 +118,7 @@ def __init__( ) self._block_table_buffer = torch.empty( scheduler_config.max_num_seqs, - max_logical_blocks * self.indexer_blocks_per_logical_block, + max_logical_blocks, dtype=torch.int32, device=device, ) @@ -193,12 +159,9 @@ def build( elif common_attn_metadata.seq_lens_cpu is not None: seq_lens_cpu = common_attn_metadata.seq_lens_cpu[:num_reqs] else: - seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].to("cpu") - seq_lens_cpu = torch.div( - seq_lens_cpu, - self.compress_ratio, - rounding_mode="floor", - ) + seq_lens_cpu = None + if seq_lens_cpu is not None: + seq_lens_cpu = torch.div(seq_lens_cpu, self.compress_ratio, rounding_mode="floor") expanded_block_table = common_attn_metadata.block_table_tensor[:num_reqs] split = self.kernel_blocks_per_logical_block if expanded_block_table.shape[1] % split: @@ -207,51 +170,30 @@ def build( f"table: width={expanded_block_table.shape[1]}, split={split}." ) logical_width = expanded_block_table.shape[1] // split - expanded_width = logical_width * self.indexer_blocks_per_logical_block - if expanded_width > self._block_table_buffer.shape[1]: + if logical_width > self._block_table_buffer.shape[1]: raise ValueError( "GLM-Next indexer block table exceeds its persistent buffer: " - f"required={expanded_width}, capacity=" + f"required={logical_width}, capacity=" f"{self._block_table_buffer.shape[1]}." ) - block_table = self._block_table_buffer[:num_reqs, :expanded_width] + block_table = self._block_table_buffer[:num_reqs, :logical_width] # The common full-group table is expanded for the C128 SFA kernel: # scheduler block N becomes [split*N, ..., split*N+split-1]. The # compressed indexer owns one physical page per scheduler block, so it # must recover N rather than treating the SFA sub-blocks as pages. - base = torch.div( + torch.div( expanded_block_table[:, ::split], split, rounding_mode="floor", + out=block_table, ) - k = self.indexer_blocks_per_logical_block - if k == 1: - block_table.copy_(base) - else: - base_repeated = base.repeat_interleave(k, dim=1) - offsets = ( - torch.arange( - expanded_width, - dtype=torch.int32, - device=base.device, - ) - % k - ) - valid = base_repeated >= 0 - block_table.copy_( - torch.where( - valid, - base_repeated * k + offsets, - torch.full_like(base_repeated, -1), - ) - ) return AscendIndexerKPoolMetadata( block_table=block_table, slot_mapping=slot_mapping, seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, positions=positions, - block_size=self.indexer_block_size, + block_size=self.storage_block_size, compress_ratio=self.compress_ratio, cum_query_lens=cum_query_lens, raw_seq_lens=raw_seq_lens, @@ -272,9 +214,8 @@ def get_name() -> str: @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - # The scheduler manages logical token blocks; the physical cache uses - # storage_block_size, split into CANN-compatible sub-blocks when passed - # to pool_key_indexer. + # The scheduler manages logical token blocks. Triton consumes complete + # compressed storage pages with their actual size and strides. return [MultipleOf(1)] @staticmethod diff --git a/vllm_ascend/attention/indexer_kpool_backend.py b/vllm_ascend/attention/indexer_kpool_backend.py new file mode 100644 index 000000000000..f20c14555671 --- /dev/null +++ b/vllm_ascend/attention/indexer_kpool_backend.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Attention backend for the GLM-Next pooled sparse indexer.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import get_forward_context + +from vllm_ascend.attention.indexer import AscendSFAIndexerBackend +from vllm_ascend.attention.indexer_kpool import ( + AscendIndexerKPoolMetadata, + AscendIndexerKPoolStateMetadata, +) +from vllm_ascend.device.hardware_profile import AttentionBackendFamily, get_current_hardware_profile +from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import ( + SparseAttnIndexerKpool, +) + + +class Glm5NextKPoolIndexerBackend(AscendSFAIndexerBackend): + """Triton KPool implementation of the unified seven-argument indexer API.""" + + def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: + nn.Module.__init__(self) + if qk_rope_head_dim != 0: + raise ValueError( + f"GLM-Next KPool indexing supports NoPE queries only, got qk_rope_head_dim={qk_rope_head_dim}." + ) + parallel_config = vllm_indexer.vllm_config.parallel_config + if parallel_config.prefill_context_parallel_size > 1 or parallel_config.decode_context_parallel_size > 1: + raise NotImplementedError("GLM-Next KPool indexing does not support PCP or DCP.") + + if get_current_hardware_profile().attention_backend_family is AttentionBackendFamily.COMPATIBILITY: + raise NotImplementedError("KPool sparse attention requires Ascend A2, A3 or A5.") + + self.n_head: int = vllm_indexer.n_head + self.head_dim: int = vllm_indexer.head_dim + self.topk_tokens: int = vllm_indexer.topk_tokens + self.q_lora_rank: int = vllm_indexer.q_lora_rank + self.index_kpool: int = vllm_indexer.index_kpool + self.wq_b = vllm_indexer.wq_b + self.wk_weights_proj = vllm_indexer.wk_weights_proj + self.k_norm = vllm_indexer.k_norm + self.softmax_scale = vllm_indexer.softmax_scale + self.index_kpool_compress_ape = vllm_indexer.index_kpool_compress_ape + self.index_kpool_compress_gate = vllm_indexer.index_kpool_compress_gate + self.k_cache: Any = vllm_indexer.k_cache + self.state_cache: Any = vllm_indexer.state_cache + self.topk_indices_buffer: torch.Tensor | None = vllm_indexer.topk_indices_buffer + self.indexer_op = SparseAttnIndexerKpool(self.topk_tokens, self.head_dim) + self.enable_sparse_li_c8 = False + for name in ("_wk_weight_f32", "_gate_weight_f32", "_norm_weight_f32", "_norm_bias_f32"): + self.register_buffer(name, None, persistent=False) + + @property + def topk_output_width(self) -> int: + return self.topk_tokens + self.index_kpool - 1 + + def get_topk_lengths(self, positions: torch.Tensor) -> torch.Tensor: + visible = (positions + 1).clamp_min(0) + history = (visible // self.index_kpool * self.index_kpool).clamp(max=self.topk_tokens) + return history + visible % self.index_kpool + + @property + def num_cache_tensors(self) -> int: + return 1 + + def process_weights_after_loading(self) -> None: + self._wk_weight_f32 = self.wk_weights_proj.weight.detach().float() + self._gate_weight_f32 = self.index_kpool_compress_gate.detach().float() + self._norm_weight_f32 = self.k_norm.weight.detach().float() if self.k_norm.weight is not None else None + self._norm_bias_f32 = self.k_norm.bias.detach().float() if self.k_norm.bias is not None else None + + @staticmethod + def _bound_cache(layer: Any) -> torch.Tensor: + context = get_forward_context() + cache = layer.kv_cache + if isinstance(cache, (list, tuple)): + virtual_engine = getattr(context, "virtual_engine", 0) or 0 + if virtual_engine >= len(cache): + raise IndexError(f"Cache virtual engine {virtual_engine} is out of range.") + cache = cache[virtual_engine] + if isinstance(cache, (list, tuple)): + if len(cache) != 1: + raise TypeError("GLM KPool cache must contain one tensor.") + cache = cache[0] + if not isinstance(cache, torch.Tensor): + raise TypeError(f"GLM KPool cache {layer.prefix!r} is not bound to a tensor.") + return cache + + def forward( + self, + hidden_states: torch.Tensor, + q_c: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + cos: torch.Tensor | None, + sin: torch.Tensor | None, + k_hidden_states: torch.Tensor, + indexer_metadata: Any, + compute_topk: bool = True, + ) -> torch.Tensor | None: + del cos, sin + if not isinstance(indexer_metadata, AscendIndexerKPoolMetadata): + raise TypeError("GLM KPool backend requires AscendIndexerKPoolMetadata.") + context = get_forward_context() + if not isinstance(context.attn_metadata, dict): + raise TypeError("GLM KPool backend requires per-layer metadata.") + state_metadata = context.attn_metadata[self.state_cache.prefix] + if not isinstance(state_metadata, AscendIndexerKPoolStateMetadata): + raise TypeError("GLM KPool backend requires compressor-state metadata.") + + num_tokens = hidden_states.shape[0] + if context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + num_tokens = min(num_tokens, indexer_metadata.num_actual_tokens) + hidden = hidden_states[:num_tokens] + k_hidden = k_hidden_states[:num_tokens] + if self._wk_weight_f32 is None: + self.process_weights_after_loading() + assert self._wk_weight_f32 is not None + hidden_f32 = hidden.float() + k_hidden_f32 = hidden_f32 if k_hidden_states is hidden_states else k_hidden.float() + projected = F.linear(k_hidden_f32, self._wk_weight_f32) + k = F.layer_norm( + projected[:, : self.head_dim], + (self.head_dim,), + self._norm_weight_f32, + self._norm_bias_f32, + getattr(self.k_norm, "eps", getattr(self.k_norm, "variance_epsilon", 1e-6)), + ) + gate_score = F.linear(k_hidden_f32, self._gate_weight_f32) + q_values = None + weights = None + if compute_topk: + if isinstance(q_c, tuple): + raise TypeError("GLM KPool backend requires an unquantized q_c tensor.") + q_values = self.wq_b(q_c[:num_tokens])[0].view(num_tokens, self.n_head, self.head_dim) + weights = ( + projected[:, self.head_dim :] + if k_hidden_states is hidden_states + else F.linear(hidden_f32, self._wk_weight_f32[self.head_dim :]) + ).to(q_values.dtype) + weights = weights * (self.softmax_scale * self.n_head**-0.5) + + indexer_cache = self._bound_cache(self.k_cache) + state_cache = self._bound_cache(self.state_cache) + positions = indexer_metadata.positions[:num_tokens] + result = self.indexer_op( + k, + q_values, + weights, + positions, + indexer_cache, + state_cache, + indexer_metadata, + state_metadata, + gate_score=gate_score, + compress_ape=self.index_kpool_compress_ape, + index_kpool=self.index_kpool, + max_pool_seq_len=( + indexer_metadata.block_table.shape[1] * indexer_cache.shape[1] + if context.cudagraph_runtime_mode == CUDAGraphMode.FULL or indexer_metadata.seq_lens_cpu is None + else int(indexer_metadata.seq_lens_cpu.max()) + if indexer_metadata.seq_lens_cpu.numel() + else 0 + ), + compute_topk=compute_topk, + ) + if result is None or self.topk_indices_buffer is None: + return result + + if num_tokens > self.topk_indices_buffer.shape[0]: + raise RuntimeError( + f"GLM KPool output exceeds the top-k buffer rows: {num_tokens} > {self.topk_indices_buffer.shape[0]}." + ) + output = self.topk_indices_buffer[:num_tokens] + output.fill_(-1) + if result.shape[-1] > output.shape[-1]: + raise RuntimeError( + f"GLM KPool output exceeds the top-k buffer width: {result.shape[-1]} > {output.shape[-1]}." + ) + output[:, : result.shape[-1]].copy_(result[:, 0]) + return result diff --git a/vllm_ascend/models/glm5next/attention.py b/vllm_ascend/models/glm5next/attention.py index 1cb56e163821..94eb88256556 100644 --- a/vllm_ascend/models/glm5next/attention.py +++ b/vllm_ascend/models/glm5next/attention.py @@ -33,9 +33,6 @@ Glm5NextIndexerCache, Glm5NextStateCache, ) -from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import ( - SparseAttnIndexerKpool, -) class Indexer(nn.Module): @@ -100,12 +97,6 @@ def __init__( self.scale_fmt = None self.quant_block_size = self.head_dim self.topk_indices_buffer = topk_indices_buffer - # FP32 weight copies for the fp32 K/gate/weights projections. Cached - # lazily after the checkpoint weights are loaded so the quantitative - # work of the first forward matches the last. - self._wk_weight_f32: torch.Tensor | None = None - self._gate_weight_f32: torch.Tensor | None = None - # Completed pools store BF16 vectors without quantization scales. self.k_cache = Glm5NextIndexerCache( head_dim=self.head_dim, @@ -123,23 +114,17 @@ def __init__( cache_config=cache_config, compress_ratio=self.index_kpool, ) - self.max_model_len = vllm_config.model_config.max_model_len self.prefix = prefix - from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size - self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) - self.indexer_op = SparseAttnIndexerKpool( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - tail_cache=self.state_cache, + def get_ascend_indexer_backend_cls(self): + # Lazy import keeps the model module independent from the shared ops + # registry during process startup. + from vllm_ascend.attention.indexer_kpool_backend import ( + Glm5NextKPoolIndexerBackend, ) + return Glm5NextKPoolIndexerBackend + def forward( self, hidden_states: torch.Tensor, @@ -147,64 +132,8 @@ def forward( positions, rotary_emb, ) -> torch.Tensor: - # Match glm-next-0806's kpool contract: the scoring chain consumes a - # raw bf16 Q (no FWHT/fp8 rotation), an fp32 K that survives the norm - # and the compressed-K insertion, raw head weights scaled only by the - # softmax/head factors, and an fp32 gate. FP32 K/gate keep near-tie - # pool rankings stable on long-context tasks. - q, _ = self.wq_b(qr) - q = q.view(-1, self.n_head, self.head_dim) - - hidden_f32 = hidden_states.float() - if self._wk_weight_f32 is None: - self._wk_weight_f32 = self.wk_weights_proj.weight.detach().float() - self._gate_weight_f32 = self.index_kpool_compress_gate.detach().float() - kw = torch.mm(hidden_f32, self._wk_weight_f32.t()) - weights = kw[:, self.head_dim :].to(torch.bfloat16) - k = self.k_norm(kw[:, : self.head_dim]) - - if self.rope_dim > 0: - q_pe, q_nope = torch.split(q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1) - k_pe, k_nope = torch.split( - k, - [self.rope_dim, self.head_dim - self.rope_dim], - dim=-1, - ) - # RoPE runs on the bf16 half; the nope half keeps fp32 so the cache - # write matches the FP32 state the compressor reads back. - k_pe = k_pe.to(torch.bfloat16) - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) - # Note: RoPE (NeoX) can introduce extra leading dimensions during - # compilation so we need to reshape back to token-flattened shapes. - q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) - k_pe = k_pe.reshape(-1, 1, self.rope_dim).float() - # `rotary_emb` is shape-preserving; `q_pe` is already - # [num_tokens, n_head, rope_dim]. - q = torch.cat([q_pe, q_nope], dim=-1) - # `k_pe` is [num_tokens, 1, rope_dim] (MQA); keep the pe rounded to - # bf16 but embed it back into the fp32 K row. - k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) - # else: qk_rope_head_dim=0 — no rope component. q is already - # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all - # nope), so skip the rope split / rotary / cat entirely; otherwise the - # split/reshape would build 0-element tensors (breaks dynamo tracing). - - q = q.to(torch.bfloat16) - weights = weights * (self.softmax_scale * self.n_head**-0.5) - gate_weight_f32 = self._gate_weight_f32 - assert gate_weight_f32 is not None - gate_score = torch.mm(hidden_f32, gate_weight_f32.t()) - - return self.indexer_op( - hidden_states, - q, - k, - weights, - gate_score=gate_score, - compress_ape=self.index_kpool_compress_ape, - index_kpool=self.index_kpool, - positions=positions, - ) + del hidden_states, qr, positions, rotary_emb + raise RuntimeError("GLM-Next Indexer must run through IndexerWrapper's Ascend backend.") class Glm5NextMLAAttention(nn.Module): diff --git a/vllm_ascend/models/glm5next/ops/kpool_compress.py b/vllm_ascend/models/glm5next/ops/kpool_compress.py deleted file mode 100644 index 0688ca038f77..000000000000 --- a/vllm_ascend/models/glm5next/ops/kpool_compress.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Ascend replacements for the GLM-5.3-Flash kpool indexer's fused kernels. - -The upstream implementation fuses the Hadamard-128 rotation with the block-128 -ue8m0 FP8 quantization into a single Triton kernel that keeps the rotated tensor -in registers. That kernel is CUDA-only, so the NPU path expresses the same -numerics in torch: the rotation becomes a matmul against a cached normalized -Hadamard matrix, and the quantization stays elementwise. Every step is -device-side, so the sequence remains ACL-graph safe. -""" - -import torch - -# The indexer query is quantized against the e4m3 maximum, and the resulting -# scale is restricted to a power of two (ue8m0), matching the cached K basis. -FP8_E4M3_MAX = 448.0 - -# Guards rows whose rotated vector is all but zero, so log2 stays finite. -_MIN_ABSMAX = 1e-4 - -_HADAMARD_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} - - -def _normalized_hadamard(dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - """Return a cached ``dim x dim`` Hadamard matrix scaled by ``dim ** -0.5``. - - Folding the normalization into the matrix keeps the rotation a single - matmul. The matrix is symmetric, so it is used without a transpose. - """ - key = (dim, device, dtype) - cached = _HADAMARD_CACHE.get(key) - if cached is not None: - return cached - - try: - from scipy.linalg import hadamard # type: ignore[import-untyped] - except ImportError as err: - raise ImportError( - "The GLM-5.3-Flash kpool indexer requires SciPy for the Hadamard transform. Please install scipy." - ) from err - - matrix = torch.tensor(hadamard(dim, dtype=float), dtype=dtype, device=device) * (dim**-0.5) - _HADAMARD_CACHE[key] = matrix - return matrix - - -def fwht128_quant_fp8(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Rotate each 128-wide row by the Hadamard-128 transform, then FP8-quant. - - Args: - q: ``[rows, 128]`` bf16 -- one head vector per row. - - Returns: - (q_fp8 ``[rows, 128]`` float8_e4m3fn, scale ``[rows, 1]`` float32). - """ - assert q.ndim == 2 and q.shape[1] == 128, q.shape - - rows, dim = q.shape - if rows == 0: - return ( - torch.empty((0, dim), dtype=torch.float8_e4m3fn, device=q.device), - torch.empty((0, 1), dtype=torch.float32, device=q.device), - ) - - hadamard = _normalized_hadamard(dim, q.device, torch.float32) - rotated = q.float() @ hadamard - # The upstream kernel materializes bf16 between the rotation and the quant, - # so the fp8 operand carries the same rounding on both backends. - rotated = rotated.to(torch.bfloat16).to(torch.float32) - - absmax = rotated.abs().amax(dim=-1, keepdim=True).clamp_min(_MIN_ABSMAX) - scale = torch.exp2(torch.ceil(torch.log2(absmax / FP8_E4M3_MAX))) - q_fp8 = (rotated / scale).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX).to(torch.float8_e4m3fn) - return q_fp8, scale diff --git a/vllm_ascend/models/glm5next/sparse_attn_indexer_kpool.py b/vllm_ascend/models/glm5next/sparse_attn_indexer_kpool.py index 390a91c26152..4aaeef3ac390 100644 --- a/vllm_ascend/models/glm5next/sparse_attn_indexer_kpool.py +++ b/vllm_ascend/models/glm5next/sparse_attn_indexer_kpool.py @@ -1,93 +1,141 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Sparse attention indexer layer for the GLM-5.3-Flash kpool indexer. +"""Tensor orchestration for the GLM-Next Triton KPool indexer.""" -GLM-5.3-Flash enables its sparse indexer only when the checkpoint sets -``index_topk``; with ``index_topk`` unset the model runs as dense NoPE MLA plus -KDA, which is the configuration vLLM Ascend currently supports. +from __future__ import annotations -This module keeps the layer so the sparse configuration still builds its index-K -and tail caches (the KV cache specs are wired up and PD transfer works), but the -scoring path is not implemented on Ascend. The upstream implementation is a set -of fused CUDA kernels -- block-FP8 MQA logits through DeepGEMM, paged MQA logits, -and radix top-k over a device workspace -- with no NPU equivalent yet. Rather -than carry that unreachable code, ``forward_oot`` reports the gap so a sparse -checkpoint fails at load with an actionable message instead of dispatching into -a CUDA-only kernel. -""" +from typing import TYPE_CHECKING import torch -from vllm.logger import logger -from vllm.model_executor.custom_op import CustomOp +from torch import nn -_UNSUPPORTED_MESSAGE = ( - "GLM-5.3-Flash sparse (kpool) attention indexing is not implemented on " - "Ascend NPU. The dense NoPE MLA path is supported: serve a checkpoint whose " - "config leaves `index_topk` unset." +from vllm_ascend.ops.triton.glm5_next_kpool_state_compress import ( # type: ignore[import-untyped] + glm5_next_kpool_state_compress_and_write_cache_triton, ) +from vllm_ascend.ops.triton.glm5_next_lightning_indexer import ( # type: ignore[import-untyped] + glm5_next_lightning_indexer_triton, +) + +if TYPE_CHECKING: + from vllm_ascend.attention.indexer_kpool import ( + AscendIndexerKPoolMetadata, + AscendIndexerKPoolStateMetadata, + ) + + +def append_causal_tail( + indices: torch.Tensor, + positions: torch.Tensor, + topk_tokens: int, + pool_size: int, +) -> None: + """Append unpooled tokens to the valid prefix required by CANN SFA.""" + tail_width = pool_size - 1 + if tail_width == 0: + return + positions = positions.to(torch.int64) + tail_start = torch.div(positions + 1, pool_size, rounding_mode="floor") * pool_size + tail_cols = torch.arange(tail_width, device=indices.device, dtype=torch.int64) + tail_tokens = tail_start.unsqueeze(1) + tail_cols + tail_values = torch.where( + tail_cols < (positions + 1 - tail_start).unsqueeze(1), + tail_tokens, + -1, + ).to(indices.dtype) + # PKI packs complete pools at the front. Short requests have fewer than + # topk_tokens history entries; placing the tail at that fixed column would + # leave invalid holes, and SFA would skip the unpooled tokens. + indices[:, topk_tokens:] = -1 + tail_offsets = tail_start.clamp(max=topk_tokens).unsqueeze(1) + tail_cols + indices.scatter_(1, tail_offsets, tail_values) -@CustomOp.register("sparse_attn_indexer_kpool") -class SparseAttnIndexerKpool(CustomOp): - """Sparse attention indexer op for the GLM-5.3-Flash kpool indexer. +class SparseAttnIndexerKpool(nn.Module): + """Update KPool caches and optionally select sparse token indices. - The op is kept as a ``CustomOp`` so the Ascend scoring path can be added - later as a plain ``forward_oot`` implementation, matching how the other - Ascend custom ops are wired up. + Cache binding and forward-context lookup belong to the model-side backend. + This helper receives explicit tensors and typed metadata so the cache update + can be tested independently from the vLLM attention wrapper. """ - def __init__( - self, - k_cache, - quant_block_size: int, - scale_fmt: str | None, - topk_tokens: int, - head_dim: int, - max_model_len: int, - max_total_seq_len: int, - topk_indices_buffer: torch.Tensor, - skip_k_cache_insert: bool = False, - use_fp4_cache: bool = False, - tail_cache=None, - ): + def __init__(self, topk_tokens: int, head_dim: int) -> None: super().__init__() - self.k_cache = k_cache - self.tail_cache = tail_cache - self.quant_block_size = quant_block_size - self.scale_fmt = scale_fmt self.topk_tokens = topk_tokens self.head_dim = head_dim - self.max_model_len = max_model_len - self.max_total_seq_len = max_total_seq_len - self.topk_indices_buffer = topk_indices_buffer - self.skip_k_cache_insert = skip_k_cache_insert - self.use_fp4_cache = use_fp4_cache - logger.warning_once(_UNSUPPORTED_MESSAGE) - def forward_oot( + def forward( self, - hidden_states: torch.Tensor, - q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], k: torch.Tensor, - weights: torch.Tensor, + q_values: torch.Tensor | None, + weights: torch.Tensor | None, + positions: torch.Tensor, + indexer_cache: torch.Tensor, + state_cache: torch.Tensor, + indexer_metadata: AscendIndexerKPoolMetadata, + state_metadata: AscendIndexerKPoolStateMetadata, *, - gate_score: torch.Tensor | None = None, - compress_ape: torch.Tensor | None = None, - index_kpool: int = 1, - positions: torch.Tensor | None = None, - ) -> torch.Tensor: - raise NotImplementedError(_UNSUPPORTED_MESSAGE) + gate_score: torch.Tensor, + compress_ape: torch.Tensor, + index_kpool: int, + max_pool_seq_len: int, + compute_topk: bool, + ) -> torch.Tensor | None: + num_tokens = k.shape[0] + if index_kpool <= 0 or self.topk_tokens % index_kpool: + raise ValueError("KPool top-k must be divisible by its positive pool size.") + if num_tokens == 0: + return ( + None + if not compute_topk + else torch.empty((0, 1, self.topk_tokens + index_kpool - 1), dtype=torch.int32, device=k.device) + ) + if indexer_metadata.cum_query_lens is None or indexer_metadata.raw_seq_lens is None: + raise ValueError("GLM KPool metadata requires cum_query_lens and raw_seq_lens.") + if indexer_cache.dtype != torch.bfloat16: + raise TypeError("GLM KPool compressed cache must be bfloat16.") + if state_cache.dtype != torch.float32 or k.dtype != torch.float32 or gate_score.dtype != torch.float32: + raise TypeError("GLM KPool keys, gates and compressor state must be float32.") + if state_cache.ndim == 4: + state_cache = state_cache.view(state_cache.shape[0], state_cache.shape[1], -1) + if state_cache.shape[-1] != 2 * self.head_dim or gate_score.shape != k.shape: + raise ValueError("GLM KPool state stores one key and gate vector per token.") + if compress_ape.shape != (index_kpool, self.head_dim) or compress_ape.dtype != torch.float32: + raise ValueError("GLM KPool APE must be FP32 with shape [pool_size, head_dim].") - def forward_native( - self, - hidden_states: torch.Tensor, - q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - weights: torch.Tensor, - *, - gate_score: torch.Tensor | None = None, - compress_ape: torch.Tensor | None = None, - index_kpool: int = 1, - positions: torch.Tensor | None = None, - ) -> torch.Tensor: - raise NotImplementedError(_UNSUPPORTED_MESSAGE) + glm5_next_kpool_state_compress_and_write_cache_triton( + state_cache, + indexer_cache, + k, + gate_score, + compress_ape, + positions, + indexer_metadata.cum_query_lens, + indexer_metadata.raw_seq_lens, + state_metadata.slot_mapping[:num_tokens], + state_metadata.block_table, + indexer_metadata.slot_mapping[:num_tokens], + index_kpool, + ) + # Sharing top-k still advances both state caches. + if not compute_topk: + return None + if q_values is None or weights is None: + raise ValueError("GLM KPool top-k requires query and head weights.") + indices = glm5_next_lightning_indexer_triton( + q_values, + indexer_cache, + weights.to(q_values.dtype), + indexer_metadata.cum_query_lens, + indexer_metadata.seq_lens, + indexer_metadata.block_table, + positions, + index_topk=self.topk_tokens, + index_kpool=index_kpool, + max_pool_seq_len=max_pool_seq_len, + ) + # A2/A3 SFA requires a contiguous valid prefix; the reference indexer + # puts the running tail at the fixed top-k column for short requests. + append_causal_tail(indices[:, 0], positions, self.topk_tokens, index_kpool) + valid = torch.arange(num_tokens, device=k.device) < indexer_metadata.cum_query_lens[-1] + indices.masked_fill_(~valid[:, None, None], -1) + return indices diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 1064a3514bb3..09975694cc85 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -228,15 +228,6 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: use_compress = getattr(attn_selector_config, "use_compress", False) use_mla = attn_selector_config.use_mla use_sparse = attn_selector_config.use_sparse - # index_kpool GLM is not DeepSeek SFA; keep MLA backend. - try: - from vllm.config import get_current_vllm_config - from vllm_ascend.utils import enable_sfa - - if use_sparse and not enable_sfa(get_current_vllm_config()): - use_sparse = False - except Exception: - pass key = (use_mla, use_sparse) backend_key = (*key, use_compress) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index 95c5c517d659..19c7a57b9dbe 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -122,8 +122,8 @@ def model_uses_kpool_indexer(model_config: Any | None) -> bool: """Return True for GLM-5.3-Flash style kpool indexer models. Those models expose ``index_topk`` like DeepSeek SFA but use a kpool - indexer over a hybrid MLA + KDA cache, so they must not be routed through - the SFA / DSA layouts. + indexer over a hybrid MLA + KDA cache. Its cache geometry is + independent of the LightningIndexer layout used by other SFA models. """ return any(hasattr(getattr(model_config, attr, None), "index_kpool") for attr in ("hf_text_config", "hf_config")) From b0fff61cfb3bbb3b316c07373f6844691290b9d7 Mon Sep 17 00:00:00 2001 From: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:26:47 +0800 Subject: [PATCH 2/3] refactor(attention): consolidate KPool indexer backend Keep KPool cache metadata and model execution in one attention module. Update model and test imports and defer KPool operator loading until backend initialization. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> --- .../test_glm5next_kpool_model_backend.py | 4 +- tests/ut/models/test_glm5next_sfa_routing.py | 14 +- vllm_ascend/attention/indexer_kpool.py | 177 ++++++++++++++++- .../attention/indexer_kpool_backend.py | 187 ------------------ vllm_ascend/models/glm5next/attention.py | 2 +- 5 files changed, 181 insertions(+), 203 deletions(-) delete mode 100644 vllm_ascend/attention/indexer_kpool_backend.py diff --git a/tests/ut/models/test_glm5next_kpool_model_backend.py b/tests/ut/models/test_glm5next_kpool_model_backend.py index b46632fa26ae..96d5eab71c57 100644 --- a/tests/ut/models/test_glm5next_kpool_model_backend.py +++ b/tests/ut/models/test_glm5next_kpool_model_backend.py @@ -8,9 +8,11 @@ import torch from torch import nn +import vllm_ascend.attention.indexer_kpool as backend_module from vllm_ascend.attention.indexer_kpool import ( AscendIndexerKPoolMetadata, AscendIndexerKPoolStateMetadata, + Glm5NextKPoolIndexerBackend, ) # The operator PR owns these kernels. Each orchestration test below replaces @@ -22,9 +24,7 @@ "vllm_ascend.ops.triton.glm5_next_lightning_indexer": MagicMock(), }, ): - import vllm_ascend.attention.indexer_kpool_backend as backend_module import vllm_ascend.models.glm5next.sparse_attn_indexer_kpool as kpool_module - from vllm_ascend.attention.indexer_kpool_backend import Glm5NextKPoolIndexerBackend from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import SparseAttnIndexerKpool, append_causal_tail diff --git a/tests/ut/models/test_glm5next_sfa_routing.py b/tests/ut/models/test_glm5next_sfa_routing.py index 01c6cf1b5d5d..a6a9933ab6b8 100644 --- a/tests/ut/models/test_glm5next_sfa_routing.py +++ b/tests/ut/models/test_glm5next_sfa_routing.py @@ -2,26 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace -from unittest.mock import MagicMock, patch import pytest import torch import vllm.config as vllm_config_module +import vllm_ascend.attention.indexer_kpool as backend_module import vllm_ascend.platform as platform +from vllm_ascend.attention.indexer_kpool import Glm5NextKPoolIndexerBackend from vllm_ascend.device.hardware_profile import AttentionBackendFamily, DeviceAdaptorFamily -# Routing and visible-length checks do not execute the external kernels. -with patch.dict( - "sys.modules", - { - "vllm_ascend.ops.triton.glm5_next_kpool_state_compress": MagicMock(), - "vllm_ascend.ops.triton.glm5_next_lightning_indexer": MagicMock(), - }, -): - import vllm_ascend.attention.indexer_kpool_backend as backend_module - from vllm_ascend.attention.indexer_kpool_backend import Glm5NextKPoolIndexerBackend - @pytest.mark.parametrize("kpool", [False, True]) @pytest.mark.parametrize("fp8_device", [False, True]) diff --git a/vllm_ascend/attention/indexer_kpool.py b/vllm_ascend/attention/indexer_kpool.py index efe52e621668..6870db677401 100644 --- a/vllm_ascend/attention/indexer_kpool.py +++ b/vllm_ascend/attention/indexer_kpool.py @@ -1,11 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Independent pooled-indexer and compressor-state metadata for GLM-Next.""" +"""Cache metadata and execution backends for the GLM-Next pooled indexer.""" from dataclasses import dataclass +from typing import Any import torch +import torch.nn.functional as F +from torch import nn from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import get_forward_context from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import ( AttentionBackend, @@ -16,11 +21,13 @@ ) from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm_ascend.attention.indexer import AscendSFAIndexerBackend from vllm_ascend.core.kv_cache_interface import ( AscendIndexerKPoolStateSpec, get_kv_cache_compression_ratio, get_storage_block_size, ) +from vllm_ascend.device.hardware_profile import AttentionBackendFamily, get_current_hardware_profile from vllm_ascend.models.glm5next.kv_cache import ( format_indexer_kpool_slot_mapping, ) @@ -326,3 +333,171 @@ def get_kv_cache_shape( if num_kv_heads != 1: raise ValueError(f"Indexer KPool state cache requires one KV head, got {num_kv_heads}.") return (num_blocks, block_size, head_size) + + +class Glm5NextKPoolIndexerBackend(AscendSFAIndexerBackend): + """Triton KPool implementation of the unified seven-argument indexer API.""" + + def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: + nn.Module.__init__(self) + if qk_rope_head_dim != 0: + raise ValueError( + f"GLM-Next KPool indexing supports NoPE queries only, got qk_rope_head_dim={qk_rope_head_dim}." + ) + parallel_config = vllm_indexer.vllm_config.parallel_config + if parallel_config.prefill_context_parallel_size > 1 or parallel_config.decode_context_parallel_size > 1: + raise NotImplementedError("GLM-Next KPool indexing does not support PCP or DCP.") + + if get_current_hardware_profile().attention_backend_family is AttentionBackendFamily.COMPATIBILITY: + raise NotImplementedError("KPool sparse attention requires Ascend A2, A3 or A5.") + + self.n_head: int = vllm_indexer.n_head + self.head_dim: int = vllm_indexer.head_dim + self.topk_tokens: int = vllm_indexer.topk_tokens + self.q_lora_rank: int = vllm_indexer.q_lora_rank + self.index_kpool: int = vllm_indexer.index_kpool + self.wq_b = vllm_indexer.wq_b + self.wk_weights_proj = vllm_indexer.wk_weights_proj + self.k_norm = vllm_indexer.k_norm + self.softmax_scale = vllm_indexer.softmax_scale + self.index_kpool_compress_ape = vllm_indexer.index_kpool_compress_ape + self.index_kpool_compress_gate = vllm_indexer.index_kpool_compress_gate + self.k_cache: Any = vllm_indexer.k_cache + self.state_cache: Any = vllm_indexer.state_cache + self.topk_indices_buffer: torch.Tensor | None = vllm_indexer.topk_indices_buffer + # Load KPool operators only when constructing the model-side backend; + # cache metadata is also imported during engine initialization. + from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import SparseAttnIndexerKpool + + self.indexer_op = SparseAttnIndexerKpool(self.topk_tokens, self.head_dim) + self.enable_sparse_li_c8 = False + for name in ("_wk_weight_f32", "_gate_weight_f32", "_norm_weight_f32", "_norm_bias_f32"): + self.register_buffer(name, None, persistent=False) + + @property + def topk_output_width(self) -> int: + return self.topk_tokens + self.index_kpool - 1 + + def get_topk_lengths(self, positions: torch.Tensor) -> torch.Tensor: + visible = (positions + 1).clamp_min(0) + history = (visible // self.index_kpool * self.index_kpool).clamp(max=self.topk_tokens) + return history + visible % self.index_kpool + + @property + def num_cache_tensors(self) -> int: + return 1 + + def process_weights_after_loading(self) -> None: + self._wk_weight_f32 = self.wk_weights_proj.weight.detach().float() + self._gate_weight_f32 = self.index_kpool_compress_gate.detach().float() + self._norm_weight_f32 = self.k_norm.weight.detach().float() if self.k_norm.weight is not None else None + self._norm_bias_f32 = self.k_norm.bias.detach().float() if self.k_norm.bias is not None else None + + @staticmethod + def _bound_cache(layer: Any) -> torch.Tensor: + context = get_forward_context() + cache = layer.kv_cache + if isinstance(cache, (list, tuple)): + virtual_engine = getattr(context, "virtual_engine", 0) or 0 + if virtual_engine >= len(cache): + raise IndexError(f"Cache virtual engine {virtual_engine} is out of range.") + cache = cache[virtual_engine] + if isinstance(cache, (list, tuple)): + if len(cache) != 1: + raise TypeError("GLM KPool cache must contain one tensor.") + cache = cache[0] + if not isinstance(cache, torch.Tensor): + raise TypeError(f"GLM KPool cache {layer.prefix!r} is not bound to a tensor.") + return cache + + def forward( + self, + hidden_states: torch.Tensor, + q_c: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + cos: torch.Tensor | None, + sin: torch.Tensor | None, + k_hidden_states: torch.Tensor, + indexer_metadata: Any, + compute_topk: bool = True, + ) -> torch.Tensor | None: + del cos, sin + if not isinstance(indexer_metadata, AscendIndexerKPoolMetadata): + raise TypeError("GLM KPool backend requires AscendIndexerKPoolMetadata.") + context = get_forward_context() + if not isinstance(context.attn_metadata, dict): + raise TypeError("GLM KPool backend requires per-layer metadata.") + state_metadata = context.attn_metadata[self.state_cache.prefix] + if not isinstance(state_metadata, AscendIndexerKPoolStateMetadata): + raise TypeError("GLM KPool backend requires compressor-state metadata.") + + num_tokens = hidden_states.shape[0] + if context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + num_tokens = min(num_tokens, indexer_metadata.num_actual_tokens) + hidden = hidden_states[:num_tokens] + k_hidden = k_hidden_states[:num_tokens] + if self._wk_weight_f32 is None: + self.process_weights_after_loading() + assert self._wk_weight_f32 is not None + hidden_f32 = hidden.float() + k_hidden_f32 = hidden_f32 if k_hidden_states is hidden_states else k_hidden.float() + projected = F.linear(k_hidden_f32, self._wk_weight_f32) + k = F.layer_norm( + projected[:, : self.head_dim], + (self.head_dim,), + self._norm_weight_f32, + self._norm_bias_f32, + getattr(self.k_norm, "eps", getattr(self.k_norm, "variance_epsilon", 1e-6)), + ) + gate_score = F.linear(k_hidden_f32, self._gate_weight_f32) + q_values = None + weights = None + if compute_topk: + if isinstance(q_c, tuple): + raise TypeError("GLM KPool backend requires an unquantized q_c tensor.") + q_values = self.wq_b(q_c[:num_tokens])[0].view(num_tokens, self.n_head, self.head_dim) + weights = ( + projected[:, self.head_dim :] + if k_hidden_states is hidden_states + else F.linear(hidden_f32, self._wk_weight_f32[self.head_dim :]) + ).to(q_values.dtype) + weights = weights * (self.softmax_scale * self.n_head**-0.5) + + indexer_cache = self._bound_cache(self.k_cache) + state_cache = self._bound_cache(self.state_cache) + positions = indexer_metadata.positions[:num_tokens] + result = self.indexer_op( + k, + q_values, + weights, + positions, + indexer_cache, + state_cache, + indexer_metadata, + state_metadata, + gate_score=gate_score, + compress_ape=self.index_kpool_compress_ape, + index_kpool=self.index_kpool, + max_pool_seq_len=( + indexer_metadata.block_table.shape[1] * indexer_cache.shape[1] + if context.cudagraph_runtime_mode == CUDAGraphMode.FULL or indexer_metadata.seq_lens_cpu is None + else int(indexer_metadata.seq_lens_cpu.max()) + if indexer_metadata.seq_lens_cpu.numel() + else 0 + ), + compute_topk=compute_topk, + ) + if result is None or self.topk_indices_buffer is None: + return result + + if num_tokens > self.topk_indices_buffer.shape[0]: + raise RuntimeError( + f"GLM KPool output exceeds the top-k buffer rows: {num_tokens} > {self.topk_indices_buffer.shape[0]}." + ) + output = self.topk_indices_buffer[:num_tokens] + output.fill_(-1) + if result.shape[-1] > output.shape[-1]: + raise RuntimeError( + f"GLM KPool output exceeds the top-k buffer width: {result.shape[-1]} > {output.shape[-1]}." + ) + output[:, : result.shape[-1]].copy_(result[:, 0]) + return result diff --git a/vllm_ascend/attention/indexer_kpool_backend.py b/vllm_ascend/attention/indexer_kpool_backend.py deleted file mode 100644 index f20c14555671..000000000000 --- a/vllm_ascend/attention/indexer_kpool_backend.py +++ /dev/null @@ -1,187 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Attention backend for the GLM-Next pooled sparse indexer.""" - -from __future__ import annotations - -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn -from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import get_forward_context - -from vllm_ascend.attention.indexer import AscendSFAIndexerBackend -from vllm_ascend.attention.indexer_kpool import ( - AscendIndexerKPoolMetadata, - AscendIndexerKPoolStateMetadata, -) -from vllm_ascend.device.hardware_profile import AttentionBackendFamily, get_current_hardware_profile -from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import ( - SparseAttnIndexerKpool, -) - - -class Glm5NextKPoolIndexerBackend(AscendSFAIndexerBackend): - """Triton KPool implementation of the unified seven-argument indexer API.""" - - def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: - nn.Module.__init__(self) - if qk_rope_head_dim != 0: - raise ValueError( - f"GLM-Next KPool indexing supports NoPE queries only, got qk_rope_head_dim={qk_rope_head_dim}." - ) - parallel_config = vllm_indexer.vllm_config.parallel_config - if parallel_config.prefill_context_parallel_size > 1 or parallel_config.decode_context_parallel_size > 1: - raise NotImplementedError("GLM-Next KPool indexing does not support PCP or DCP.") - - if get_current_hardware_profile().attention_backend_family is AttentionBackendFamily.COMPATIBILITY: - raise NotImplementedError("KPool sparse attention requires Ascend A2, A3 or A5.") - - self.n_head: int = vllm_indexer.n_head - self.head_dim: int = vllm_indexer.head_dim - self.topk_tokens: int = vllm_indexer.topk_tokens - self.q_lora_rank: int = vllm_indexer.q_lora_rank - self.index_kpool: int = vllm_indexer.index_kpool - self.wq_b = vllm_indexer.wq_b - self.wk_weights_proj = vllm_indexer.wk_weights_proj - self.k_norm = vllm_indexer.k_norm - self.softmax_scale = vllm_indexer.softmax_scale - self.index_kpool_compress_ape = vllm_indexer.index_kpool_compress_ape - self.index_kpool_compress_gate = vllm_indexer.index_kpool_compress_gate - self.k_cache: Any = vllm_indexer.k_cache - self.state_cache: Any = vllm_indexer.state_cache - self.topk_indices_buffer: torch.Tensor | None = vllm_indexer.topk_indices_buffer - self.indexer_op = SparseAttnIndexerKpool(self.topk_tokens, self.head_dim) - self.enable_sparse_li_c8 = False - for name in ("_wk_weight_f32", "_gate_weight_f32", "_norm_weight_f32", "_norm_bias_f32"): - self.register_buffer(name, None, persistent=False) - - @property - def topk_output_width(self) -> int: - return self.topk_tokens + self.index_kpool - 1 - - def get_topk_lengths(self, positions: torch.Tensor) -> torch.Tensor: - visible = (positions + 1).clamp_min(0) - history = (visible // self.index_kpool * self.index_kpool).clamp(max=self.topk_tokens) - return history + visible % self.index_kpool - - @property - def num_cache_tensors(self) -> int: - return 1 - - def process_weights_after_loading(self) -> None: - self._wk_weight_f32 = self.wk_weights_proj.weight.detach().float() - self._gate_weight_f32 = self.index_kpool_compress_gate.detach().float() - self._norm_weight_f32 = self.k_norm.weight.detach().float() if self.k_norm.weight is not None else None - self._norm_bias_f32 = self.k_norm.bias.detach().float() if self.k_norm.bias is not None else None - - @staticmethod - def _bound_cache(layer: Any) -> torch.Tensor: - context = get_forward_context() - cache = layer.kv_cache - if isinstance(cache, (list, tuple)): - virtual_engine = getattr(context, "virtual_engine", 0) or 0 - if virtual_engine >= len(cache): - raise IndexError(f"Cache virtual engine {virtual_engine} is out of range.") - cache = cache[virtual_engine] - if isinstance(cache, (list, tuple)): - if len(cache) != 1: - raise TypeError("GLM KPool cache must contain one tensor.") - cache = cache[0] - if not isinstance(cache, torch.Tensor): - raise TypeError(f"GLM KPool cache {layer.prefix!r} is not bound to a tensor.") - return cache - - def forward( - self, - hidden_states: torch.Tensor, - q_c: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - cos: torch.Tensor | None, - sin: torch.Tensor | None, - k_hidden_states: torch.Tensor, - indexer_metadata: Any, - compute_topk: bool = True, - ) -> torch.Tensor | None: - del cos, sin - if not isinstance(indexer_metadata, AscendIndexerKPoolMetadata): - raise TypeError("GLM KPool backend requires AscendIndexerKPoolMetadata.") - context = get_forward_context() - if not isinstance(context.attn_metadata, dict): - raise TypeError("GLM KPool backend requires per-layer metadata.") - state_metadata = context.attn_metadata[self.state_cache.prefix] - if not isinstance(state_metadata, AscendIndexerKPoolStateMetadata): - raise TypeError("GLM KPool backend requires compressor-state metadata.") - - num_tokens = hidden_states.shape[0] - if context.cudagraph_runtime_mode != CUDAGraphMode.FULL: - num_tokens = min(num_tokens, indexer_metadata.num_actual_tokens) - hidden = hidden_states[:num_tokens] - k_hidden = k_hidden_states[:num_tokens] - if self._wk_weight_f32 is None: - self.process_weights_after_loading() - assert self._wk_weight_f32 is not None - hidden_f32 = hidden.float() - k_hidden_f32 = hidden_f32 if k_hidden_states is hidden_states else k_hidden.float() - projected = F.linear(k_hidden_f32, self._wk_weight_f32) - k = F.layer_norm( - projected[:, : self.head_dim], - (self.head_dim,), - self._norm_weight_f32, - self._norm_bias_f32, - getattr(self.k_norm, "eps", getattr(self.k_norm, "variance_epsilon", 1e-6)), - ) - gate_score = F.linear(k_hidden_f32, self._gate_weight_f32) - q_values = None - weights = None - if compute_topk: - if isinstance(q_c, tuple): - raise TypeError("GLM KPool backend requires an unquantized q_c tensor.") - q_values = self.wq_b(q_c[:num_tokens])[0].view(num_tokens, self.n_head, self.head_dim) - weights = ( - projected[:, self.head_dim :] - if k_hidden_states is hidden_states - else F.linear(hidden_f32, self._wk_weight_f32[self.head_dim :]) - ).to(q_values.dtype) - weights = weights * (self.softmax_scale * self.n_head**-0.5) - - indexer_cache = self._bound_cache(self.k_cache) - state_cache = self._bound_cache(self.state_cache) - positions = indexer_metadata.positions[:num_tokens] - result = self.indexer_op( - k, - q_values, - weights, - positions, - indexer_cache, - state_cache, - indexer_metadata, - state_metadata, - gate_score=gate_score, - compress_ape=self.index_kpool_compress_ape, - index_kpool=self.index_kpool, - max_pool_seq_len=( - indexer_metadata.block_table.shape[1] * indexer_cache.shape[1] - if context.cudagraph_runtime_mode == CUDAGraphMode.FULL or indexer_metadata.seq_lens_cpu is None - else int(indexer_metadata.seq_lens_cpu.max()) - if indexer_metadata.seq_lens_cpu.numel() - else 0 - ), - compute_topk=compute_topk, - ) - if result is None or self.topk_indices_buffer is None: - return result - - if num_tokens > self.topk_indices_buffer.shape[0]: - raise RuntimeError( - f"GLM KPool output exceeds the top-k buffer rows: {num_tokens} > {self.topk_indices_buffer.shape[0]}." - ) - output = self.topk_indices_buffer[:num_tokens] - output.fill_(-1) - if result.shape[-1] > output.shape[-1]: - raise RuntimeError( - f"GLM KPool output exceeds the top-k buffer width: {result.shape[-1]} > {output.shape[-1]}." - ) - output[:, : result.shape[-1]].copy_(result[:, 0]) - return result diff --git a/vllm_ascend/models/glm5next/attention.py b/vllm_ascend/models/glm5next/attention.py index 94eb88256556..14e3dc53f3b4 100644 --- a/vllm_ascend/models/glm5next/attention.py +++ b/vllm_ascend/models/glm5next/attention.py @@ -119,7 +119,7 @@ def __init__( def get_ascend_indexer_backend_cls(self): # Lazy import keeps the model module independent from the shared ops # registry during process startup. - from vllm_ascend.attention.indexer_kpool_backend import ( + from vllm_ascend.attention.indexer_kpool import ( Glm5NextKPoolIndexerBackend, ) From 1ff19ad48e546b68ef06a81810ad6380903c5754 Mon Sep 17 00:00:00 2001 From: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:33:58 +0800 Subject: [PATCH 3/3] fix(attention): keep KPool metadata imports independent Use the self-contained KPool execution module without inheriting the generic SFA indexer. This avoids the DeviceOperator import cycle during metadata-only startup. Add regression coverage with execution dependencies unavailable. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com> --- .../test_glm5next_kpool_model_backend.py | 20 +++++++++++++++++++ vllm_ascend/attention/indexer_kpool.py | 11 ++++++---- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/tests/ut/models/test_glm5next_kpool_model_backend.py b/tests/ut/models/test_glm5next_kpool_model_backend.py index 96d5eab71c57..0ef4730556c2 100644 --- a/tests/ut/models/test_glm5next_kpool_model_backend.py +++ b/tests/ut/models/test_glm5next_kpool_model_backend.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.util from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -28,6 +29,25 @@ from vllm_ascend.models.glm5next.sparse_attn_indexer_kpool import SparseAttnIndexerKpool, append_causal_tail +def test_cache_metadata_import_does_not_require_indexer_operators() -> None: + module_name = f"{backend_module.__name__}_import_test" + spec = importlib.util.spec_from_file_location(module_name, backend_module.__file__) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + with patch.dict( + "sys.modules", + { + module_name: module, + "vllm_ascend.attention.indexer": None, + "vllm_ascend.models.glm5next.sparse_attn_indexer_kpool": None, + }, + ): + spec.loader.exec_module(module) + + assert module.AscendIndexerKPoolBackend.get_builder_cls() is module.AscendIndexerKPoolMetadataBuilder + assert module.AscendIndexerKPoolStateBackend.get_builder_cls() is module.AscendIndexerKPoolStateMetadataBuilder + + @pytest.mark.parametrize("pool_size", [1, 4]) def test_causal_tail_follows_valid_history_without_holes(pool_size: int) -> None: topk = 8 diff --git a/vllm_ascend/attention/indexer_kpool.py b/vllm_ascend/attention/indexer_kpool.py index 6870db677401..39bee8ec1ce4 100644 --- a/vllm_ascend/attention/indexer_kpool.py +++ b/vllm_ascend/attention/indexer_kpool.py @@ -21,7 +21,6 @@ ) from vllm.v1.kv_cache_interface import MLAAttentionSpec -from vllm_ascend.attention.indexer import AscendSFAIndexerBackend from vllm_ascend.core.kv_cache_interface import ( AscendIndexerKPoolStateSpec, get_kv_cache_compression_ratio, @@ -335,11 +334,15 @@ def get_kv_cache_shape( return (num_blocks, block_size, head_size) -class Glm5NextKPoolIndexerBackend(AscendSFAIndexerBackend): - """Triton KPool implementation of the unified seven-argument indexer API.""" +class Glm5NextKPoolIndexerBackend(nn.Module): + """Model-side implementation of the unified seven-argument indexer API. + + The cache-only backends above own the engine-facing attention contracts. + Keep execution independent of the standard SFA indexer and its operators. + """ def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: - nn.Module.__init__(self) + super().__init__() if qk_rope_head_dim != 0: raise ValueError( f"GLM-Next KPool indexing supports NoPE queries only, got qk_rope_head_dim={qk_rope_head_dim}."