From ca8744e09850e4cab18c7a09f7e6c075a8de97a1 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 16:42:48 +0800 Subject: [PATCH 001/120] feat(hisparse): add Ascend backend foundation Split the unquantized NPU data path from PR19109 onto current Ascend main. Reuse merged Core backend, HMA, Coordinator and lifetime interfaces, expose device cache operations and view hooks, and document independent contribution areas. Exclude C8 tuning, P/D transport, native optimized kernels and experiment profiles. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 123 ++++ mkdocs.yml | 1 + .../one_card/test_hisparse_attention.py | 105 +++ .../one_card/test_hisparse_kv_preprocess.py | 204 ++++++ .../one_card/test_hisparse_ops.py | 626 +++++++++++++++++ .../one_card/test_hisparse_prefill_staging.py | 115 +++ .../one_card/test_hisparse_selection_order.py | 85 +++ .../test_hisparse_foundation_cpu.py | 238 +++++++ tests/ut/ops/test_mla_weight_hook_cpu.py | 54 ++ .../test_hisparse_capture_mirror_cpu.py | 70 ++ vllm_ascend/attention/indexer.py | 45 ++ vllm_ascend/attention/sfa_hisparse.py | 333 +++++++++ vllm_ascend/attention/sfa_v1.py | 8 + .../kv_transfer/hisparse/__init__.py | 3 + .../kv_transfer/hisparse/backend.py | 90 +++ .../kv_transfer/hisparse/binding.py | 72 ++ vllm_ascend/ops/hisparse.py | 660 ++++++++++++++++++ vllm_ascend/ops/mla.py | 18 +- vllm_ascend/ops/triton/hisparse.py | 546 +++++++++++++++ vllm_ascend/patch/__init__.py | 11 + .../platform/patch_kv_cache_coordinator.py | 6 + vllm_ascend/patch/worker/patch_deepseek_v2.py | 3 +- .../patch/worker/patch_hisparse_streams.py | 141 ++++ vllm_ascend/platform.py | 4 + vllm_ascend/worker/v2/attn_utils.py | 15 + vllm_ascend/worker/worker.py | 3 + 26 files changed, 3572 insertions(+), 7 deletions(-) create mode 100644 docs/source/user_guide/feature_guide/hisparse.md create mode 100644 tests/e2e/pull_request/one_card/test_hisparse_attention.py create mode 100644 tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py create mode 100644 tests/e2e/pull_request/one_card/test_hisparse_ops.py create mode 100644 tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py create mode 100644 tests/e2e/pull_request/one_card/test_hisparse_selection_order.py create mode 100644 tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py create mode 100644 tests/ut/ops/test_mla_weight_hook_cpu.py create mode 100644 tests/ut/worker/test_hisparse_capture_mirror_cpu.py create mode 100644 vllm_ascend/attention/sfa_hisparse.py create mode 100644 vllm_ascend/distributed/kv_transfer/hisparse/__init__.py create mode 100644 vllm_ascend/distributed/kv_transfer/hisparse/backend.py create mode 100644 vllm_ascend/distributed/kv_transfer/hisparse/binding.py create mode 100644 vllm_ascend/ops/hisparse.py create mode 100644 vllm_ascend/ops/triton/hisparse.py create mode 100644 vllm_ascend/patch/worker/patch_hisparse_streams.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md new file mode 100644 index 000000000000..5db0f473da96 --- /dev/null +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -0,0 +1,123 @@ +# HiSparse foundation on Ascend + +## Scope + +This is the foundational split of [#19109](https://github.com/vllm-project/vllm-ascend/pull/19109). +It supplies an Ascend backend for the merged vLLM HiSparse API and the minimal +SFA data path. It is a contribution starting point, with new-head NPU and full-model +qualification pending. Historical optimized C8 performance does not qualify this split. + +Core owns HMA allocation, Coordinator admission/eviction, prefix-cache state, +request lifetimes, prefetch control flow, mirror completion and release/rollback. +Ascend implements device operations and cache-layout adaptation. Attention uses +the existing CANN sparse-flash-attention and fused RMSNorm/RoPE operations; +this PR adds no Attention core or native C++/AscendC kernel. + +```text +Core sparse index group -> logical TopK -> NPU residency/refill plan + | + Resident / Hot hit / mapped-Host miss + | + packed KV -> existing native SFA operation +``` + +This first data path covers unquantized main/indexer KV with the RoPE SFA +layout and Model Runner V2. Model-weight quantization is independent. Private +Host pools are used per rank because NPU IPC events are not implemented. +The mapped-Host capability is checked at registration; hardware evidence for +the original primitives is from Ascend 950. A3 and the new source pair still +need their own qualification. + +The SFA kernel requires TopK divisible by 128. The resolver's int16 LRU format +limits aligned Hot capacity to 32768 rows, and Hot must cover the complete +verification-row union. Context parallelism is not implemented in this split. +Request counts, lengths, layer/group counts and scratch capacities come from +configuration, rather than benchmark-specific constants. + +## Configuration + +Use the matching merged Core pin in `.github/vllm-main-verified.commit`. +No unmerged Core PR, copied scheduler, experiment profile or source-file hash +allowlist is required. As with other Ascend backends, paired-version ABI changes +still need review and device validation. + +Set `VLLM_USE_V2_MODEL_RUNNER=1` and add these options to a compatible model's +normal launch command: + +```bash +--block-size 128 \ +--kv-cache-dtype bfloat16 \ +--attention-config '{"indexer_kv_dtype":"bf16","hisparse_config":{"eager_host_mirror":true}}' \ +--kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":4}}' +``` + +The example reserves 4 GiB Host storage per rank. Size it for the model and context +budget; it is not a model capacity claim. Keep Core's eager mirror policy for +this foundation. External P/D transfer is not provided by this PR. + +## Extension points for contributors + +Use existing interfaces, rather than adding unused placeholder hooks: + +| Boundary | Current consumer and responsibility | Follow-up work | +| --- | --- | --- | +| `NPUPlatform.get_hisparse_backend_cls` | Core selects the device backend only for HiSparse | Backend capability/version support | +| `HiSparseBackend.create_cache_ops` | Runtime invokes resolver, refill and invalidation operations | Faster resolver/refill, vectorization and balancing | +| `validate_kv_cache_config` / `prepare_kv_cache_views` | Core validates before allocation and binds views inside rollback scope | Packed/quantized KV formats with layout tests | +| Backend `pin_tensor` / `unpin_tensor` / `swap_blocks_batch` | Core owns completion, explicit registration release and transfer lifetime | DMA implementation and supported memory capabilities | +| `AscendHiSparseMetadataBuilder` / `AscendHiSparseImpl` | SFA factories select metadata and the layout bridge | C8 and a smaller layout-conversion cost | + +`ops/hisparse.py` owns operation validation and persistent per-group workspace; +`ops/triton/hisparse.py` implements the reference NPU path. Followers share their +leader's scratch only on that group's ordered stream. Independent groups and +current request/state mappings remain separate. + +The stream facade is an isolated PTA compatibility adapter in the two common +HiSparse modules. It preserves event placement and does not replace the actual +global Torch module. Remove it when generic PTA capture/workspace behavior is +equivalent or Core provides the matching factories. + +Follow-up contributions should be separate and measurable: packed C8/indexer +formats, native resolver/refill improvements, Host-writeback optimization, metadata +reuse, additional hardware and P/D transport. PR19109 retains these experiments. +AVSG, lazy-MTP policy, native optimized kernels, private Core adapters and +machine-specific benchmarking tools are excluded here. + +Preserve completion-before-unregister, graph allocation owners, invalidation on +overwrite and protection of every selected token before choosing an LRU victim. +A same-shaped batch is not permission to reuse old lengths or request mappings. + +## Validation + +Portable, source-isolated CPU boundary checks: + +```bash +python -m pytest -q --noconftest \ + tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py \ + tests/ut/ops/test_mla_weight_hook_cpu.py \ + tests/ut/worker/test_hisparse_capture_mirror_cpu.py +``` + +These check current Host offsets, failed-unregister ownership, group scratch +isolation, strided indexer aliases, planner alignment, unsupported formats, +Core coordinator forwarding, loader ABI and capture-only mirror suppression. +They do not exercise an installed NPU package or qualify device execution. + +Real-NPU cases are included for resolver/refill/overwrite invalidation, independent +FP32 Attention references, selection order/duplicates, native BF16 preprocessing +and Resident/Host prefill staging, in eager and changing-graph cases: + +```bash +pytest --confcutdir=tests/e2e/pull_request/one_card \ + tests/e2e/pull_request/one_card/test_hisparse_ops.py \ + tests/e2e/pull_request/one_card/test_hisparse_attention.py \ + tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py \ + tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py \ + tests/e2e/pull_request/one_card/test_hisparse_selection_order.py +``` + +These cases require an appropriate NPU environment and are not claimed to have +passed on this new split. Fresh full-model correctness, dataset accuracy, +eager/Graph lifecycle and same-workload performance remain prerequisites for +marking the foundation ready. Existing C8 benchmark results are intentionally +not presented as this implementation's performance. diff --git a/mkdocs.yml b/mkdocs.yml index e0f3f0526339..59544a078608 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -289,6 +289,7 @@ nav: - user_guide/configuration/encoder_cache_manager.md - Feature Guide: - user_guide/feature_guide/index.md + - user_guide/feature_guide/hisparse.md - user_guide/feature_guide/graph_mode.md - user_guide/feature_guide/cpu_binding.md - user_guide/feature_guide/Ai_QoS_introduction_en.md diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py new file mode 100644 index 000000000000..192e73e8f73e --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check the actual SFA bridge against an independent FP32 CPU reference.""" + +import math +from functools import partial +from types import SimpleNamespace + +import pytest +import torch +import torch_npu # noqa: F401 + + +@pytest.fixture(autouse=True) +def _require_ascend950_sfa(): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse SFA integration currently targets Ascend 950") + torch.npu.set_device(0) + from vllm_ascend.utils import bootstrap_custom_op_env + + # Load the required SFA extension without enabling all A5 custom paths. + bootstrap_custom_op_env() + import vllm_ascend.vllm_ascend_C # noqa: F401 + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("num_queries", [1, 7, 35]) +def test_sparse_attention_reference(num_queries, strided, capture): + from vllm_ascend.attention.sfa_hisparse import _ATTENTION_CHUNK_ROWS, AscendHiSparseImpl + + torch.manual_seed(47) + device = torch.device("npu:0") + topk, latent_dim, rope_dim, heads = 2048, 512, 64, 8 + block_size, blocks = 128, 16 + capacity = _ATTENTION_CHUNK_ROWS + scale = 1 / math.sqrt(latent_dim + rope_dim) + workspace = ( + torch.empty((capacity, topk, 1, latent_dim), device=device, dtype=torch.bfloat16), + torch.empty((capacity, topk, 1, rope_dim), device=device, dtype=torch.bfloat16), + torch.arange(capacity * topk // block_size, dtype=torch.int32, device=device).view(capacity, -1), + torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), + torch.arange(1, capacity + 1, dtype=torch.int32, device=device), + torch.full((capacity,), topk, dtype=torch.int32, device=device), + ) + impl = SimpleNamespace(kv_lora_rank=latent_dim, qk_rope_head_dim=rope_dim, scale=scale) + impl._workspace = lambda dtype: workspace + # Quantize before taking the FP32 reference so both paths see the same inputs. + source = torch.randn((blocks, block_size, latent_dim + rope_dim)).to(torch.bfloat16) + query = torch.randn((num_queries, heads, latent_dim)).to(torch.bfloat16) + query_rope = torch.randn((num_queries, heads, rope_dim)).to(torch.bfloat16) + physical = torch.randint(0, blocks * block_size, (num_queries, topk), dtype=torch.int32) + physical[:, 3::7] = -1 + physical[:, 1] = physical[:, 0] + # A short row makes duplicate multiplicity and invalid-padding mistakes large. + physical[0] = -1 + physical[0, :3] = torch.tensor([0, 0, 1], dtype=torch.int32) + if num_queries > 1: + physical[-1] = -1 + if strided: + backing = torch.empty((blocks, 2, block_size, latent_dim + rope_dim), device=device, dtype=source.dtype) + cache = backing[:, 1] + cache.copy_(source) + else: + cache = source.to(device) + query_npu, rope_npu, physical_npu = query.to(device), query_rope.to(device), physical.to(device) + forward = partial(AscendHiSparseImpl._attention_physical, impl, query_npu, rope_npu, cache, physical_npu) + + graph = None + try: + if capture: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + for _ in range(2): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + captured = forward() + # Change both payload and selection length after graph capture. + query.neg_() + physical[0, :5] = torch.tensor([3, 3, 8, 12, 14], dtype=torch.int32) + query_npu.copy_(query) + physical_npu.copy_(physical) + for _ in range(3): + graph.replay() + actual = captured.float().cpu() + else: + actual = forward().float().cpu() + + selected = source.flatten(0, 1)[physical.clamp_min(0).long()].float() + logits = torch.einsum("thd,tkd->thk", query.float(), selected[..., :latent_dim]) + logits += torch.einsum("thd,tkd->thk", query_rope.float(), selected[..., latent_dim:]) + logits *= scale + logits.masked_fill_(~(physical >= 0)[:, None, :], -torch.inf) + probabilities = torch.softmax(logits, dim=-1).nan_to_num(0) + expected = torch.einsum("thk,tkd->thd", probabilities, selected[..., :latent_dim]) + torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.025) + if num_queries > 1: + assert torch.count_nonzero(actual[-1]) == 0 + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py new file mode 100644 index 000000000000..6433402b71b2 --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Keep HiSparse KV preprocessing numerically identical to native SFA.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch_npu + +# One fixed BF16 activation row with unit gamma and epsilon 1e-5. Its RMS +# normalization lands within 8e-8 of three BF16 rounding midpoints. Repeating +# this row reproduces the standalone-vs-fused difference; no model is loaded. +_ROUNDING_ROW_HEX = ( + "bc3df03d173e0d3c1bbe9cbcefbdb73822bde4ba14bd2ebda73dff3d003d1e3d1e3d053eec3da83d21be63bd8d3d40bd1cbda03bb1bd1abe04bb583c" + "563ec63bfdbdd1bd6b3d70bde73d103e343e0e3de13d1ebc323e03bd77bdec3c3abd2abca6bd043d4ebdf5bd1e3ce13da73dbebd083da73d063ec03d" + "1cbe71bced3c92bd3cbdd13d71bdc83dc13ca2bb85bcb1bd533d4a3dbabd57bdb43cc33c0bbb163e683d953dbeba65bdcb3ac73c56bd23bd86bd12bd" + "9ebd2abeb93def3ce33d7f3d783dca3cc3bc9b3d6b3d513d043d1d3dd33da03c9c3d023d59bd943dc9bd29bd2dbd8f3dcc3c0c3cba3d863d1f3ea23d" + "3b3db73d083df7bc9d3cc8bd43be19bef5bd91bd4ebe00bd65b9403e7f3dc5bcf6bb1b3d82bc883d25bd81bbe2bd553dc7bc673c0d3ce3bd74bc36bd" + "ba3d213c9d3b2dbbee3d6a3d823d8abb1dbd4abd883d383d85bdd73dacbdc5bae2bd57bdd3bd32bec9bd05be06bd6a3cd93d333dbf3c8ebcb73dc73a" + "11bd89bd213b0abc09bd663b5abccdbd0b3c623d753d603eb3bdc6bd813d9d3b2c3cb6bb9c3d3c3d853c96bb903d3fbdf33cd2ba0ebd25be05bd353d" + "edbdd3bb5b3e843d0e3dd53ca93cefbc9cbc633d463b57bd033cc8bda13dbe3dc4bd2f3e41bdeabd0f3dc9bd7a3cc13df8bd5b3c263d223ed83dcd3d" + "353d8ebdc53d673d0dbeccbda1bd1bbd4abdfbbca7bd3bbd00bd2bbc4d3c70bc74bdf13dd7bc2b3e31bd4abcd0bdb8bd9fbce4bba13b95bd2d3e51bd" + "1c3e9cbddfbae73d37bdddbce5bcaebae7bd39bda3bd0ebc0f3c83bd85bc34bbcf3cb53cc93c27bbabbdd6bd923c013ea43d3a3d2ebd643da13de33d" + "2e3c8dbd943c9ebdc63d92bb9ebd673d14bd283d8fbce43d323d883c493c5d3d66b85cbdafbdb1bd2d3c373d873d4abda5bdb9bcb8bd71bd76bd14be" + "443b083daf3cdf3d743d91bd913ddf3d3cbd2d3d53bd1cbe95bd15be88bdf53b8c3c763b9e3de33d2e3e82bc3e3d9ebd8dbd073ce3bc483d353d023e" + "4cbbb83d0b3e3c3e0abe74bd04be0fbdaabdd6bd3a3bf93a573c5f3e0c3d283dcc3d2b3d043eae3ca8bd90bd0abcb33dbebdc9bdee3cfbbdbd397dbd" + "243dbe3df13d203d84bd043dc0bdb0bca6bd06bd313dacbdf7bc05be263ec93dc03dd73c05bdad3d32bbcfbd31be5cbd84bda7bdd43d1d3d143e6a3d" + "053ead3b57bd173dd0bb83bdf73ce3bd773e323e99bbea3d793d5bbd90bce53d07bed2bc79bd98bd81bc713d41bd2dbc153d443c493e2b3e29be53bd" + "40be02bed83d13bdba3da2bc7f3d2bbded3dde3dba3cdcbd9dbc533e18beacbd54bd8fbb003db2bdc63df5bc033d763e1f3e13bcd6bc47bd33bd27be" + "ae3d323d183dc0bdbfbd50bd89bcc43cfebdacbdbabcf8bcea3cce3d22be963d093e153e2c3d2d3b2abcdc3d0cbc48bdb5bd1d3d40bdddbdb7bce23d" + "063e123b" +) + + +def _setup(): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse integration currently targets Ascend 950") + torch.npu.set_device(0) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_fused_preprocess_preserves_native_rounding_and_packed_writes(monkeypatch, capture): + _setup() + from vllm_ascend.attention import sfa_hisparse + + tokens, width, capacity = 17, 576, 128 + row = torch.frombuffer(bytearray.fromhex(_ROUNDING_ROW_HEX), dtype=torch.bfloat16).clone() + assert row.shape == (512,) + values = torch.empty(tokens, 2624, dtype=torch.bfloat16, device="npu")[:, 2048:] + source = torch.cat((row, torch.arange(64).to(torch.bfloat16) / 64)).expand(tokens, -1).clone() + values.copy_(source.to("npu")) + cos = torch.ones(tokens, 1, 1, 64, dtype=values.dtype, device="npu") + sin = torch.zeros_like(cos) + gamma = torch.ones(512, dtype=values.dtype, device="npu") + local_slots = torch.arange(tokens, dtype=torch.int64, device="npu") + reference_k = torch.empty(1, 128, 1, 512, dtype=values.dtype, device="npu") + reference_r = torch.empty(1, 128, 1, 64, dtype=values.dtype, device="npu") + + def native(): + torch_npu.npu_kv_rmsnorm_rope_cache( + values.view(tokens, 1, 1, width), + gamma, + cos, + sin, + local_slots, + reference_r, + reference_k, + epsilon=1e-5, + cache_mode="PA", + ) + return torch.cat((reference_k.view(-1, 512)[:tokens], reference_r.view(-1, 64)[:tokens]), -1).cpu() + + reference = native() + old_split, _ = torch_npu.npu_rms_norm(values[:, :512], gamma, epsilon=1e-5) + # A random oracle can miss this rare rounding boundary. This fixture must + # demonstrate the old path's disagreement before checking the replacement. + assert not torch.equal(old_split.cpu(), reference[:, :512]) + + resources = {} + + def persistent(key, allocate): + if key not in resources: + resources[key] = allocate() + return resources[key] + + monkeypatch.setattr( + sfa_hisparse, "current_workspace_manager", lambda: SimpleNamespace(get_persistent_resource=persistent) + ) + raw = torch.empty(4 * 6 * 128 * width * 2, dtype=torch.int8, device="npu") + backing = raw.view(values.dtype) + target = backing.as_strided((4, 128, width), (6 * 128 * width, width, 1), 2 * 128 * width) + slots_cpu = 256 + torch.arange(tokens, dtype=torch.int64) + slots_cpu[::4] = -1 + slots = slots_cpu.to("npu") + ends = torch.tensor([0, tokens], dtype=torch.int32, device="npu") + mirrored = [] + handle = SimpleNamespace( + write_target=lambda count, slot_count: (target, slots, count), + finish_kv_update=lambda: mirrored.append(True), + ) + impl = SimpleNamespace( + kv_lora_rank=512, + qk_rope_head_dim=64, + num_kv_heads=1, + _kv_staging_capacity=capacity, + index_group=SimpleNamespace(logical_topk_indices=slots), + kv_a_layernorm=SimpleNamespace(weight=gamma, variance_epsilon=1e-5), + hisparse_cache=handle, + ) + impl._kv_workspace = lambda dtype: sfa_hisparse.AscendHiSparseImpl._kv_workspace(impl, dtype) + first = impl._kv_workspace(values.dtype) + assert all(a is b for a, b in zip(first, impl._kv_workspace(values.dtype))) + assert first[2].dtype == torch.int64 and first[2].numel() == capacity + + def forward(): + return sfa_hisparse.AscendHiSparseImpl.exec_kv( + impl, values, cos, sin, None, slots, SimpleNamespace(query_start_loc=ends) + ) + + graph = None + try: + backing.fill_(-11) + if capture: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + rope, latent = forward() + else: + rope, latent = forward() + torch.testing.assert_close(latent.cpu(), reference[:, :512], rtol=0, atol=0) + torch.testing.assert_close(rope.cpu(), reference[:, 512:], rtol=0, atol=0) + assert mirrored + for iteration in range(3): + values.copy_((source * (iteration + 2)).to("npu")) + reference = native() + slots_cpu = 128 + torch.arange(tokens, dtype=torch.int64) + slots_cpu[1::4] = -1 + slots.copy_(slots_cpu) + valid = tokens - iteration - 1 + ends[1].fill_(valid) + backing.fill_(-11) + if capture: + graph.replay() + else: + rope, latent = forward() + expected_backing = torch.full_like(backing, -11).cpu() + expected = expected_backing.as_strided(target.shape, target.stride(), target.storage_offset()) + for row_id, slot in enumerate(slots_cpu.tolist()): + if row_id < valid and slot >= 0: + expected[slot // 128, slot % 128].copy_(reference[row_id]) + torch.testing.assert_close(backing.cpu(), expected_backing, rtol=0, atol=0) + torch.testing.assert_close(latent.cpu(), reference[:, :512], rtol=0, atol=0) + torch.testing.assert_close(rope.cpu(), reference[:, 512:], rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + + +def test_nope_preprocess_uses_native_layernorm(monkeypatch): + _setup() + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + + values = torch.arange(3 * 512, dtype=torch.float32, device="npu").view(3, 512).to(torch.bfloat16) + target = torch.full((1, 128, 512), -7, dtype=values.dtype, device="npu") + slots = torch.tensor([1, -1, 4], dtype=torch.int64, device="npu") + calls = [] + + def norm(x): + calls.append(True) + return x * 2 + + def unexpected(*args, **kwargs): + raise AssertionError("NoPE must not call the RoPE fused kernel") + + monkeypatch.setattr(torch_npu, "npu_kv_rmsnorm_rope_cache", unexpected) + impl = SimpleNamespace( + kv_lora_rank=512, + qk_rope_head_dim=0, + kv_a_layernorm=norm, + hisparse_cache=SimpleNamespace( + write_target=lambda count, slot_count: (target, slots, count), + finish_kv_update=lambda: None, + ), + ) + rope, latent = AscendHiSparseImpl.exec_kv( + impl, + values, + None, + None, + None, + slots, + SimpleNamespace(query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device="npu")), + ) + expected = torch.full_like(target, -7).cpu() + expected[0, 1].copy_((values[0] * 2).cpu()) + torch.testing.assert_close(target.cpu(), expected, rtol=0, atol=0) + assert calls and rope.shape == (3, 0) and latent.shape == values.shape diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py new file mode 100644 index 000000000000..396fafaff923 --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -0,0 +1,626 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real-NPU HiSparse state and mapped-payload tests against a CPU oracle.""" + +import json + +import pytest +import torch +import torch_npu # noqa: F401 + +from vllm_ascend.ops.hisparse import ( + HiSparseNPUOps, + MappedHostTensor, + ResolverWorkspace, + allocate_mapped_host, + pack_attention_rows, + scatter_packed_kv, +) + + +@pytest.fixture(autouse=True) +def _require_mapped_host_device(): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse mapped-host operators currently target Ascend 950") + torch.npu.set_device(0) + + +def reference(topk, requests, states, source, resident, hot_table, dgi, lru, block): + dgi = dgi.clone() + lru = lru.clone() + out = torch.full_like(topk, -1) + counts = torch.zeros(len(topk), dtype=torch.int32) + sh = torch.full_like(topk, -9) + sd = torch.full_like(topk, -9) + sc = torch.zeros(len(topk), dtype=torch.int32) + stats = torch.zeros(2, dtype=torch.int64) + for request in sorted(set(requests.tolist())): + if request < 0: + continue + rows = [i for i, r in enumerate(requests.tolist()) if r == request] + state = int(states[request]) + union = [] + references = [] + for row in rows: + for col, token in enumerate(topk[row].tolist()): + if token < 0: + continue + page, offset = divmod(token, block) + res = int(resident[state, page]) + if res: + out[row, col] = res * block + offset + else: + host = int(source[request, page]) * block + offset + references.append((row, col, host)) + if host not in union: + union.append(host) + counts[row] += 1 + order = lru[state].tolist() + hits = [slot for slot in order if int(dgi[state, slot]) in union] + evict = [slot for slot in order if slot not in hits] + mapping = {int(dgi[state, slot]): slot for slot in hits} + misses = [] + for host in union: + if host not in mapping: + slot = evict[len(misses)] + mapping[host] = slot + dgi[state, slot] = host + misses.append((host, slot)) + if union: + lru[state] = torch.tensor(evict[len(misses) :] + evict[: len(misses)] + hits, dtype=torch.int16) + for row, col, host in references: + slot = mapping[host] + out[row, col] = int(hot_table[request, slot // block]) * block + slot % block + for n, (host, slot) in enumerate(misses): + row = rows[0] + n // topk.size(1) + col = n % topk.size(1) + sh[row, col] = host + sd[row, col] = int(hot_table[request, slot // block]) * block + slot % block + sc[row] += 1 + stats += torch.tensor([len(references) - len(misses), len(misses)]) + return out, counts, sh, sd, sc, dgi, lru, stats + + +@pytest.mark.parametrize("device_row_count", [False, True]) +def test_resolver_strides_padding_graph(device_row_count): + torch.npu.set_device(0) + raw = torch.empty(8192, dtype=torch.uint8) + offset = (-raw.data_ptr()) % 4096 + registered = raw[offset : offset + 4096] + mapping = MappedHostTensor(registered) + host = registered.view(torch.float32).view(64, 16) + host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) + # Deliberate gaps between physical blocks exercise fused layer strides. + backing = torch.full((10, 2, 4, 16), -17.0, device="npu") + hot = backing[:, 1] + topk = torch.tensor( + [ + [0, 4, 8, 12, -1, -1, -1, -1], + [5, 9, 13, 12, -1, -1, -1, -1], + [0, 1, -1, 2, -1, -1, -1, -1], + [0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=torch.int32, + ) + requests = torch.tensor([0, 0, 1, -1], dtype=torch.int32) + states = torch.tensor([1, 0], dtype=torch.int32) + source = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.int32) + resident = torch.tensor([[0, 0, 0, 0], [5, 0, 0, 0]], dtype=torch.int32) + hot_table = torch.tensor([[1, 2, 6, 7], [3, 4, 8, 9]], dtype=torch.int32) + dgi = torch.tensor([[20, -1, -1, -1, -1, -1, -1, -1], [8, 99, -1, -1, -1, -1, -1, -1]], dtype=torch.int32) + lru = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7], [3, 0, 5, 1, 7, 4, 2, 6]], dtype=torch.int16) + dgi = torch.cat((dgi, torch.full((2, 8), -1, dtype=torch.int32)), dim=1) + lru = torch.cat((lru, torch.arange(8, 16, dtype=torch.int16)[None, :].expand(2, -1)), dim=1) + expected = reference(topk, requests, states, source, resident, hot_table, dgi, lru, 4) + workspace = ResolverWorkspace(4, 16, 16, torch.device("npu:0")) + ops = HiSparseNPUOps(mapping, workspace) + topk_n, requests_n, states_n, source_n, resident_n, hot_table_n, dgi_n, lru_n = [ + t.npu() for t in (topk, requests, states, source, resident, hot_table, dgi, lru) + ] + out = torch.empty_like(topk_n) + attn = torch.empty_like(topk_n) + counts = torch.empty(4, dtype=torch.int32, device="npu") + sh, sd = torch.full_like(out, -9), torch.full_like(out, -9) + sc = torch.empty_like(counts) + stats = torch.zeros(2, dtype=torch.int64, device="npu") + valid_rows = torch.tensor(3, dtype=torch.int32, device="npu") if device_row_count else None + if device_row_count: + # Stale graph padding looks like another valid request row. Only the + # device scalar separates it from current query rows. + requests_n[-1] = 1 + + def forward(): + ops.hisparse_resolve_residency( + host, + hot, + hot_table_n, + topk_n, + out, + dgi_n, + lru_n, + states_n, + 16, + 16, + None, + stats, + attn, + 8, + requests_n, + source_n, + 4, + None, + counts, + sh, + sd, + sc, + resident_n, + 4, + 0, + valid_rows, + ) + ops.hisparse_gather_compact(host, hot, sh, sd, sc) + + forward() + torch.npu.synchronize() + actual = (out, counts, sh, sd, sc, dgi_n, lru_n, stats) + for name, a, e in zip( + ("indices", "counts", "swap_host", "swap_device", "swap_counts", "dgi", "lru", "stats"), actual, expected + ): + torch.testing.assert_close(a.cpu(), e, msg=name) + expected_attn = torch.where(expected[0] >= 0, expected[0] // 4 * 8 + expected[0] % 4, -1) + torch.testing.assert_close(attn.cpu(), expected_attn) + for row in range(4): + for col in range(int(expected[4][row])): + src, dst = int(expected[2][row, col]), int(expected[3][row, col]) + torch.testing.assert_close(hot[dst // 4, dst % 4].cpu(), host[src]) + torch.testing.assert_close(backing[:, 0].cpu(), torch.full((10, 4, 16), -17.0)) + print(json.dumps({"resolver_mtp_union_lru_compact_gather": "PASS"}), flush=True) + latent = torch.empty((4, 8, 1, 12), device="npu") + rope = torch.empty((4, 8, 1, 4), device="npu") + pack_attention_rows(hot, out, latent, rope, 12) + torch.npu.synchronize() + hot_cpu = hot.cpu() + for row in range(4): + for col in range(8): + slot = int(expected[0][row, col]) + value = torch.zeros(16) if slot < 0 else hot_cpu[slot // 4, slot % 4] + torch.testing.assert_close(latent[row, col, 0].cpu(), value[:12]) + torch.testing.assert_close(rope[row, col, 0].cpu(), value[12:]) + new_latent = torch.arange(36, dtype=torch.float32).view(3, 1, 1, 12).npu() + new_rope = torch.arange(12, dtype=torch.float32).view(3, 1, 1, 4).npu() + slots = torch.tensor([0, -1, 2], dtype=torch.int64, device="npu") + valid = torch.tensor([2], dtype=torch.int32, device="npu") + scatter_packed_kv(new_latent, new_rope, hot, slots, valid) + torch.npu.synchronize() + torch.testing.assert_close(hot[0, 0].cpu(), torch.cat((new_latent[0].flatten().cpu(), new_rope[0].flatten().cpu()))) + torch.testing.assert_close(hot[0, 2].cpu(), hot_cpu[0, 2]) + print(json.dumps({"pack_scatter_strides_padding": "PASS"}), flush=True) + # Capture warmed kernels, reset state outside capture, then verify replay + # after changing the mapped CPU payload. No host selection work runs here. + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + for _ in range(2): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + forward() + dgi_n.copy_(dgi.npu()) + lru_n.copy_(lru.npu()) + stats.zero_() + sh.fill_(-9) + sd.fill_(-9) + host.add_(5000) + graph.replay() + torch.npu.synchronize() + for name, a, e in zip( + ("indices", "counts", "swap_host", "swap_device", "swap_counts", "dgi", "lru", "stats"), actual, expected + ): + torch.testing.assert_close(a.cpu(), e, msg="graph " + name) + for row in range(4): + for col in range(int(expected[4][row])): + src, dst = int(expected[2][row, col]), int(expected[3][row, col]) + torch.testing.assert_close(hot[dst // 4, dst % 4].cpu(), host[src]) + print(json.dumps({"acl_graph_replay_mapped_host_update": "PASS"}), flush=True) + for _ in range(100): + graph.replay() + torch.npu.synchronize() + torch.testing.assert_close(out.cpu(), expected[0]) + torch.testing.assert_close(dgi_n.cpu(), expected[5]) + torch.testing.assert_close(lru_n.cpu(), expected[6]) + print(json.dumps({"repeated_graph_replay": "PASS", "iterations": 100}), flush=True) + if device_row_count: + for active_rows in [0, 1, 2, 3]: + valid_rows.fill_(active_rows) + dgi_n.copy_(dgi.npu()) + lru_n.copy_(lru.npu()) + stats.zero_() + sh.fill_(-9) + sd.fill_(-9) + oracle_requests = requests.clone() + oracle_requests[active_rows:] = -1 + oracle = reference(topk, oracle_requests, states, source, resident, hot_table, dgi, lru, 4) + graph.replay() + torch.npu.synchronize() + for actual_value, expected_value in zip(actual, oracle): + torch.testing.assert_close(actual_value.cpu(), expected_value) + graph.reset() + mapping.close() + + +@pytest.mark.parametrize("query_rows,hot_size", [(2, 4096), (4, 10240)]) +def test_topk_2048_speculative_union_graph(query_rows, hot_size): + topk, width, block = 2048, 576, 128 + union = query_rows * topk + blocks = hot_size // block + 1 + host = allocate_mapped_host((hot_size + block, width), torch.bfloat16) + mapping = host._hisparse_registration + host.copy_((torch.arange(host.size(0)) % 97).to(torch.bfloat16)[:, None].expand_as(host)) + hot = torch.zeros((blocks, block, width), dtype=torch.bfloat16, device="npu") + indices = torch.arange(union, dtype=torch.int32, device="npu").view(query_rows, topk) + requests = torch.zeros(query_rows, dtype=torch.int32, device="npu") + states = torch.zeros(1, dtype=torch.int32, device="npu") + table = torch.arange(1, blocks, dtype=torch.int32, device="npu")[None, :] + resident = torch.zeros_like(table) + dgi = torch.full((1, hot_size), -1, dtype=torch.int32, device="npu") + lru = torch.arange(hot_size, dtype=torch.int32, device="npu").to(torch.int16)[None, :] + workspace = ResolverWorkspace(query_rows, hot_size, union, torch.device("npu:0")) + ops = HiSparseNPUOps(mapping, workspace) + out = torch.empty_like(indices) + counts = torch.empty(query_rows, dtype=torch.int32, device="npu") + sh, sd, sc = torch.empty_like(out), torch.empty_like(out), torch.empty_like(counts) + stats = torch.zeros(2, dtype=torch.int64, device="npu") + + def forward(): + ops.hisparse_resolve_residency( + host, + hot, + table, + indices, + out, + dgi, + lru, + states, + hot_size, + union, + None, + stats, + None, + block, + requests, + table, + block, + None, + counts, + sh, + sd, + sc, + resident, + block, + 0, + None, + ) + ops.hisparse_gather_compact(host, hot, sh, sd, sc) + + forward() + torch.npu.synchronize() + expected = torch.arange(block, block + union, dtype=torch.int32).view(query_rows, topk) + torch.testing.assert_close(out.cpu(), expected) + torch.testing.assert_close(sc.cpu(), torch.full((query_rows,), topk, dtype=torch.int32)) + torch.testing.assert_close(stats.cpu(), torch.tensor([0, union], dtype=torch.int64)) + torch.testing.assert_close(hot[1:].cpu().reshape(hot_size, width)[:union], host[block : block + union]) + expected_lru = torch.cat((torch.arange(union, hot_size), torch.arange(union))).to(torch.int16)[None, :] + torch.testing.assert_close(lru.cpu(), expected_lru) + forward() + torch.npu.synchronize() + torch.testing.assert_close(sc.cpu(), torch.zeros(query_rows, dtype=torch.int32)) + torch.testing.assert_close(stats.cpu(), torch.tensor([union, union], dtype=torch.int64)) + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + forward() + dgi.fill_(-1) + lru.copy_(torch.arange(hot_size, dtype=torch.int16)[None, :]) + stats.zero_() + graph.replay() + torch.npu.synchronize() + torch.testing.assert_close(out.cpu(), expected) + torch.testing.assert_close(sc.cpu(), torch.full((query_rows,), topk, dtype=torch.int32)) + torch.testing.assert_close(stats.cpu(), torch.tensor([0, union], dtype=torch.int64)) + for _ in range(10): + graph.replay() + torch.npu.synchronize() + torch.testing.assert_close(sc.cpu(), torch.zeros(query_rows, dtype=torch.int32)) + torch.testing.assert_close(stats.cpu(), torch.tensor([10 * union, union], dtype=torch.int64)) + torch.testing.assert_close(hot[1:].cpu().reshape(hot_size, width)[:union], host[block : block + union]) + torch.testing.assert_close(lru.cpu(), expected_lru) + graph.reset() + mapping.close() + + +@pytest.mark.parametrize("capture", [False, True]) +def test_random_request_unions(capture): + """Vary query counts, request order, stable state, duplicates and residency.""" + rows, topk, hot_size, block, width, requests_count = 40, 8, 16, 4, 16, 8 + host = allocate_mapped_host((132, width), torch.float32) + mapping = host._hisparse_registration + host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) + # 1..32 are hot pages and 33..64 are resident pages; layer 0 is a sentinel. + backing = torch.full((65, 2, block, width), -17.0, device="npu") + hot = backing[:, 1] + workspace = ResolverWorkspace(rows, hot_size, hot_size, torch.device("npu:0")) + ops = HiSparseNPUOps(mapping, workspace) + indices = torch.empty((rows, topk), dtype=torch.int32, device="npu") + requests = torch.empty(rows, dtype=torch.int32, device="npu") + states = torch.empty(requests_count, dtype=torch.int32, device="npu") + source = torch.arange(1, 33, dtype=torch.int32).view(requests_count, 4) + hot_table = source.npu() + source_n = source.npu() + resident = torch.empty_like(hot_table) + dgi = torch.empty((requests_count, hot_size), dtype=torch.int32, device="npu") + lru = torch.empty_like(dgi, dtype=torch.int16) + out, sh, sd = (torch.empty_like(indices) for _ in range(3)) + counts = torch.empty(rows, dtype=torch.int32, device="npu") + sc = torch.empty_like(counts) + stats = torch.empty(2, dtype=torch.int64, device="npu") + miss = torch.empty_like(indices, dtype=torch.uint8) + + def forward(): + ops.hisparse_resolve_residency( + host, + hot, + hot_table, + indices, + out, + dgi, + lru, + states, + hot_size, + hot_size, + miss, + stats, + None, + block, + requests, + source_n, + block, + None, + counts, + sh, + sd, + sc, + resident, + block, + 0, + None, + ) + ops.hisparse_gather_compact(host, hot, sh, sd, sc) + + graph = None + try: + for seed in range(16): + generator = torch.Generator().manual_seed(seed) + state_cpu = torch.randperm(requests_count, generator=generator).to(torch.int32) + requests_cpu = torch.full((rows,), -1, dtype=torch.int32) + indices_cpu = torch.randint(-1, 16, (rows, topk), generator=generator, dtype=torch.int32) + resident_cpu = torch.zeros((requests_count, 4), dtype=torch.int32) + dgi_cpu = torch.full((requests_count, hot_size), -1, dtype=torch.int32) + lru_cpu = torch.stack([torch.randperm(hot_size, generator=generator) for _ in range(requests_count)]).to( + torch.int16 + ) + cache_cpu = torch.full(hot.shape, -17.0) + cursor = 0 + for request in torch.randperm(requests_count, generator=generator).tolist(): + query_rows = 1 + (request + seed) % 4 + requests_cpu[cursor : cursor + query_rows] = request + cursor += query_rows + 1 # Adjacent padded rows share no state. + state = int(state_cpu[request]) + for page in range(4): + is_resident = request % 3 == 0 or (request % 3 == 2 and page < 2) + if is_resident: + physical = 33 + 4 * request + page + resident_cpu[state, page] = physical + cache_cpu[physical] = host[ + int(source[request, page]) * block : int(source[request, page]) * block + block + ] + selected = torch.randperm(hot_size, generator=generator)[:8] + 4 + request * hot_size + slots = torch.randperm(hot_size, generator=generator)[:8] + dgi_cpu[state, slots] = selected.to(torch.int32) + for slot, host_id in zip(slots.tolist(), selected.tolist()): + physical = int(source[request, slot // block]) + cache_cpu[physical, slot % block] = host[host_id] + expected = reference( + indices_cpu, requests_cpu, state_cpu, source, resident_cpu, source, dgi_cpu, lru_cpu, block + ) + for destination, value in [ + (indices, indices_cpu), + (requests, requests_cpu), + (states, state_cpu), + (resident, resident_cpu), + (dgi, dgi_cpu), + (lru, lru_cpu), + (hot, cache_cpu), + ]: + destination.copy_(value) + sh.fill_(-9) + sd.fill_(-9) + stats.zero_() + if capture and graph is None: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + forward() + dgi.copy_(dgi_cpu) + lru.copy_(lru_cpu) + hot.copy_(cache_cpu) + sh.fill_(-9) + sd.fill_(-9) + stats.zero_() + graph.replay() if capture else forward() + torch.npu.synchronize() + for label, actual, wanted in zip( + ("indices", "counts", "swap_host", "swap_device", "swap_counts", "dgi", "lru", "stats"), + (out, counts, sh, sd, sc, dgi, lru, stats), + expected, + ): + torch.testing.assert_close(actual.cpu(), wanted, msg=f"seed {seed}: {label}") + # Every selected token, whether resident, cached or copied, has + # exactly the canonical host payload. Sentinel layer stays intact. + cache = hot.cpu() + for row in range(rows): + request = int(requests_cpu[row]) + if request < 0: + continue + for col, token in enumerate(indices_cpu[row].tolist()): + if token < 0: + continue + physical = int(expected[0][row, col]) + host_id = int(source[request, token // block]) * block + token % block + torch.testing.assert_close(cache[physical // block, physical % block], host[host_id]) + torch.testing.assert_close(backing[:, 0].cpu(), torch.full((65, block, width), -17.0)) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + mapping.close() + + +def test_plan_gather_and_written_slot_invalidation(): + host = allocate_mapped_host((32, 16), torch.float32) + mapping = host._hisparse_registration + host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) + hot = torch.full((8, 4, 16), -17.0, device="npu") + ops = HiSparseNPUOps(mapping) + source = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]], device="npu", dtype=torch.int32) + destination = source + 4 + mask = torch.ones_like(source, dtype=torch.int32) + state = torch.tensor([0, -1], device="npu", dtype=torch.int32) + attention = torch.empty_like(source) + try: + ops.hisparse_gather_plan(host, hot, source, destination, mask, state, attention, 8) + torch.npu.synchronize() + cache = hot.cpu().view(32, 16) + torch.testing.assert_close(cache[5:13], host[1:9]) + torch.testing.assert_close(attention[1].cpu(), torch.full((8,), -1, dtype=torch.int32)) + dgi = torch.tensor([[4, 7, 9, -1], [11, 13, 15, -1]], dtype=torch.int32, device="npu") + states = torch.tensor([1, 0], dtype=torch.int32, device="npu") + reqs = torch.tensor([0, 1, 1, -1], dtype=torch.int32, device="npu") + written = torch.tensor([13, 4, 9, 11], dtype=torch.int64, device="npu") + ops.hisparse_invalidate_written_slots(dgi, states, reqs, written) + torch.npu.synchronize() + torch.testing.assert_close(dgi.cpu(), torch.tensor([[-1, 7, -1, -1], [11, -1, 15, -1]], dtype=torch.int32)) + finally: + torch.npu.synchronize() + mapping.close() + + +def test_upstream_prefill_staging_plan(): + from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan + + host = allocate_mapped_host((32, 16), torch.float32) + mapping = host._hisparse_registration + host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) + table = torch.tensor([[2, 1, 0], [3, 0, 0]], dtype=torch.int32, device="npu") + lengths = torch.tensor([5, 4], dtype=torch.int32, device="npu") + try: + plan = build_hisparse_prefill_staging_plan(table, lengths, 4, 3) + assert plan.miss_mask.dtype == torch.int32 + hot = torch.full((plan.row_ids.size(1) // 4, 4, 16), -17.0, device="npu") + ops = HiSparseNPUOps(mapping) + ops.hisparse_gather_plan(host, hot, plan.row_ids, plan.dst_rows, plan.miss_mask, None, None, 0) + torch.npu.synchronize() + source = plan.row_ids.cpu()[0] + expected = torch.full((source.numel(), 16), -17.0) + expected[source >= 0] = host[source[source >= 0].long()] + torch.testing.assert_close(hot.cpu().view_as(expected), expected) + finally: + torch.npu.synchronize() + mapping.close() + + +def test_invalidate_large_hot_capacity_graph(): + dgi = torch.arange(20480, dtype=torch.int32, device="npu").view(2, 10240) + initial = dgi.cpu() + states = torch.tensor([1, 0], dtype=torch.int32, device="npu") + requests = torch.tensor([0, 1, 1, -1], dtype=torch.int32, device="npu") + written = torch.tensor([20479, 0, 10239, 10240], dtype=torch.int64, device="npu") + expected = initial.clone() + expected[1, -1] = expected[0, 0] = expected[0, -1] = -1 + + def forward(): + HiSparseNPUOps.hisparse_invalidate_written_slots(dgi, states, requests, written) + + forward() + torch.npu.synchronize() + torch.testing.assert_close(dgi.cpu(), expected) + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + try: + with torch.npu.graph(graph, stream=stream): + forward() + for _ in range(10): + dgi.copy_(initial) + graph.replay() + torch.npu.synchronize() + torch.testing.assert_close(dgi.cpu(), expected) + finally: + torch.npu.synchronize() + graph.reset() + + +@pytest.mark.parametrize("groups", [1, 3, 21]) +def test_common_group_invalidation_changed_graph_replay(groups): + """Exercise the current common pointer-table caller with changing writes.""" + from types import SimpleNamespace + + from vllm.v1.hisparse.runtime import HiSparseSlotInvalidator + + initial = torch.arange(48, dtype=torch.int32).view(3, 16) + tables = [(initial + group * 7).npu() for group in range(groups)] + states = torch.tensor([2, 0, 1], dtype=torch.int32, device="npu") + requests = torch.tensor([0, 1, 2, -1], dtype=torch.int32, device="npu") + written = torch.tensor([35, 4, 19, 1], dtype=torch.int64, device="npu") + leaders = [ + SimpleNamespace(index_group=SimpleNamespace(device_global_indices=table), cache_ops=HiSparseNPUOps) + for table in tables + ] + invalidator = HiSparseSlotInvalidator(states) + + def forward(): + invalidator(leaders, written, requests) + + forward() + torch.npu.synchronize() + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + graph = torch.npu.NPUGraph() + try: + with torch.npu.graph(graph, stream=stream): + forward() + for values in [[35, 4, 19, 1], [47, 8, 23, 1], [-1, -1, -1, 1]]: + written.copy_(torch.tensor(values, dtype=torch.int64)) + for group, table in enumerate(tables): + table.copy_(initial + group * 7) + graph.replay() + torch.npu.synchronize() + for group, table in enumerate(tables): + expected = initial + group * 7 + for row, value in zip([2, 0, 1], values[:3]): + expected[row][expected[row] == value] = -1 + torch.testing.assert_close(table.cpu(), expected) + finally: + torch.npu.synchronize() + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py new file mode 100644 index 000000000000..8a69b1ad1daa --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Resident prefill plans must preserve each host page's physical row mapping.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch_npu # noqa: F401 +from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView, build_hisparse_prefill_staging_plan + +from vllm_ascend.distributed.kv_transfer.hisparse.backend import AscendHiSparseBackend +from vllm_ascend.ops.hisparse import HiSparseNPUOps, allocate_mapped_host, scatter_packed_kv + + +@pytest.mark.parametrize("resident_block_size", [128, 64, 32]) +@pytest.mark.parametrize("host_pages", [1, 2]) +def test_resident_prefill_plan_shared_layers(resident_block_size, host_pages): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse mapped-host staging currently targets Ascend 950") + torch.npu.set_device(0) + device = "npu:0" + block_size, width, layers, blocks = 128, 576, 4, 16 + page_elements = block_size * width + block_stride = 6 * page_elements + raw = torch.zeros(blocks * block_stride * 2, dtype=torch.int8, device=device) + # Match the shared HMA slab: create and retain every layer view before use. + views = [ + PagedCacheView.bind( + raw, + dtype=torch.bfloat16, + row_width=width, + byte_offset=layer * resident_block_size * width * 2, + block_stride=block_stride * 2, + num_blocks=blocks, + block_size=resident_block_size, + ) + for layer in range(layers) + ] + host = allocate_mapped_host(((host_pages + 1) * block_size, width), torch.bfloat16) + host.zero_() + host[block_size:].fill_(-23) + ops = HiSparseNPUOps(host._hisparse_registration) + try: + subpages = block_size // resident_block_size + resident_ids = torch.arange(2, 2 + host_pages * subpages, dtype=torch.int32) + if host_pages == 2: + # The last subpage is a real host miss; other rows must stay on NPU. + resident_ids[-1] = 0 + resident_table = resident_ids[None, :].to(device) + host_table = torch.arange(1, host_pages + 1, dtype=torch.int32, device=device)[None, :] + plan = build_hisparse_prefill_staging_plan( + host_table, torch.tensor([host_pages * block_size - 7], device=device), block_size, host_pages + ) + state = torch.tensor([0], dtype=torch.int32, device=device) + logical_rows = torch.arange(host_pages * block_size) + physical = resident_ids.long()[logical_rows // resident_block_size] * resident_block_size + physical += logical_rows % resident_block_size + hits = resident_ids[logical_rows // resident_block_size] > 0 + physical[~hits] = -1 + expected_gpu = torch.cat((torch.full((block_size,), -1), physical)).int() + expected_mask = (expected_gpu < 0).int() + original_rows = None + for layer, view in enumerate(views): + columns = torch.arange(width)[None, :] % 16 + payload = (layer * 16 + logical_rows[:, None].float() / 8 + columns / 16 + 1).bfloat16() + on_device = payload.to(device) + scatter_packed_kv(on_device[:, :512], on_device[:, 512:], view.cache, physical.int().to(device)) + plan.ensure_gpu_sources(resident_table, state, resident_block_size) + torch.testing.assert_close(plan.gpu_row_ids.cpu().flatten(), expected_gpu, rtol=0, atol=0) + torch.testing.assert_close(plan.miss_mask.cpu().flatten(), expected_mask, rtol=0, atol=0) + if original_rows is None: + original_rows = plan.gpu_row_ids + else: + # Layers in an index group reuse the same resolved page plan. + assert plan.gpu_row_ids is original_rows + actual = HiSparseRuntime.gather_prefill_cache( + SimpleNamespace(cache_ops=ops, backend=AscendHiSparseBackend()), + host.view(-1, block_size, width), + plan, + view.cache, + ) + expected = host.clone() + expected[block_size:][hits] = payload[hits] + torch.testing.assert_close(actual.cpu().flatten(0, 1), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + host._hisparse_registration.close() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_backend_resident_payload_graph_replay(dtype): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse resident gather currently targets Ascend 950") + torch.npu.set_device(0) + width = 656 if dtype != torch.bfloat16 else 576 + stride = 3 * width + raw = torch.arange(7 * width * dtype.itemsize, dtype=torch.int64).to(torch.uint8).to("npu:0") + source = raw.view(dtype).as_strided((2, 2, width), (stride, width, 1), storage_offset=width) + rows = torch.tensor([3, 0, 2, -1], dtype=torch.int32, device="npu:0") + destination = torch.empty((4, width), dtype=dtype, device="npu:0") + backend = AscendHiSparseBackend() + backend.gather_resident_rows(source, destination, rows) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + backend.gather_resident_rows(source, destination, rows) + for values in ([3, 0, 2, -1], [0, 3, 1, 2]): + rows.copy_(torch.tensor(values, dtype=torch.int32, device="npu:0")) + raw.bitwise_xor_(1) + graph.replay() + torch.npu.synchronize() + safe = torch.tensor(values).clamp_min(0) + expected = source.view(torch.uint8).cpu()[safe // 2, safe % 2] + assert torch.equal(destination.view(torch.uint8).cpu(), expected) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py b/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py new file mode 100644 index 000000000000..ea224efab1ba --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The SFA bridge must preserve valid selection order across graph replay.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch_npu # noqa: F401 + + +@pytest.mark.parametrize("capture", [False, True]) +def test_valid_order_duplicates_and_interleaved_padding(monkeypatch, capture): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse SFA integration currently targets Ascend 950") + torch.npu.set_device(0) + from vllm_ascend.utils import bootstrap_custom_op_env + + bootstrap_custom_op_env() + import vllm_ascend.vllm_ascend_C # noqa: F401 + + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.ops import hisparse + + torch.manual_seed(61) + device, topk, latent, rope, capacity = "npu:0", 2048, 512, 64, 32 + workspace = ( + torch.empty((capacity, topk, 1, latent), dtype=torch.bfloat16, device=device), + torch.empty((capacity, topk, 1, rope), dtype=torch.bfloat16, device=device), + torch.arange(capacity * topk // 128, dtype=torch.int32, device=device).view(capacity, -1), + torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), + torch.arange(1, capacity + 1, dtype=torch.int32, device=device), + torch.full((capacity,), topk, dtype=torch.int32, device=device), + ) + impl = SimpleNamespace( + topk=topk, kv_lora_rank=latent, qk_rope_head_dim=rope, scale=0.0625, _workspace=lambda _: workspace + ) + query = torch.randn(3, 8, latent, dtype=torch.bfloat16, device=device) + query_rope = torch.randn(3, 8, rope, dtype=torch.bfloat16, device=device) + cache = torch.randn(1, 128, latent + rope, dtype=torch.bfloat16, device=device) + indices = torch.full((3, topk), -1, dtype=torch.int32) + indices[0, :6] = torch.tensor([7, -1, 2, 7, -1, 0]) + indices[2, :6] = torch.tensor([0, 5, -1, 3, 5, 1]) + device_indices = indices.to(device) + observed = torch.empty_like(device_indices) + original_pack = hisparse.pack_attention_rows + + def observe(cache, physical, selected_latent, selected_rope, latent_dim): + observed.copy_(physical) + return original_pack(cache, physical, selected_latent, selected_rope, latent_dim) + + monkeypatch.setattr(hisparse, "pack_attention_rows", observe) + + def forward(): + return AscendHiSparseImpl._attention_physical(impl, query, query_rope, cache, device_indices) + + graph = None + try: + if capture: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + result = forward() + indices[0, :6] = torch.tensor([-1, 9, 1, 9, 4, -1]) + device_indices.copy_(indices) + for _ in range(3): + graph.replay() + else: + result = forward() + expected = torch.full_like(indices, -1) + for row in range(3): + valid = indices[row][indices[row] >= 0] + expected[row, : valid.numel()] = valid + torch.testing.assert_close(observed.cpu(), expected, rtol=0, atol=0) + assert torch.isfinite(result.cpu()).all() + assert torch.count_nonzero(result[1].cpu()) == 0 + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py new file mode 100644 index 000000000000..ad6949406559 --- /dev/null +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Source-isolated CPU boundary checks; these do not qualify NPU execution.""" + +import ast +import ctypes +import inspect +import math +import sys +from dataclasses import dataclass, replace +from pathlib import Path +from types import SimpleNamespace as NS + +import pytest +import torch + +ROOT = Path(__file__).resolve().parents[4] + + +def definitions(relative, names, **bindings): + path = ROOT / "vllm_ascend" / relative + tree = ast.parse(path.read_bytes()) + nodes = [n for n in tree.body if isinstance(n, (ast.ClassDef, ast.FunctionDef)) and n.name in names] + assert len(nodes) == len(names) + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + namespace = {"torch": torch, **bindings} + module = ast.fix_missing_locations(ast.Module(body=[future, *nodes], type_ignores=[])) + exec(compile(module, str(path), "exec"), namespace) + return namespace + + +def backend_type(): + return definitions( + "distributed/kv_transfer/hisparse/backend.py", {"AscendHiSparseBackend"}, HiSparseBackend=object + )["AscendHiSparseBackend"] + + +def test_merged_core_resolver_call_keeps_optional_outputs_in_their_current_positions(): + operation = definitions("ops/hisparse.py", {"HiSparseNPUOps"})["HiSparseNPUOps"].hisparse_resolve_residency + # Argument positions from the reviewed merged Core Runtime.swap_in call. + # The former pre-merge method rejected this call or shifted output buffers. + arguments = [object() for _ in range(26)] + bound = inspect.signature(operation).bind(object(), *arguments).arguments + for name, position in { + "miss_mask": 10, + "stats": 11, + "source_block_size": 16, + "resolved_global_indices": 17, + "valid_counts": 18, + "resident_null_block": 24, + "num_valid_rows": 25, + }.items(): + assert bound[name] is arguments[position] + + +def test_failed_unregister_retains_host_owner_for_safe_retry(monkeypatch): + tensor = torch.empty(64, dtype=torch.uint8) + events = [] + + class Registration: + def __init__(self, allocation): + self.tensor = allocation + + def close(self): + events.append("close") + if len(events) == 1: + raise RuntimeError("driver did not retire registration") + + monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(MappedHostTensor=Registration)) + backend = backend_type()() + backend.pin_tensor(tensor) + registration = tensor._hisparse_registration + assert registration.tensor is tensor + with pytest.raises(RuntimeError, match="driver did not retire"): + backend.unpin_tensor(tensor) + assert tensor._hisparse_registration is registration + backend.unpin_tensor(tensor) + assert not hasattr(tensor, "_hisparse_registration") + assert events == ["close", "close"] + + +def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): + cls = definitions("ops/hisparse.py", {"MappedHostTensor"}, ctypes=ctypes)["MappedHostTensor"] + backing = torch.arange(128, dtype=torch.uint8) + mapping = cls.__new__(cls) + mapping.tensor = backing + mapping.host_address = backing.data_ptr() + mapping.nbytes = backing.numel() + mapping.device_ptr = 4096 + mapping.closed = False + assert mapping.address_of(backing[17:43]) == 4113 + with pytest.raises(ValueError, match="outside"): + mapping.address_of(torch.empty(256, dtype=torch.uint8)) + with pytest.raises(ValueError, match="contiguous"): + mapping.address_of(backing[::2]) + mapping.closed = True + with pytest.raises(ValueError, match="live mapped"): + mapping.address_of(backing) + + +@pytest.mark.parametrize("max_rows,hot_rows,union_rows", [(1, 16, 8), (9, 48, 24), (37, 128, 72)]) +def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( + monkeypatch, max_rows, hot_rows, union_rows +): + allocated = [] + + class Workspace: + def __init__(self, **geometry): + self.geometry = geometry + allocated.append(self) + + class Ops: + def __init__(self, mapping, workspace): + self.mapping, self.workspace = mapping, workspace + + monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(ResolverWorkspace=Workspace, HiSparseNPUOps=Ops)) + backend = backend_type()() + + def group(): + owner = NS() + leader = NS( + index_group=owner, + max_swap_rows=max_rows, + region_stride=hot_rows, + max_union_rows=union_rows, + device=torch.device("cpu"), + registered_host_pool=NS(_hisparse_registration=object()), + ) + owner.leader = leader + leader.cache_ops = backend.create_cache_ops(leader) + follower = NS(**vars(leader)) + return leader, backend.create_cache_ops(follower) + + first, follower = group() + second, _ = group() + assert follower.workspace is first.cache_ops.workspace + assert follower.mapping is first.registered_host_pool._hisparse_registration + assert second.cache_ops.workspace is not first.cache_ops.workspace + assert len(allocated) == 2 + assert allocated[0].geometry == dict( + max_rows=max_rows, hot_size=hot_rows, max_union_rows=union_rows, device=torch.device("cpu") + ) + + +@dataclass(frozen=True) +class IndexerSpec: + block_size: int = 4 + num_kv_heads: int = 1 + head_size: int = 4 + block_stride_alignment: int = 3 + scale_dim: int = 0 + cache_sparse_li_c8: bool = False + cache_sparse_li_c4: bool = False + + @property + def page_size_bytes(self): + return self.block_size * self.num_kv_heads * self.head_size * torch.float32.itemsize + + def get_block_stride_alignment(self): + return self.block_stride_alignment + + +def binding_functions(): + return definitions( + "distributed/kv_transfer/hisparse/binding.py", + {"align_hisparse_indexer_specs", "get_hisparse_indexer_specs", "_bind_indexer_views"}, + AscendSFAIndexerCacheSpec=IndexerSpec, + UniformTypeKVCacheSpecs=type("Uniform", (), {}), + math=math, + replace=replace, + ) + + +@pytest.mark.parametrize("blocks,stride", [(1, 1), (3, 2), (5, 3)]) +def test_expanded_indexer_alias_preserves_shared_slab_offsets(blocks, stride): + functions = binding_functions() + spec = IndexerSpec() + page = 16 + backing = torch.arange(blocks * stride * page + page, dtype=torch.float32) + original = backing.as_strided((blocks, 1, 4, 4), (stride * page, page, 4, 1), storage_offset=page) + layer = NS() + caches = {"indexer": original} + functions["_bind_indexer_views"](caches, {"indexer": spec}, {"indexer": layer}) + view = caches["indexer"][0] + assert view.data_ptr() == original.data_ptr() + assert view.storage_offset() == original.storage_offset() + assert view.untyped_storage().data_ptr() == backing.data_ptr() + assert layer.hisparse_block_stride == stride + for block in range(blocks): + torch.testing.assert_close(view[block * stride, :, 0], original[block, 0]) + view[(blocks - 1) * stride, 1, 0, 2] = -73 + assert original[-1, 0, 1, 2] == -73 + + +def test_alignment_updates_only_indexer_specs_without_changing_inputs(): + functions = binding_functions() + source, indexer = object(), IndexerSpec() + values = {"source": source, "indexer": indexer} + result = functions["align_hisparse_indexer_specs"](values) + assert result["source"] is source + assert result["indexer"].get_block_stride_alignment() == math.lcm(3, indexer.page_size_bytes) + assert values["indexer"] is indexer and indexer.block_stride_alignment == 3 + + +@pytest.mark.parametrize("field", ["scale_dim", "cache_sparse_li_c8", "cache_sparse_li_c4"]) +def test_unsupported_indexer_formats_are_rejected_before_pool_binding(field): + functions = binding_functions() + spec = replace(IndexerSpec(), **{field: 1}) + config = NS(kv_cache_groups=[NS(layer_names=["indexer"], kv_cache_spec=spec)]) + with pytest.raises(ValueError, match="unquantized"): + functions["get_hisparse_indexer_specs"](config) + + +@pytest.mark.parametrize("hash_lookahead", [0, 7]) +def test_hisparse_roles_preserve_core_admission_and_explicit_hash_budget(hash_lookahead): + calls = [] + marker = object() + functions = definitions( + "patch/platform/patch_kv_cache_coordinator.py", + {"get_kv_cache_coordinator"}, + KVCacheGroupRole=NS(HISPARSE_SOURCE="source"), + _select_kv_token_budget=lambda maximum, inflight: maximum if inflight is None else inflight, + _orig_get_kv_cache_coordinator=lambda **kwargs: (calls.append(kwargs), marker)[1], + ) + config = NS(kv_cache_groups=[NS(role="source"), NS(role="resident"), NS(role="hot")]) + result = functions["get_kv_cache_coordinator"]( + config, + max_model_len=65536, + max_in_flight_tokens=129, + num_hash_lookahead_tokens=hash_lookahead, + enable_caching=False, + ) + assert result is marker and len(calls) == 1 + assert calls[0]["kv_cache_config"] is config + assert calls[0]["max_in_flight_tokens"] == 129 + assert calls[0]["num_hash_lookahead_tokens"] == hash_lookahead + assert calls[0]["enable_caching"] is False + assert "eagle_attn_layer_names" not in calls[0] diff --git a/tests/ut/ops/test_mla_weight_hook_cpu.py b/tests/ut/ops/test_mla_weight_hook_cpu.py new file mode 100644 index 000000000000..5c1f9a450d53 --- /dev/null +++ b/tests/ut/ops/test_mla_weight_hook_cpu.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Exercise the real weight-hook wrapper with both loader call signatures.""" + +import ast +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.parametrize("sparse", [False, True]) +@pytest.mark.parametrize("explicit_dtype", [False, True]) +@pytest.mark.parametrize("replicated", [False, True]) +def test_post_load_hook_preserves_activation_dtype(monkeypatch, sparse, explicit_dtype, replicated): + source = Path(__file__).resolve().parents[3] / "vllm_ascend/ops/mla.py" + tree = ast.parse(source.read_bytes()) + wrapper = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "wrapped_process_weights" + ) + calls = [] + + class SFA: + def process_weights_after_loading(self, dtype): + calls.append(("SFA", dtype)) + + class Dense: + def process_weights_after_loading(self, dtype): + calls.append(("replicated", dtype)) + + def old_hook(dtype): + calls.append(("original", dtype)) + + def current_hook(): + calls.append(("original", torch.bfloat16)) + + monkeypatch.setitem(sys.modules, "vllm_ascend.attention.sfa_v1", SimpleNamespace(AscendSFAImpl=SFA)) + namespace = { + "torch": torch, + "self": SimpleNamespace( + mla_attn=SimpleNamespace(impl=SFA() if sparse else Dense(), act_dtype=torch.bfloat16), + dcp_q_replicate_enabled=replicated, + ), + "original_process_weights": old_hook if explicit_dtype else current_hook, + } + exec(compile(ast.Module(body=[wrapper], type_ignores=[]), str(source), "exec"), namespace) + if explicit_dtype: + namespace["wrapped_process_weights"](torch.float16) + else: + namespace["wrapped_process_weights"]() + expected = torch.float16 if explicit_dtype else torch.bfloat16 + assert calls == [("SFA" if sparse else "replicated" if replicated else "original", expected)] diff --git a/tests/ut/worker/test_hisparse_capture_mirror_cpu.py b/tests/ut/worker/test_hisparse_capture_mirror_cpu.py new file mode 100644 index 000000000000..3ab15b32f19a --- /dev/null +++ b/tests/ut/worker/test_hisparse_capture_mirror_cpu.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Capture warmups must not publish a real connector step's Host mirrors.""" + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest + +ROOT = Path(__file__).resolve().parents[3] +BASE = ROOT / "vllm_ascend/attention/sfa_hisparse.py" + + +def load_helper(): + tree = ast.parse(BASE.read_text()) + node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "finish_hisparse_kv_update") + namespace = {} + exec(compile(ast.Module(body=[node], type_ignores=[]), str(BASE), "exec"), namespace) + return namespace[node.name] + + +@pytest.mark.parametrize("filename", ["sfa_hisparse.py"]) +def test_capture_builder_marks_only_its_fresh_metadata_and_forwards_arguments(filename): + path = BASE.with_name(filename) + tree = ast.parse(path.read_text()) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "AscendHiSparseMetadataBuilder") + method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "build_for_cudagraph_capture") + calls = [] + + class Parent: + def build_for_cudagraph_capture(self, common, **kwargs): + calls.append((common, kwargs)) + return self.build() + + def build(self): + return SimpleNamespace(hisparse_for_capture=False) + + cls = ast.ClassDef( + name=cls.name, bases=[ast.Name(id="Parent", ctx=ast.Load())], keywords=[], body=[method], decorator_list=[] + ) + namespace = {"Parent": Parent} + exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), str(path), "exec"), namespace) + builder = namespace[cls.name]() + common, draft = object(), object() + capture = builder.build_for_cudagraph_capture(common, draft_index=draft) + assert calls == [(common, {"draft_index": draft})] + assert capture.hisparse_for_capture is True + assert builder.build().hisparse_for_capture is False + + +def test_repeated_capture_forwards_do_not_submit_host_mirrors(): + finish = load_helper() + calls = [] + cache = SimpleNamespace(finish_kv_update=lambda: calls.append(True)) + for _ in range(5): + finish(cache, SimpleNamespace(hisparse_for_capture=True)) + assert calls == [] + finish(cache, SimpleNamespace(hisparse_for_capture=False)) + finish(cache, SimpleNamespace()) + assert calls == [True, True] + + +def test_real_forward_duplicate_errors_remain_visible(): + finish = load_helper() + + def duplicate(): + raise RuntimeError("HiSparse layer 0 mirrored twice.") + + with pytest.raises(RuntimeError, match="mirrored twice"): + finish(SimpleNamespace(finish_kv_update=duplicate), SimpleNamespace(hisparse_for_capture=False)) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 33e183fc2da4..d4eafe8b589a 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -154,6 +154,8 @@ def get_name() -> str: @staticmethod def get_builder_cls(): + if get_current_vllm_config().attention_config.hisparse_config is not None: + return AscendHiSparseIndexerMetadataBuilder return AscendSFAIndexerMetadataBuilder @staticmethod @@ -965,3 +967,46 @@ def _build( num_decode_tokens=num_decode_tokens, uses_pcp_sharded_layout=is_pcp_sharded, ) + + +class AscendHiSparseIndexerMetadataBuilder(AscendSFAIndexerMetadataBuilder): + """Map logical indexer pages into HiSparse's shared device allocation.""" + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + context = vllm_config.compilation_config.static_forward_context + self._cache_layers = [context[name] for name in layer_names] + self._physical_metadata_buffers: dict[object, tuple[torch.Tensor, torch.Tensor]] = {} + scheduler = vllm_config.scheduler_config + max_capture = vllm_config.compilation_config.max_cudagraph_capture_size or 0 + self._table_capacity = ( + max(scheduler.max_num_seqs, max_capture) + 1, + _round_up(vllm_config.model_config.max_model_len, self.kernel_block_size) // self.kernel_block_size, + ) + + def _build(self, common_attn_metadata, buffer_key, use_cached_rope, copy_rope, **kwargs): + metadata = super()._build(common_attn_metadata, buffer_key, use_cached_rope, copy_rope, **kwargs) + strides = {getattr(layer, "hisparse_block_stride", 1) for layer in self._cache_layers} + if len(strides) != 1: + raise RuntimeError("HiSparse indexers sharing metadata must have the same physical block stride") + stride = strides.pop() + if stride == 1: + return metadata + if buffer_key not in self._physical_metadata_buffers: + self._physical_metadata_buffers[buffer_key] = ( + torch.empty(self._table_capacity, dtype=metadata.block_table.dtype, device=self.device), + torch.empty(self._slot_capacity, dtype=metadata.slot_mapping.dtype, device=self.device), + ) + table_buffer, slot_buffer = self._physical_metadata_buffers[buffer_key] + rows, columns = metadata.block_table.shape + table = table_buffer[:rows, :columns] + slots = slot_buffer[: metadata.slot_mapping.numel()] + torch.mul(metadata.block_table, stride, out=table) + logical = metadata.slot_mapping + nonnegative = logical.clamp_min(0) + physical = (nonnegative // self.kernel_block_size) * (self.kernel_block_size * stride) + physical += nonnegative % self.kernel_block_size + slots.copy_(torch.where(logical >= 0, physical, -1)) + metadata.block_table = table + metadata.slot_mapping = slots + return metadata diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py new file mode 100644 index 000000000000..56eec31e7555 --- /dev/null +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ascend SFA execution over upstream HiSparse's resident and hot cache pool.""" + +from dataclasses import dataclass + +import torch +import torch_npu +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder +from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan +from vllm.v1.worker.workspace import current_workspace_manager + +from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder + +# Bound the temporary D2D packing workspace independently of prefill length. +_ATTENTION_CHUNK_ROWS = 32 + + +def finish_hisparse_kv_update(cache, metadata) -> None: + """Publish real writes; capture forwards have no connector step to mirror. + + V2 runs both FULL-graph warmup and capture with runtime mode NONE. + Preserve the device KV writes in those forwards, but leave Host mirror + submission to the real connector step after graph replay. + """ + if not getattr(metadata, "hisparse_for_capture", False): + cache.finish_kv_update() + + +@dataclass +class AscendHiSparsePrefillMetadata: + host_staging_plan: HiSparsePrefillStagingPlan + + +@dataclass +class AscendHiSparseMetadata(AscendSFAMetadata): + req_id_per_token: torch.Tensor | None = None + num_reqs: int = 0 + prefill: AscendHiSparsePrefillMetadata | None = None + hisparse_for_capture: bool = False + + +class AscendHiSparseMetadataBuilder(AscendSFAMetadataBuilder): + def build_for_cudagraph_capture(self, common_attn_metadata, **kwargs): + metadata = super().build_for_cudagraph_capture(common_attn_metadata, **kwargs) + metadata.hisparse_for_capture = True + return metadata + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): + super().__init__(kv_cache_spec, layer_names, vllm_config, device, metadata_cls=AscendHiSparseMetadata, **kwargs) + self._request_row_buffers: dict[int | None, torch.Tensor] = {} + self._row_capacity = max( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) + + def _build(self, common_attn_metadata, draft_index=None, pcp_context=None): + metadata = super()._build(common_attn_metadata, draft_index, pcp_context) + metadata.num_reqs = common_attn_metadata.num_reqs + rows = torch.arange(metadata.num_input_tokens, device=self.device, dtype=torch.int32) + if draft_index not in self._request_row_buffers: + self._request_row_buffers[draft_index] = torch.empty( + self._row_capacity, device=self.device, dtype=torch.int32 + ) + metadata.req_id_per_token = self._request_row_buffers[draft_index][: metadata.num_input_tokens] + metadata.req_id_per_token.copy_(torch.searchsorted(metadata.query_start_loc[1:].contiguous(), rows, right=True)) + metadata.req_id_per_token.clamp_(max=max(metadata.num_reqs - 1, 0)) + if metadata.num_prefills: + seq_lens_cpu = metadata.seq_lens_cpu + if seq_lens_cpu is None: + # Capacity may use an upper bound; the staging plan receives + # current device lengths to decide the actual copied rows. + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + if seq_lens_cpu is None: + raise RuntimeError("HiSparse prefill requires scheduler sequence-length bounds on CPU") + prefill_lens = seq_lens_cpu[metadata.num_decodes :] + capacity = int(((prefill_lens + metadata.block_size - 1) // metadata.block_size).sum()) + metadata.prefill = AscendHiSparsePrefillMetadata( + build_hisparse_prefill_staging_plan( + metadata.block_table[metadata.num_decodes :], + metadata.seq_lens[metadata.num_decodes :], + metadata.block_size, + capacity, + ) + ) + return metadata + + +class AscendHiSparseImpl(AscendSFAImpl): + """Keep HiSparse residency decisions on device, with an SFA layout bridge.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.enable_sparse_sfa_c8 or self.enable_sparse_sfa_turboquant: + raise ValueError("Ascend HiSparse currently requires an unquantized main MLA KV cache") + parallel = self.vllm_config.parallel_config + if parallel.prefill_context_parallel_size > 1 or parallel.decode_context_parallel_size > 1: + raise ValueError("Ascend HiSparse does not support context parallelism") + builder = kwargs.get("index_group_builder") + if builder is None: + if self.topk_indices_buffer is None: + raise ValueError("Ascend HiSparse requires a model sparse-index group") + builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) + from vllm_ascend.patch.worker.patch_hisparse_streams import install + + install() + self.index_group, self.index_group_index = builder.register_layer( + self.has_indexer and not self.skip_topk, + self.vllm_config, + row_width=self.kv_lora_rank + self.qk_rope_head_dim, + kv_cache_dtype=self.vllm_config.cache_config.cache_dtype, + ) + if not isinstance(self.index_group, HiSparseMLAIndexGroup): + raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") + self.hisparse_cache = self.index_group.cache(self.index_group_index) + self.topk = self.index_group.logical_topk_indices.shape[-1] + if self.topk % 128: + raise ValueError("Ascend HiSparse SFA requires top-k to be divisible by 128") + max_rows = max( + self.vllm_config.scheduler_config.max_num_batched_tokens, + self.vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) + self._kv_staging_capacity = (max_rows + 127) // 128 * 128 + + def _fused_preprocess_type(self): + # Fused prologs write the ordinary split cache directly. HiSparse owns + # both the resident destination and optional host-mirror staging. This + # also keeps quantized linear weight layout under its normal owner. + return None + + def _compose_sfa_kv_cache(self, kv_cache): + return (kv_cache,) + + def _workspace(self, dtype): + key = ("ascend_hisparse_attention", self.topk, self.kv_lora_rank, self.qk_rope_head_dim, dtype) + + def allocate(): + device = self.index_group.logical_topk_indices.device + shape = (_ATTENTION_CHUNK_ROWS, self.topk, 1) + return ( + torch.empty((*shape, self.kv_lora_rank), dtype=dtype, device=device), + torch.empty((*shape, self.qk_rope_head_dim), dtype=dtype, device=device), + torch.arange(_ATTENTION_CHUNK_ROWS * self.topk // 128, dtype=torch.int32, device=device).view( + _ATTENTION_CHUNK_ROWS, self.topk // 128 + ), + torch.arange(self.topk, dtype=torch.int32, device=device).view(1, 1, self.topk), + torch.arange(1, _ATTENTION_CHUNK_ROWS + 1, dtype=torch.int32, device=device), + torch.full((_ATTENTION_CHUNK_ROWS,), self.topk, dtype=torch.int32, device=device), + ) + + return current_workspace_manager().get_persistent_resource(key, allocate) + + def _kv_workspace(self, dtype): + key = ( + "ascend_hisparse_kv_preprocess", + self._kv_staging_capacity, + self.kv_lora_rank, + self.qk_rope_head_dim, + dtype, + ) + + def allocate(): + device = self.index_group.logical_topk_indices.device + shape = (self._kv_staging_capacity // 128, 128, 1) + return ( + torch.empty((*shape, self.kv_lora_rank), dtype=dtype, device=device), + torch.empty((*shape, self.qk_rope_head_dim), dtype=dtype, device=device), + torch.arange(self._kv_staging_capacity, dtype=torch.int64, device=device), + ) + + return current_workspace_manager().get_persistent_resource(key, allocate) + + def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=None): + self._workspace(hidden_states.dtype) + if self.qk_rope_head_dim: + # Include preprocessing scratch in profiling and allocate it before + # graph capture, including dummy forwards without attention metadata. + self._kv_workspace(hidden_states.dtype) + self.index_group.prepare_for_batch(self.index_group_index, attn_metadata) + return super().forward(layer_name, hidden_states, kv_cache, attn_metadata, output) + + def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): + from vllm_ascend.ops.hisparse import scatter_packed_kv + + assert self.kv_a_layernorm is not None + values = kv_no_split.reshape(-1, self.kv_lora_rank + self.qk_rope_head_dim) + num_tokens = values.shape[0] + if self.qk_rope_head_dim: + if num_tokens > self._kv_staging_capacity: + raise ValueError("HiSparse KV preprocessing exceeds its profiled token capacity") + latent_cache, rope_cache, local_slots = self._kv_workspace(values.dtype) + # Match ordinary SFA's fused RMSNorm/RoPE rounding. Separate RMSNorm + # can differ by a BF16 ULP near a rounding midpoint. Local slots + # isolate padding from the real resident/hot destination below. + torch_npu.npu_kv_rmsnorm_rope_cache( + values.view(num_tokens, self.num_kv_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim), + self.kv_a_layernorm.weight, + cos, + sin, + local_slots[:num_tokens], + rope_cache, + latent_cache, + epsilon=self.kv_a_layernorm.variance_epsilon, + cache_mode="PA", + ) + latent = latent_cache.view(-1, self.kv_lora_rank)[:num_tokens] + rope = rope_cache.view(-1, self.qk_rope_head_dim)[:num_tokens] + else: + # The native NoPE path uses the layer's normalization implementation. + latent = self.kv_a_layernorm(values) + rope = values[:, self.kv_lora_rank :] + cache = self.hisparse_cache + destination, mapping, num_rows = cache.write_target(values.shape[0], slots.numel()) + num_valid_rows = attn_metadata.query_start_loc[-1:] + scatter_packed_kv(latent[:num_rows], rope[:num_rows], destination, mapping, num_valid_rows) + # The common connector mirrors resident rows after this write. Draft + # rows are deferred until the next step, and FULL graphs are mirrored + # by finish_forward after replay. + finish_hisparse_kv_update(cache, attn_metadata) + return rope, latent + + def _attention_physical(self, query, query_rope, packed_cache, physical_topk): + from vllm_ascend.ops.hisparse import pack_attention_rows + + latent, rope, block_table, indices, query_ends, seq_lens = self._workspace(query.dtype) + positions = indices.view(1, -1).to(torch.float32) + outputs = [] + for start in range(0, query.shape[0], _ATTENTION_CHUNK_ROWS): + end = min(start + _ATTENTION_CHUNK_ROWS, query.shape[0]) + count = end - start + # SFA's causal mode derives the valid selection length from KV + # lengths. Move invalid entries to the tail while preserving the + # selected order and duplicate multiplicity. Reordering valid KV + # changes BF16 reductions. Unique FP32 position keys are exact for + # supported top-k sizes and avoid integer ArgSort on AiCPU. + physical = physical_topk[start:end] + keys = positions + (physical < 0).to(torch.float32) * physical.shape[-1] + physical = physical.gather(1, keys.sort(dim=-1).indices) + valid_lengths = (physical >= 0).sum(dim=-1, dtype=torch.int32) + selected_latent, selected_rope = latent[:count], rope[:count] + pack_attention_rows(packed_cache, physical, selected_latent, selected_rope, self.kv_lora_rank) + selected_indices = torch.where(physical[:, None, :] >= 0, indices, -1) + result = torch.ops._C_ascend.npu_sparse_flash_attention( + query=query[start:end], + key=selected_latent.view(-1, 128, 1, self.kv_lora_rank), + value=selected_latent.view(-1, 128, 1, self.kv_lora_rank), + sparse_indices=selected_indices, + scale_value=self.scale, + sparse_block_size=1, + block_table=block_table[:count], + actual_seq_lengths_query=query_ends[:count], + actual_seq_lengths_kv=valid_lengths.clamp_min(1), + query_rope=query_rope[start:end] if self.qk_rope_head_dim else None, + key_rope=selected_rope.view(-1, 128, 1, self.qk_rope_head_dim) if self.qk_rope_head_dim else None, + layout_query="TND", + layout_kv="PA_BSND", + sparse_mode=3, + attention_mode=2, + ) + if isinstance(result, tuple): + result = result[0] + result = result.masked_fill(~(physical >= 0).any(dim=-1)[:, None, None], 0) + outputs.append(result) + return torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] + + def _execute_sparse_flash_attention_process( + self, + ql_nope, + q_pe, + kv_cache, + topk_indices, + attn_metadata, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table=None, + ): + metadata = attn_metadata + logical = topk_indices.squeeze(1) if topk_indices.ndim == 3 else topk_indices + assert metadata.req_id_per_token is not None + assert metadata.query_start_loc is not None + rows = torch.arange(logical.shape[0], dtype=torch.int32, device=logical.device) + logical = torch.where(rows[:, None] < metadata.query_start_loc[-1], logical, -1) + self.index_group.set_logical_topk_ready(self.index_group_index) + outputs = [] + num_decode = metadata.num_decode_tokens + if num_decode: + cache = self.hisparse_cache + assert cache.source_block_table is not None + # Keep mixed prefill/decode batches on the NPU residency resolver. + # The generic index-group resident shortcut uses the CUDA sparse + # index converter when a batch also contains prefill tokens. + physical = cache.swap_in( + metadata.req_id_per_token[:num_decode], + block_table=cache.source_block_table, + logical_topk_indices=logical[:num_decode], + block_size=metadata.block_size, + num_valid_rows=metadata.query_start_loc[-1:], + ) + outputs.append( + self._attention_physical( + ql_nope[:num_decode], + q_pe[:num_decode] if q_pe is not None else None, + self.index_group.physical_kv_cache(self.index_group_index), + physical, + ) + ) + if metadata.num_prefills: + host_cache = self.hisparse_cache.runtime.host_cache + staged, staged_table, req_ids = self.index_group.stage_prefill_rows( + self.index_group_index, + host_cache.view(-1, metadata.block_size, host_cache.shape[-1]), + metadata, + ) + selected = logical[num_decode:] + block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size + request_tables = staged_table[req_ids.to(torch.int64)] + physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size + physical = physical + selected.clamp_min(0) % metadata.block_size + physical = physical.masked_fill(selected < 0, -1).to(torch.int32) + outputs.append( + self._attention_physical( + ql_nope[num_decode:], + q_pe[num_decode:] if q_pe is not None else None, + staged, + physical, + ) + ) + if not outputs: + return torch.zeros_like(ql_nope) + result = torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] + if result.shape[0] < ql_nope.shape[0]: + result = torch.cat((result, torch.zeros_like(ql_nope[result.shape[0] :])), dim=0) + return result diff --git a/vllm_ascend/attention/sfa_v1.py b/vllm_ascend/attention/sfa_v1.py index a21b8c50794b..c8b4fc01121c 100644 --- a/vllm_ascend/attention/sfa_v1.py +++ b/vllm_ascend/attention/sfa_v1.py @@ -370,6 +370,10 @@ def supports_batch_invariance(cls) -> bool: @staticmethod def get_builder_cls(): + if get_current_vllm_config().attention_config.hisparse_config is not None: + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder + + return AscendHiSparseMetadataBuilder if get_ascend_config().sparse_kv_offload_config.enabled: from vllm_ascend.attention.sfa_kv_offload import AscendSFAKVOffloadMetadataBuilder @@ -390,6 +394,10 @@ def get_kv_cache_shape( @staticmethod def get_impl_cls() -> type["AscendSFAImpl"]: + if get_current_vllm_config().attention_config.hisparse_config is not None: + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + + return AscendHiSparseImpl if get_ascend_config().sparse_kv_offload_config.enabled: from vllm_ascend.attention.sfa_kv_offload import AscendSFAKVOffloadImpl diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/__init__.py b/vllm_ascend/distributed/kv_transfer/hisparse/__init__.py new file mode 100644 index 000000000000..4b8ccbf15b4c --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/hisparse/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Ascend device operations for the upstream HiSparse connector.""" diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py new file mode 100644 index 000000000000..5a4950b2ee06 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""NPU streams, mapped host memory and DMA for the common HiSparse runtime.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING, Any + +import torch +from vllm.v1.hisparse.backend import HiSparseBackend + +if TYPE_CHECKING: + from vllm.v1.hisparse.runtime import HiSparseRuntime + from vllm.v1.kv_cache_interface import KVCacheConfig + + +class AscendHiSparseBackend(HiSparseBackend): + # NPU event IPC is not implemented here. Keep private pools per TP rank; + # the common connector still owns transfer completion and TP coordination. + supports_shared_host_pool = False + stats_dtype = torch.int64 + + def check_available(self) -> None: + from vllm_ascend.ops.hisparse import HiSparseNPUOps + from vllm_ascend.utils import bootstrap_custom_op_env + + if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): + # A5 disables the global custom-op selection policy. HiSparse still + # needs its explicit DMA dependency; load the extension lazily here + # without changing that policy for unrelated model operators. + bootstrap_custom_op_env() + try: + importlib.import_module("vllm_ascend.vllm_ascend_C") + except ImportError as error: + raise RuntimeError("Ascend HiSparse requires the compiled Ascend DMA extension.") from error + if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): + raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") + if not callable(HiSparseNPUOps): + raise RuntimeError("Ascend HiSparse device kernels are unavailable.") + + def pin_tensor(self, tensor: torch.Tensor) -> None: + from vllm_ascend.ops.hisparse import MappedHostTensor + + tensor._hisparse_registration = MappedHostTensor(tensor) + + def unpin_tensor(self, tensor: torch.Tensor) -> None: + registration = tensor._hisparse_registration + registration.close() + del tensor._hisparse_registration + + def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: + from vllm_ascend.ops.hisparse import HiSparseNPUOps, ResolverWorkspace + + group = runtime.index_group + if runtime is group.leader: + workspace = ResolverWorkspace( + max_rows=runtime.max_swap_rows, + hot_size=runtime.region_stride, + max_union_rows=runtime.max_union_rows, + device=runtime.device, + ) + else: + leader_ops = group.leader.cache_ops + if not isinstance(leader_ops, HiSparseNPUOps): + raise RuntimeError("Bind the HiSparse leader before follower caches.") + workspace = leader_ops.workspace + mapping = runtime.registered_host_pool._hisparse_registration + return HiSparseNPUOps(mapping, workspace) + + def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: + from .binding import get_hisparse_indexer_specs + + get_hisparse_indexer_specs(kv_cache_config) + + def prepare_kv_cache_views( + self, + kv_cache_config: KVCacheConfig, + kv_caches: dict[str, Any], + forward_context: dict[str, Any], + ) -> None: + from .binding import _bind_indexer_views, get_hisparse_indexer_specs + + _bind_indexer_views(kv_caches, get_hisparse_indexer_specs(kv_cache_config), forward_context) + + def is_current_stream_capturing(self) -> bool: + return torch.npu.is_current_stream_capturing() + + def swap_blocks_batch(self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor) -> None: + torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1) diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py new file mode 100644 index 000000000000..8dc0b2937560 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Bind Ascend indexer views while retaining upstream HiSparse allocation.""" + +import math +from dataclasses import replace +from typing import Any + +import torch +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheSpec, + UniformTypeKVCacheSpecs, +) + +from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec + + +def align_hisparse_indexer_specs(specs: dict[str, KVCacheSpec]) -> dict[str, KVCacheSpec]: + """Declare dense indexer-page addressing to Core's shared-slab planner. + + Core includes this alignment when sizing each physical device block, before + choosing the block count. Source MLA specs and their host layout stay intact. + """ + return { + name: replace(spec, block_stride_alignment=math.lcm(spec.get_block_stride_alignment(), spec.page_size_bytes)) + if isinstance(spec, AscendSFAIndexerCacheSpec) + else spec + for name, spec in specs.items() + } + + +def get_hisparse_indexer_specs( + kv_cache_config: KVCacheConfig, +) -> dict[str, AscendSFAIndexerCacheSpec]: + """Validate the static indexer format before common storage allocation.""" + indexer_specs = {} + for group in kv_cache_config.kv_cache_groups: + for name in group.layer_names: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + spec = spec.kv_cache_specs[name] + if isinstance(spec, AscendSFAIndexerCacheSpec): + if spec.scale_dim or spec.cache_sparse_li_c8 or spec.cache_sparse_li_c4: + raise ValueError("Ascend HiSparse currently requires an unquantized indexer cache.") + indexer_specs[name] = spec + + return indexer_specs + + +def _bind_indexer_views( + kv_caches: dict[str, Any], indexer_specs: dict[str, AscendSFAIndexerCacheSpec], forward_context: dict[str, Any] +) -> None: + for name, spec in indexer_specs.items(): + cache = kv_caches[name] + # The indexer kernels address dense pages. Expand the view over the + # shared slab and let the layer scale its block IDs and write slots. + if cache.ndim != 4 or cache.shape[1] != spec.num_kv_heads: + raise ValueError(f"Unexpected HiSparse indexer cache shape for {name}: {tuple(cache.shape)}") + block_size = cache.shape[2] + page_elements = block_size * spec.num_kv_heads * spec.head_size + if cache.stride(0) % page_elements: + raise ValueError("HiSparse indexer slab stride must be a whole number of indexer pages.") + block_stride = cache.stride(0) // page_elements + physical_blocks = (cache.shape[0] - 1) * block_stride + 1 + indexer_cache = torch.as_strided( + cache, + size=(physical_blocks, block_size, spec.num_kv_heads, spec.head_size), + stride=(page_elements, spec.num_kv_heads * spec.head_size, spec.head_size, 1), + ) + forward_context[name].hisparse_block_stride = block_stride + kv_caches[name] = (indexer_cache,) diff --git a/vllm_ascend/ops/hisparse.py b/vllm_ascend/ops/hisparse.py new file mode 100644 index 000000000000..b2bbd7d387ea --- /dev/null +++ b/vllm_ascend/ops/hisparse.py @@ -0,0 +1,660 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ascend HiSparse operators and explicit mapped-host memory ownership. + +Mapping and scratch allocation happen at cache initialization. Decode planning +and payload copies stay on the NPU, including during ACL graph replay. +""" + +from __future__ import annotations + +import ctypes +import math +import mmap + +import torch +from vllm.triton_utils import triton + +from vllm_ascend.ops.triton.hisparse import ( + assign_residency_kernel, + classify_residency_kernel, + finish_residency_kernel, + gather_rows_kernel, + invalidate_kernel, + pack_attention_rows_kernel, + partition_residency_kernel, + prepare_residency_kernel, + protect_residency_kernel, + residency_stats_kernel, + scatter_packed_kv_kernel, + update_lru_kernel, +) + +_ACL_HOST_REGISTER_MAPPED = 0 +_ACL_HAC_AIV = 3 + + +class MappedHostTensor: + """Own an ACL mapping and retain the CPU allocation until explicit close. + + Close only after all streams and graphs using this mapping have completed + and been destroyed. No destructor unmaps memory behind asynchronous work. + """ + + def __init__(self, tensor: torch.Tensor) -> None: + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError("HiSparse host mapping requires a contiguous CPU tensor") + self.tensor = tensor + self._lib = ctypes.CDLL("libascendcl.so") + self._lib.aclrtHostRegister.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_int, + ctypes.POINTER(ctypes.c_void_p), + ] + self._lib.aclrtHostRegister.restype = ctypes.c_int + self._lib.aclrtHostUnregister.argtypes = [ctypes.c_void_p] + self._lib.aclrtHostUnregister.restype = ctypes.c_int + self._lib.aclrtHostMemMapCapabilities.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.POINTER(ctypes.c_int)] + self._lib.aclrtHostMemMapCapabilities.restype = ctypes.c_int + device_index = torch.npu.current_device() + self.device = torch.device("npu", device_index) + capability = ctypes.c_int() + status = self._lib.aclrtHostMemMapCapabilities(device_index, _ACL_HAC_AIV, ctypes.byref(capability)) + if status or not capability.value: + raise RuntimeError( + f"HiSparse requires AIV host-memory mapping: status={status}, capability={capability.value}" + ) + # The upstream host-pool allocator supplies an aligned subview of a + # larger backing allocation. Register that range, not its unaligned + # storage base; retained tensor ownership still retains the backing. + self.host_address = tensor.data_ptr() + self.nbytes = tensor.numel() * tensor.element_size() + if not self.nbytes: + raise ValueError("Cannot register an empty HiSparse host allocation") + mapped = ctypes.c_void_p() + status = self._lib.aclrtHostRegister( + self.host_address, self.nbytes, _ACL_HOST_REGISTER_MAPPED, ctypes.byref(mapped) + ) + if status: + raise RuntimeError(f"aclrtHostRegister for HiSparse failed: {status}") + self.device_ptr = int(mapped.value) + self.closed = False + + def address_of(self, tensor: torch.Tensor) -> int: + offset = tensor.data_ptr() - self.host_address + size = tensor.numel() * tensor.element_size() + if self.closed or tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError("HiSparse source must be a contiguous view of a live mapped allocation") + if offset < 0 or offset + size > self.nbytes: + raise ValueError("HiSparse source is outside its registered host allocation") + return self.device_ptr + offset + + def close(self) -> None: + if not self.closed: + status = self._lib.aclrtHostUnregister(self.host_address) + if status: + raise RuntimeError(f"aclrtHostUnregister for HiSparse failed: {status}") + self.closed = True + + +def allocate_mapped_host(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + """Allocate page-aligned CPU storage and attach its explicit mapping owner.""" + count = math.prod(shape) + nbytes = count * torch.empty((), dtype=dtype).element_size() + if nbytes <= 0: + raise ValueError("HiSparse host allocation must be nonempty") + allocated_bytes = triton.cdiv(nbytes, mmap.PAGESIZE) * mmap.PAGESIZE + owner = mmap.mmap(-1, allocated_bytes) + raw = torch.frombuffer(owner, dtype=torch.uint8) + mapping = MappedHostTensor(raw) + tensor = raw[:nbytes].view(dtype).view(shape) + # The mapping owns raw; frombuffer keeps the mmap alive. Runtime owns the + # registration through this base tensor until its explicit shutdown. + tensor._hisparse_registration = mapping + return tensor + + +class ResolverWorkspace: + """Fixed scratch shared by one index group on its ordered copy stream.""" + + def __init__(self, max_rows: int, hot_size: int, max_union_rows: int, device: torch.device) -> None: + if not 0 < max_union_rows <= hot_size <= 32768 or max_rows <= 0: + raise ValueError("HiSparse requires 0 < union <= hot_size <= 32768 and positive max_rows") + self.max_rows = max_rows + self.hot_size = hot_size + self.max_union_rows = max_union_rows + # Packed int64 keys retain exact int32 host IDs and deterministic first + # reference positions. torch.sort's out= ABI reuses these NPU buffers. + self.row_keys = torch.empty(max_rows * hot_size, dtype=torch.int64, device=device) + self.sorted_rows = torch.empty_like(self.row_keys) + self.row_sort_indices = torch.empty_like(self.row_keys) + self.cache_keys = torch.empty((max_rows, hot_size), dtype=torch.int64, device=device) + self.sorted_cache = torch.empty_like(self.cache_keys) + self.cache_sort_indices = torch.empty_like(self.cache_keys) + self.first_refs = torch.empty(max_rows * hot_size, dtype=torch.int32, device=device) + self.first_misses = torch.empty_like(self.first_refs) + self.metadata = torch.empty((max_rows, 8), dtype=torch.int32, device=device) + self.row_stats = torch.empty((max_rows, 4), dtype=torch.int64, device=device) + self.lru_temp = torch.empty((max_rows, hot_size), dtype=torch.int32, device=device) + self.protected = torch.empty_like(self.lru_temp) + self.hit_prefix = torch.empty_like(self.lru_temp) + self.errors = torch.zeros((max_rows, 8), dtype=torch.int32, device=device) + + +def _cache_layout(host_cache: torch.Tensor, hot_cache: torch.Tensor) -> tuple[int, int, int, int]: + if host_cache.ndim != 2 or not host_cache.is_contiguous() or host_cache.device.type != "cpu": + raise ValueError("HiSparse host cache must be a contiguous CPU [rows, width] view") + _check_packed_cache(hot_cache) + row_bytes = hot_cache.size(2) * hot_cache.element_size() + if hot_cache.stride(1) != hot_cache.size(2) or host_cache.size(1) * host_cache.element_size() != row_bytes: + raise ValueError("HiSparse cache rows must be dense and have matching byte widths") + return ( + row_bytes, + hot_cache.size(1), + hot_cache.size(0) * hot_cache.size(1), + hot_cache.stride(0) * hot_cache.element_size(), + ) + + +def _check_tensor(tensor, name, device, dtypes, *, ndim=None, shape=None, contiguous=True) -> None: + """Validate metadata without copying values or synchronizing the device.""" + if tensor.device != device or tensor.dtype not in dtypes: + raise ValueError(f"{name} must have dtype {dtypes} on {device}") + if ndim is not None and tensor.ndim != ndim: + raise ValueError(f"{name} must have {ndim} dimensions") + if shape is not None and tuple(tensor.shape) != tuple(shape): + raise ValueError(f"{name} must have shape {tuple(shape)}") + if contiguous and not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _check_packed_cache(cache) -> None: + if cache.device.type != "npu" or cache.ndim != 3 or min(cache.shape) <= 0: + raise ValueError("HiSparse device cache must be nonempty NPU [blocks, block_size, width]") + if cache.stride(2) != 1 or cache.stride(1) != cache.size(2): + raise ValueError("HiSparse device cache must have dense token rows") + if cache.stride(0) < cache.size(1) * cache.size(2): + raise ValueError("HiSparse cache blocks must not overlap") + if cache.size(0) * cache.size(1) >= 2**31: + raise ValueError("HiSparse physical row IDs must fit int32") + + +def _check_payload_rows(tensor, name, device, dtype, rows) -> None: + _check_tensor(tensor, name, device, (dtype,), contiguous=False) + if tensor.ndim < 2 or tensor.size(0) < rows: + raise ValueError(f"{name} must contain at least {rows} token rows") + stride = 1 + for dim in range(tensor.ndim - 1, 0, -1): + if tensor.size(dim) > 1 and tensor.stride(dim) != stride: + raise ValueError(f"{name} must have dense token rows") + stride *= tensor.size(dim) + if tensor.stride(0) < stride: + raise ValueError(f"{name} token rows must not overlap") + + +class HiSparseNPUOps: + """Runtime-owned implementations of the upstream HiSparse operator ABI.""" + + def __init__(self, host_mapping: MappedHostTensor, workspace: ResolverWorkspace | None = None) -> None: + self.host_mapping = host_mapping + self.workspace = workspace + + def hisparse_resolve_residency( + self, + host_cache, + hot_cache, + hot_block_table, + global_indices, + hot_indices, + device_global_indices, + lru_slots, + request_state_indices, + region_stride, + max_union_rows, + miss_mask, + stats, + attention_indices, + attention_block_stride, + request_ids, + source_block_table, + source_block_size, + resolved_global_indices, + valid_counts, + swap_host_physical_rows, + swap_device_physical_rows, + swap_counts, + resident_block_table, + resident_block_size, + resident_null_block=0, + num_valid_rows=None, + ) -> None: + """Resolve logical selections and prepare resident/hot-cache refill plans. + + Protect every selected token in a request's query-row union before + evicting LRU slots. Update residency and output plans in place using the + persistent workspace; host KV payloads are copied by the gather stage. + """ + workspace = self.workspace + if workspace is None: + raise RuntimeError("HiSparse resolver workspace must be prepared before forward") + _cache_layout(host_cache, hot_cache) + device = hot_cache.device + _check_tensor(global_indices, "global_indices", device, (torch.int32,), ndim=2, contiguous=False) + if global_indices.stride(1) != 1: + raise ValueError("HiSparse top-k columns must be contiguous") + _check_tensor(device_global_indices, "device_global_indices", device, (torch.int32,), ndim=2) + _check_tensor(lru_slots, "lru_slots", device, (torch.int16,), shape=device_global_indices.shape) + rows, topk = global_indices.shape + if num_valid_rows is not None: + _check_tensor(num_valid_rows, "num_valid_rows", device, (torch.int32,)) + if num_valid_rows.numel() != 1: + raise ValueError("num_valid_rows must contain one device scalar") + hot_size = lru_slots.size(1) + if rows == 0: + return + if rows > workspace.max_rows or hot_size != workspace.hot_size or max_union_rows != workspace.max_union_rows: + raise ValueError("HiSparse resolver input exceeds its prepared workspace") + if region_stride != device_global_indices.size(1) or region_stride != hot_size or hot_size < topk: + raise ValueError("HiSparse region and LRU widths must match and cover top-k") + if topk <= 0 or topk % 8 or hot_size % 16: + raise ValueError("Ascend HiSparse requires top-k aligned to 8 and hot capacity aligned to 16") + if workspace.row_keys.device != device: + raise ValueError("HiSparse resolver workspace must be on the cache device") + if host_cache.size(0) >= 2**31: + raise ValueError("HiSparse host row IDs must fit int32") + for name, tensor in ( + ("hot_indices", hot_indices), + ("resolved_global_indices", resolved_global_indices), + ("swap_host_physical_rows", swap_host_physical_rows), + ("swap_device_physical_rows", swap_device_physical_rows), + ): + if tensor is not None: + _check_tensor(tensor, name, device, (torch.int32,), shape=(rows, topk)) + if request_ids is not None: + _check_tensor(request_ids, "request_ids", device, (torch.int32,), shape=(rows,)) + if request_state_indices is not None: + _check_tensor(request_state_indices, "request_state_indices", device, (torch.int32,), ndim=1) + _check_tensor(hot_block_table, "hot_block_table", device, (torch.int32,), ndim=2, contiguous=False) + request_capacity = ( + device_global_indices.size(0) if request_state_indices is None else request_state_indices.numel() + ) + if ( + hot_block_table.stride(1) != 1 + or hot_block_table.size(0) < request_capacity + or hot_block_table.size(1) < triton.cdiv(hot_size, hot_cache.size(1)) + ): + raise ValueError("HiSparse hot block table does not cover all state slots") + for name, table, block_size in ( + ("source_block_table", source_block_table, source_block_size), + ("resident_block_table", resident_block_table, resident_block_size), + ): + if table is not None: + _check_tensor(table, name, device, (torch.int32,), ndim=2, contiguous=False) + if table.stride(1) != 1 or block_size <= 0: + raise ValueError(f"{name} requires contiguous columns and a positive block size") + if attention_indices is not None: + _check_tensor( + attention_indices, "attention_indices", device, (torch.int32,), shape=(rows, topk), contiguous=False + ) + if attention_indices.stride(1) != 1 or attention_block_stride < hot_cache.size(1): + raise ValueError("HiSparse attention indices require dense columns and a valid block stride") + if miss_mask is not None: + _check_tensor(miss_mask, "miss_mask", device, (torch.uint8, torch.bool, torch.int32), shape=(rows, topk)) + if valid_counts is not None: + _check_tensor(valid_counts, "valid_counts", device, (torch.int32,), contiguous=False) + if tuple(valid_counts.shape) not in ((rows,), (rows, 1)): + raise ValueError("valid_counts must have one entry per query row") + if swap_counts is not None: + _check_tensor(swap_counts, "swap_counts", device, (torch.int32,), shape=(rows,)) + if stats is not None: + _check_tensor(stats, "stats", device, (torch.int64,), shape=(2,)) + dummy = hot_indices + source = source_block_table + resident = resident_block_table + has_swaps = swap_counts is not None + if ( + len({tensor is not None for tensor in (swap_counts, swap_host_physical_rows, swap_device_physical_rows)}) + != 1 + ): + raise ValueError("All three compact HiSparse swap outputs are required together") + row_keys = workspace.row_keys[: rows * topk].view(rows, topk) + sorted_rows = workspace.sorted_rows[: rows * topk].view(rows, topk) + row_sort_indices = workspace.row_sort_indices[: rows * topk].view(rows, topk) + cache_keys = workspace.cache_keys[:rows] + sorted_cache = workspace.sorted_cache[:rows] + prepare_residency_kernel[(rows,)]( + global_indices, + request_ids if request_ids is not None else dummy, + request_state_indices if request_state_indices is not None else dummy, + source if source is not None else dummy, + resident if resident is not None else dummy, + device_global_indices, + hot_indices, + resolved_global_indices if resolved_global_indices is not None else dummy, + valid_counts if valid_counts is not None else dummy, + row_keys, + cache_keys, + workspace.metadata, + num_valid_rows if num_valid_rows is not None else dummy, + ROWS=rows, + TOP_K=topk, + HOT_SIZE=hot_size, + INPUT_STRIDE=global_indices.stride(0), + COUNT_STRIDE=valid_counts.stride(0) if valid_counts is not None else 1, + STATE_ROWS=device_global_indices.size(0), + STATE_COUNT=request_state_indices.numel() if request_state_indices is not None else 0, + HOST_ROWS=host_cache.size(0), + SOURCE_STRIDE=source.stride(0) if source is not None else 0, + SOURCE_ROWS=source.size(0) if source is not None else 0, + SOURCE_COLS=source.size(1) if source is not None else 0, + SOURCE_BLOCK=source_block_size if source is not None else 1, + RESIDENT_STRIDE=resident.stride(0) if resident is not None else 0, + RESIDENT_ROWS=resident.size(0) if resident is not None else 0, + RESIDENT_COLS=resident.size(1) if resident is not None else 0, + RESIDENT_BLOCK=resident_block_size if resident is not None else 1, + RESIDENT_NULL=resident_null_block, + HAS_REQUESTS=request_ids is not None, + HAS_STATE=request_state_indices is not None, + HAS_SOURCE=source is not None, + HAS_RESIDENT=resident is not None, + HAS_RESOLVED=resolved_global_indices is not None, + HAS_COUNTS=valid_counts is not None, + HAS_VALID_ROWS=num_valid_rows is not None, + BK=triton.next_power_of_2(topk), + multibuffer=False, + ) + torch.sort(row_keys, dim=-1, out=(sorted_rows, row_sort_indices)) + torch.sort(cache_keys, dim=-1, out=(sorted_cache, workspace.cache_sort_indices[:rows])) + classify_residency_kernel[(rows,)]( + row_keys, + sorted_rows, + sorted_cache, + workspace.metadata, + hot_block_table, + hot_indices, + workspace.first_refs, + workspace.first_misses, + workspace.row_stats, + TOP_K=topk, + HOT_SIZE=hot_size, + HOT_BLOCK=hot_cache.size(1), + TABLE_STRIDE=hot_block_table.stride(0), + BK=triton.next_power_of_2(topk), + multibuffer=False, + ) + protect_residency_kernel[(rows, triton.cdiv(hot_size, 256))]( + sorted_rows, + workspace.metadata, + device_global_indices, + lru_slots, + workspace.protected, + TOP_K=topk, + HOT_SIZE=hot_size, + STATE_ROWS=device_global_indices.size(0), + BLOCK=256, + multibuffer=False, + ) + torch.cumsum(workspace.protected[:rows], dim=-1, dtype=torch.int32, out=workspace.hit_prefix[:rows]) + partition_residency_kernel[(rows, triton.cdiv(hot_size, 256))]( + workspace.metadata, + workspace.protected, + workspace.hit_prefix, + workspace.row_stats, + lru_slots, + workspace.lru_temp, + workspace.errors, + HOT_SIZE=hot_size, + MAX_UNION=max_union_rows, + STATE_ROWS=device_global_indices.size(0), + BLOCK=256, + multibuffer=False, + ) + update_lru_kernel[(rows, triton.cdiv(hot_size, 256))]( + workspace.metadata, + workspace.lru_temp, + lru_slots, + workspace.errors, + HOT_SIZE=hot_size, + BLOCK=256, + multibuffer=False, + ) + assign_residency_kernel[(rows,)]( + row_keys, + workspace.metadata, + workspace.first_misses, + hot_block_table, + hot_indices, + device_global_indices, + workspace.lru_temp, + workspace.errors, + swap_host_physical_rows if has_swaps else dummy, + swap_device_physical_rows if has_swaps else dummy, + swap_counts if has_swaps else dummy, + TOP_K=topk, + HOT_SIZE=hot_size, + HOT_BLOCK=hot_cache.size(1), + TABLE_STRIDE=hot_block_table.stride(0), + HAS_SWAPS=has_swaps, + BK=triton.next_power_of_2(topk), + multibuffer=False, + ) + finish_residency_kernel[(rows,)]( + workspace.metadata, + workspace.first_refs, + workspace.first_misses, + hot_indices, + attention_indices if attention_indices is not None else dummy, + miss_mask if miss_mask is not None else dummy, + valid_counts if valid_counts is not None else dummy, + workspace.errors, + TOP_K=topk, + HOT_BLOCK=hot_cache.size(1), + ATTENTION_BLOCK=attention_block_stride, + ATTENTION_STRIDE=attention_indices.stride(0) if attention_indices is not None else topk, + COUNT_STRIDE=valid_counts.stride(0) if valid_counts is not None else 1, + HAS_ATTENTION=attention_indices is not None, + HAS_MISS_MASK=miss_mask is not None, + HAS_COUNTS=valid_counts is not None, + BK=triton.next_power_of_2(topk), + multibuffer=False, + ) + if stats is not None: + residency_stats_kernel[(1,)]( + workspace.row_stats, + stats, + ROWS=rows, + BLOCK=triton.next_power_of_2(rows), + multibuffer=False, + ) + + def hisparse_gather_compact( + self, host_cache, hot_cache, miss_global_indices, miss_hot_indices, miss_counts + ) -> None: + self._gather(host_cache, hot_cache, miss_global_indices, miss_hot_indices, miss_counts, None, None, None, 0) + + def hisparse_gather_plan( + self, + host_cache, + hot_cache, + global_indices, + hot_indices, + miss_mask, + request_state_indices=None, + attention_indices=None, + attention_block_stride=0, + ) -> None: + self._gather( + host_cache, + hot_cache, + global_indices, + hot_indices, + None, + miss_mask, + request_state_indices, + attention_indices, + attention_block_stride, + ) + + def _gather( + self, host_cache, hot_cache, source, destination, counts, mask, state, attention, attention_block + ) -> None: + row_bytes, block_size, hot_rows, hot_stride = _cache_layout(host_cache, hot_cache) + device = hot_cache.device + if self.host_mapping.device != device: + raise ValueError("HiSparse host registration must belong to the cache device") + _check_tensor(source, "source", device, (torch.int32,), ndim=2) + _check_tensor(destination, "destination", device, (torch.int32,), shape=source.shape) + rows, topk = source.shape + if counts is not None: + _check_tensor(counts, "counts", device, (torch.int32,), shape=(rows,)) + elif mask is not None: + _check_tensor(mask, "mask", device, (torch.uint8, torch.bool, torch.int32), shape=source.shape) + else: + raise ValueError("HiSparse gather requires compact counts or a miss mask") + if state is not None: + _check_tensor(state, "state", device, (torch.int32,), shape=(rows,)) + if attention is not None: + _check_tensor(attention, "attention", device, (torch.int32,), shape=source.shape) + if attention_block < block_size: + raise ValueError("HiSparse attention block stride must cover a physical block") + if rows == 0: + return + gather_rows_kernel[(rows, topk)]( + self.host_mapping.address_of(host_cache), + hot_cache.view(torch.uint8), + source, + destination, + counts if counts is not None else source, + mask if mask is not None else source, + attention if attention is not None else source, + state if state is not None else source, + HOST_ROWS=host_cache.size(0), + HOT_ROWS=hot_rows, + ROW_BYTES=row_bytes, + HOT_BLOCK=block_size, + HOT_STRIDE_BYTES=hot_stride, + TOP_K=topk, + ATTENTION_BLOCK=attention_block, + COMPACT=counts is not None, + HAS_ATTENTION=attention is not None, + HAS_STATE=state is not None, + BLOCK=min(1024, triton.next_power_of_2(row_bytes)), + multibuffer=False, + ) + + @staticmethod + def hisparse_invalidate_written_slots( + device_global_indices, + request_state_indices, + req_id_per_token, + written_slots, + num_state_rows=None, + region_stride=None, + ) -> None: + device = device_global_indices.device + if device.type != "npu": + raise ValueError("HiSparse invalidation requires NPU state") + grouped = num_state_rows is not None or region_stride is not None + if grouped: + if num_state_rows is None or region_stride is None or num_state_rows < 0 or region_stride < 0: + raise ValueError("Grouped HiSparse invalidation requires nonnegative state and region sizes") + _check_tensor(device_global_indices, "group_table", device, (torch.int64,), ndim=1) + groups = device_global_indices.numel() + region = region_stride + state_rows = num_state_rows + else: + _check_tensor(device_global_indices, "device_global_indices", device, (torch.int32,), ndim=2) + groups = 1 + region = device_global_indices.size(1) + state_rows = device_global_indices.size(0) + _check_tensor(request_state_indices, "request_state_indices", device, (torch.int32,), ndim=1) + _check_tensor(written_slots, "written_slots", device, (torch.int32, torch.int64), ndim=1) + _check_tensor(req_id_per_token, "req_id_per_token", device, (torch.int32,), shape=written_slots.shape) + if written_slots.numel() == 0: + return + if region == 0 or state_rows == 0 or groups == 0: + return + invalidate_kernel[(written_slots.numel(), triton.cdiv(region, 256), groups)]( + device_global_indices, + request_state_indices, + req_id_per_token, + written_slots, + REQUESTS=request_state_indices.numel(), + STATE_ROWS=state_rows, + REGION=region, + BLOCK=256, + GROUPED=grouped, + multibuffer=False, + ) + + +def pack_attention_rows(packed_cache, physical_topk, selected_latent_out, selected_rope_out, latent_dim) -> None: + """Gather packed physical KV rows into the split, dense SFA input layout.""" + _check_packed_cache(packed_cache) + device = packed_cache.device + _check_tensor(physical_topk, "physical_topk", device, (torch.int32,), ndim=2, contiguous=False) + if physical_topk.stride(1) != 1: + raise ValueError("HiSparse attention top-k columns must be contiguous") + rows, topk = physical_topk.shape + width = packed_cache.size(-1) + if not 0 < latent_dim < width: + raise ValueError("HiSparse attention requires nonempty latent and RoPE rows") + _check_tensor(selected_latent_out, "selected_latent_out", device, (packed_cache.dtype,)) + _check_tensor(selected_rope_out, "selected_rope_out", device, (packed_cache.dtype,)) + if selected_latent_out.numel() != rows * topk * latent_dim or selected_rope_out.numel() != rows * topk * ( + width - latent_dim + ): + raise ValueError("HiSparse attention staging shapes do not match the selected rows") + if rows: + pack_attention_rows_kernel[(rows, topk)]( + packed_cache, + physical_topk, + selected_latent_out, + selected_rope_out, + K=topk, + CACHE_ROWS=packed_cache.size(0) * packed_cache.size(1), + CACHE_BLOCK=packed_cache.size(1), + CACHE_STRIDE=packed_cache.stride(0), + WIDTH=width, + LATENT=latent_dim, + INDEX_STRIDE=physical_topk.stride(0), + BLOCK=triton.next_power_of_2(width), + multibuffer=False, + ) + + +def scatter_packed_kv(k_nope, k_rope, packed_cache, slot_mapping, num_valid_rows=None) -> None: + """Write current KV with explicit page strides; ignore padding slots.""" + _check_packed_cache(packed_cache) + device = packed_cache.device + _check_tensor(slot_mapping, "slot_mapping", device, (torch.int32, torch.int64), ndim=1) + rows = slot_mapping.numel() + _check_payload_rows(k_nope, "k_nope", device, packed_cache.dtype, rows) + _check_payload_rows(k_rope, "k_rope", device, packed_cache.dtype, rows) + if num_valid_rows is not None: + _check_tensor(num_valid_rows, "num_valid_rows", device, (torch.int32, torch.int64)) + if num_valid_rows.numel() != 1: + raise ValueError("num_valid_rows must contain one device scalar") + latent_dim = math.prod(k_nope.shape[1:]) + rope_dim = math.prod(k_rope.shape[1:]) + if latent_dim + rope_dim != packed_cache.size(-1): + raise ValueError("HiSparse packed row width must equal latent + RoPE widths") + if rows: + scatter_packed_kv_kernel[(rows,)]( + k_nope, + k_rope, + packed_cache, + slot_mapping, + num_valid_rows if num_valid_rows is not None else slot_mapping, + CACHE_ROWS=packed_cache.size(0) * packed_cache.size(1), + CACHE_BLOCK=packed_cache.size(1), + CACHE_STRIDE=packed_cache.stride(0), + WIDTH=latent_dim + rope_dim, + LATENT=latent_dim, + LATENT_STRIDE=k_nope.stride(0), + ROPE_STRIDE=k_rope.stride(0), + HAS_VALID=num_valid_rows is not None, + BLOCK=triton.next_power_of_2(latent_dim + rope_dim), + multibuffer=False, + ) diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index c4762dd8bd0f..ca27b31b0943 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -179,6 +179,7 @@ def __init__( # both layernorms as extra args, so this flag has no effect here and is # accepted for signature compatibility only. fuse_qkv_rmsnorm: bool = False, + index_group_builder=None, ) -> None: nn.Module.__init__(self) self.hidden_size = hidden_size @@ -225,6 +226,7 @@ def __init__( dcp_q_replicate=self.dcp_q_replicate_enabled, use_sparse=mla_modules.is_sparse, indexer=ascend_indexer, + index_group_builder=index_group_builder, skip_topk=skip_topk, topk_indices_buffer=getattr(mla_modules, "topk_indices_buffer", None), non_causal_multi_token_decode=non_causal_multi_token_decode, @@ -249,17 +251,23 @@ def __init__( original_process_weights = self.mla_attn.process_weights_after_loading - def wrapped_process_weights(act_dtype: torch.dtype): + def wrapped_process_weights(act_dtype: torch.dtype | None = None): from vllm_ascend.attention.sfa_v1 import AscendSFAImpl if self.dcp_q_replicate_enabled: - self.mla_attn.impl.process_weights_after_loading(act_dtype) + self.mla_attn.impl.process_weights_after_loading( + self.mla_attn.act_dtype if act_dtype is None else act_dtype + ) elif not isinstance(self.mla_attn.impl, AscendSFAImpl): - # Both supported vLLM versions dispatch to the impl here. - original_process_weights(act_dtype) + if act_dtype is None: + original_process_weights() + else: + original_process_weights(act_dtype) else: # SFA disposes kv_b_proj, so bypass upstream's dense packing. - self.mla_attn.impl.process_weights_after_loading(act_dtype) + self.mla_attn.impl.process_weights_after_loading( + self.mla_attn.act_dtype if act_dtype is None else act_dtype + ) self.mla_attn.process_weights_after_loading = wrapped_process_weights diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py new file mode 100644 index 000000000000..9e383e003289 --- /dev/null +++ b/vllm_ascend/ops/triton/hisparse.py @@ -0,0 +1,546 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device-only HiSparse planning and mapped-host row copies for Ascend. + +Device sorts and binary searches form the union of a request's contiguous +verification rows. All selected slots are protected before any eviction. +Scratch belongs to the caller and is allocated before graph capture. +""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _translate( + token, + request, + state, + source, + resident, + HOST_ROWS: tl.constexpr, + SOURCE_STRIDE: tl.constexpr, + SOURCE_ROWS: tl.constexpr, + SOURCE_COLS: tl.constexpr, + SOURCE_BLOCK: tl.constexpr, + RESIDENT_STRIDE: tl.constexpr, + RESIDENT_ROWS: tl.constexpr, + RESIDENT_COLS: tl.constexpr, + RESIDENT_BLOCK: tl.constexpr, + RESIDENT_NULL: tl.constexpr, + HAS_SOURCE: tl.constexpr, + HAS_RESIDENT: tl.constexpr, +): + host = token + device = tl.full(token.shape, -1, tl.int32) + if HAS_SOURCE: + col = token // SOURCE_BLOCK + valid = (token >= 0) & (request >= 0) & (request < SOURCE_ROWS) & (col < SOURCE_COLS) + block = tl.load(source + request.to(tl.int64) * SOURCE_STRIDE + col, mask=valid, other=0) + host = tl.where(valid & (block > 0), block * SOURCE_BLOCK + token % SOURCE_BLOCK, -1) + if HAS_RESIDENT: + col = token // RESIDENT_BLOCK + valid = (token >= 0) & (state >= 0) & (state < RESIDENT_ROWS) & (col < RESIDENT_COLS) + block = tl.load(resident + state.to(tl.int64) * RESIDENT_STRIDE + col, mask=valid, other=RESIDENT_NULL) + device = tl.where( + valid & (block >= 0) & (block != RESIDENT_NULL), + block * RESIDENT_BLOCK + token % RESIDENT_BLOCK, + -1, + ) + host = tl.where((host >= 0) & (host < HOST_ROWS), host, -1) + return host, device + + +@triton.jit +def _lookup(sorted_keys, value, active, N: tl.constexpr): + # Keys pack (nonnegative host row, original position). Lower-bound also + # identifies the first reference when a query contains duplicate indices. + low = tl.full(value.shape, 0, tl.int32) + high = tl.full(value.shape, N, tl.int32) + target = value.to(tl.int64) << 32 + for _ in range(0, tl.constexpr(N.bit_length())): + middle = (low + high) // 2 + key = tl.load(sorted_keys + middle, mask=active & (middle < N), other=9223372036854775807) + left = key < target + low = tl.where(left, middle + 1, low) + high = tl.where(left, high, middle) + key = tl.load(sorted_keys + low, mask=active & (low < N), other=9223372036854775807) + found = active & (low < N) & ((key >> 32) == value) + return (key & 4294967295).to(tl.int32), found + + +@triton.jit +def prepare_residency_kernel( + indices, + requests, + states, + source, + resident, + dgi, + hot_indices, + resolved, + counts, + row_keys, + cache_keys, + metadata, + num_valid_rows, + ROWS: tl.constexpr, + TOP_K: tl.constexpr, + HOT_SIZE: tl.constexpr, + INPUT_STRIDE: tl.constexpr, + COUNT_STRIDE: tl.constexpr, + STATE_ROWS: tl.constexpr, + STATE_COUNT: tl.constexpr, + HOST_ROWS: tl.constexpr, + SOURCE_STRIDE: tl.constexpr, + SOURCE_ROWS: tl.constexpr, + SOURCE_COLS: tl.constexpr, + SOURCE_BLOCK: tl.constexpr, + RESIDENT_STRIDE: tl.constexpr, + RESIDENT_ROWS: tl.constexpr, + RESIDENT_COLS: tl.constexpr, + RESIDENT_BLOCK: tl.constexpr, + RESIDENT_NULL: tl.constexpr, + HAS_REQUESTS: tl.constexpr, + HAS_STATE: tl.constexpr, + HAS_SOURCE: tl.constexpr, + HAS_RESIDENT: tl.constexpr, + HAS_RESOLVED: tl.constexpr, + HAS_COUNTS: tl.constexpr, + HAS_VALID_ROWS: tl.constexpr, + BK: tl.constexpr, +): + row = tl.program_id(0) + active_rows = tl.minimum(tl.maximum(tl.load(num_valid_rows), 0), ROWS) if HAS_VALID_ROWS else ROWS + request = tl.load(requests + row) if HAS_REQUESTS else row + request = tl.where(row < active_rows, request, -1) + first = row + end = row + 1 + if HAS_REQUESTS: + searching = first > 0 + while searching: + same = tl.load(requests + first - 1, mask=first > 0, other=-2147483648) == request + same = same & (first > 0) + first -= same.to(tl.int32) + searching = same + searching = end < active_rows + while searching: + same = tl.load(requests + end, mask=end < active_rows, other=-2147483648) == request + same = same & (end < active_rows) + end += same.to(tl.int32) + searching = same + state = tl.load(states + request, mask=(request >= 0) & (request < STATE_COUNT), other=-1) if HAS_STATE else request + valid_state = (state >= 0) & (state < STATE_ROWS) + m = tl.arange(0, 8) + meta = tl.where(m == 0, request, tl.where(m == 1, state, tl.where(m == 2, first, end))) + tl.store(metadata + row * 8 + m, meta) + i = tl.arange(0, BK) + token = tl.load(indices + row * INPUT_STRIDE + i, mask=(i < TOP_K) & valid_state, other=-1) + host, device = _translate( + token, + request, + state, + source, + resident, + HOST_ROWS, + SOURCE_STRIDE, + SOURCE_ROWS, + SOURCE_COLS, + SOURCE_BLOCK, + RESIDENT_STRIDE, + RESIDENT_ROWS, + RESIDENT_COLS, + RESIDENT_BLOCK, + RESIDENT_NULL, + HAS_SOURCE, + HAS_RESIDENT, + ) + tl.store(hot_indices + row * TOP_K + i, device, mask=i < TOP_K) + if HAS_RESOLVED: + tl.store(resolved + row * TOP_K + i, host, mask=i < TOP_K) + if HAS_COUNTS: + tl.store(counts + row * COUNT_STRIDE, tl.sum(((i < TOP_K) & ((device >= 0) | (host >= 0))).to(tl.int32), 0)) + cold = (i < TOP_K) & (device < 0) & (host >= 0) + key = tl.where(cold, (host.to(tl.int64) << 32) | i.to(tl.int64), 9223372036854775807) + tl.store(row_keys + row * TOP_K + i, key, mask=i < TOP_K) + for start in range(0, HOT_SIZE, 256): + p = start + tl.arange(0, 256) + cached = tl.load( + dgi + state.to(tl.int64) * HOT_SIZE + p, mask=valid_state & (row == first) & (p < HOT_SIZE), other=-1 + ) + cache_key = tl.where(cached >= 0, (cached.to(tl.int64) << 32) | p.to(tl.int64), 9223372036854775807) + tl.store(cache_keys + row * HOT_SIZE + p, cache_key, mask=p < HOT_SIZE) + + +@triton.jit +def classify_residency_kernel( + row_keys, + sorted_rows, + sorted_cache, + metadata, + hot_table, + hot_indices, + first_refs, + first_misses, + row_stats, + TOP_K: tl.constexpr, + HOT_SIZE: tl.constexpr, + HOT_BLOCK: tl.constexpr, + TABLE_STRIDE: tl.constexpr, + BK: tl.constexpr, +): + row = tl.program_id(0) + request = tl.load(metadata + row * 8) + first = tl.load(metadata + row * 8 + 2) + i = tl.arange(0, BK) + key = tl.load(row_keys + row * TOP_K + i, mask=i < TOP_K, other=9223372036854775807) + host = (key >> 32).to(tl.int32) + active = (i < TOP_K) & (host != 2147483647) + slot, hit = _lookup(sorted_cache + first * HOT_SIZE, host, active, HOT_SIZE) + reference = tl.full((BK,), 2147483647, tl.int32) + for query in range(first, row + 1): + pos, found = _lookup(sorted_rows + query * TOP_K, host, active, TOP_K) + reference = tl.minimum(reference, tl.where(found, query * TOP_K + pos, 2147483647)) + miss = active & ~hit & (reference == row * TOP_K + i) + tl.store(first_refs + row * TOP_K + i, tl.where(active, reference, -1), mask=i < TOP_K) + tl.store(first_misses + row * TOP_K + i, miss, mask=i < TOP_K) + block = tl.load(hot_table + request.to(tl.int64) * TABLE_STRIDE + slot // HOT_BLOCK, mask=hit, other=0) + tl.store(hot_indices + row * TOP_K + i, block * HOT_BLOCK + slot % HOT_BLOCK, mask=hit) + # Independent padded records avoid adjacent programs sharing a transfer line. + s = tl.arange(0, 4) + misses = tl.sum(miss.to(tl.int32), 0) + references = tl.sum(active.to(tl.int32), 0) + tl.store(row_stats + row * 4 + s, tl.where(s == 0, references - misses, tl.where(s == 1, misses, 0))) + + +@triton.jit +def protect_residency_kernel( + sorted_rows, + metadata, + dgi, + lru, + protected, + TOP_K: tl.constexpr, + HOT_SIZE: tl.constexpr, + STATE_ROWS: tl.constexpr, + BLOCK: tl.constexpr, +): + first = tl.program_id(0) + tile = tl.program_id(1) + leader = tl.load(metadata + first * 8 + 2) + p = tile * BLOCK + tl.arange(0, BLOCK) + hit = tl.full((BLOCK,), False, tl.int1) + if first == leader: + state = tl.load(metadata + first * 8 + 1) + end = tl.load(metadata + first * 8 + 3) + valid = (state >= 0) & (state < STATE_ROWS) & (p < HOT_SIZE) + slot = tl.load(lru + state.to(tl.int64) * HOT_SIZE + p, mask=valid, other=-1).to(tl.int32) + valid = valid & (slot >= 0) & (slot < HOT_SIZE) + cached = tl.load(dgi + state.to(tl.int64) * HOT_SIZE + slot, mask=valid, other=-1) + for row in range(first, end): + _, found = _lookup(sorted_rows + row * TOP_K, cached, valid & (cached >= 0), TOP_K) + hit = hit | found + tl.store(protected + first * HOT_SIZE + p, hit, mask=p < HOT_SIZE) + + +@triton.jit +def partition_residency_kernel( + metadata, + protected, + hit_prefix, + row_stats, + lru, + temp_lru, + errors, + HOT_SIZE: tl.constexpr, + MAX_UNION: tl.constexpr, + STATE_ROWS: tl.constexpr, + BLOCK: tl.constexpr, +): + first = tl.program_id(0) + tile = tl.program_id(1) + leader = tl.load(metadata + first * 8 + 2) + if first == leader: + state = tl.load(metadata + first * 8 + 1) + end = tl.load(metadata + first * 8 + 3) + p = tile * BLOCK + tl.arange(0, BLOCK) + valid = (state >= 0) & (state < STATE_ROWS) & (p < HOT_SIZE) + hit = tl.load(protected + first * HOT_SIZE + p, mask=valid, other=0) != 0 + hits = tl.load(hit_prefix + first * HOT_SIZE + HOT_SIZE - 1) + misses = 0 + for row in range(first, end): + misses += tl.load(row_stats + row * 4 + 1).to(tl.int32) + union = hits + misses + failed = union > MAX_UNION + if tile == 0: + err = tl.arange(0, 8) + tl.store(errors + first * 8 + err, tl.where(failed, union, 0)) + tl.store(metadata + first * 8 + 4, hits) + tl.store(metadata + first * 8 + 5, misses) + tl.device_assert(~failed, "HiSparse request union exceeds hot capacity") + if (union > 0) & ~failed: + slot = tl.load(lru + state.to(tl.int64) * HOT_SIZE + p, mask=valid, other=-1).to(tl.int32) + prefix = tl.load(hit_prefix + first * HOT_SIZE + p, mask=valid, other=0) + hit_rank = prefix - 1 + evictable = valid & ~hit + evict_rank = p - prefix + temp = temp_lru + first * HOT_SIZE + tl.store(temp + hit_rank, slot, mask=hit) + tl.store(temp + HOT_SIZE - 1 - evict_rank, slot, mask=evictable) + + +@triton.jit +def update_lru_kernel(metadata, temp_lru, lru, errors, HOT_SIZE: tl.constexpr, BLOCK: tl.constexpr): + first = tl.program_id(0) + tile = tl.program_id(1) + leader = tl.load(metadata + first * 8 + 2) + if first == leader: + hits = tl.load(metadata + first * 8 + 4) + misses = tl.load(metadata + first * 8 + 5) + failed = tl.load(errors + first * 8) != 0 + if ((hits + misses) > 0) & ~failed: + state = tl.load(metadata + first * 8 + 1) + p = tile * BLOCK + tl.arange(0, BLOCK) + remaining = HOT_SIZE - hits - misses + source = tl.where( + p < remaining, + HOT_SIZE - 1 - misses - p, + tl.where(p < remaining + misses, HOT_SIZE - 1 - (p - remaining), p - remaining - misses), + ) + updated = tl.load(temp_lru + first * HOT_SIZE + source, mask=p < HOT_SIZE, other=0) + tl.store(lru + state.to(tl.int64) * HOT_SIZE + p, updated, mask=p < HOT_SIZE) + + +@triton.jit +def assign_residency_kernel( + row_keys, + metadata, + first_misses, + hot_table, + hot_indices, + dgi, + temp_lru, + errors, + swap_host, + swap_device, + swap_counts, + TOP_K: tl.constexpr, + HOT_SIZE: tl.constexpr, + HOT_BLOCK: tl.constexpr, + TABLE_STRIDE: tl.constexpr, + HAS_SWAPS: tl.constexpr, + BK: tl.constexpr, +): + first = tl.program_id(0) + leader = tl.load(metadata + first * 8 + 2) + if first == leader: + request = tl.load(metadata + first * 8) + state = tl.load(metadata + first * 8 + 1) + end = tl.load(metadata + first * 8 + 3) + misses = tl.load(metadata + first * 8 + 5) + failed = tl.load(errors + first * 8) != 0 + if HAS_SWAPS: + for row in range(first, end): + count = tl.minimum(TOP_K, tl.maximum(0, misses - (row - first) * TOP_K)) + tl.store(swap_counts + row, tl.where(failed, 0, count)) + if (misses > 0) & ~failed: + i = tl.arange(0, BK) + consumed = 0 + temp = temp_lru + first * HOT_SIZE + for row in range(first, end): + miss = tl.load(first_misses + row * TOP_K + i, mask=i < TOP_K, other=0) != 0 + rank = tl.cumsum(miss.to(tl.int32), 0) - 1 + consumed + evict = tl.load(temp + HOT_SIZE - 1 - rank, mask=miss, other=-1) + good = miss & (evict >= 0) & (evict < HOT_SIZE) + key = tl.load(row_keys + row * TOP_K + i, mask=miss, other=0) + host = (key >> 32).to(tl.int32) + tl.store(dgi + state.to(tl.int64) * HOT_SIZE + evict, host, mask=good) + block = tl.load( + hot_table + request.to(tl.int64) * TABLE_STRIDE + evict // HOT_BLOCK, mask=good, other=0 + ) + physical = tl.where(good, block * HOT_BLOCK + evict % HOT_BLOCK, -1) + tl.store(hot_indices + row * TOP_K + i, physical, mask=miss) + if HAS_SWAPS: + tl.store(swap_host + first * TOP_K + rank, host, mask=miss) + tl.store(swap_device + first * TOP_K + rank, physical, mask=miss) + consumed += tl.sum(miss.to(tl.int32), 0) + + +@triton.jit +def finish_residency_kernel( + metadata, + first_refs, + first_misses, + hot_indices, + attention, + miss_mask, + counts, + errors, + TOP_K: tl.constexpr, + HOT_BLOCK: tl.constexpr, + ATTENTION_BLOCK: tl.constexpr, + ATTENTION_STRIDE: tl.constexpr, + COUNT_STRIDE: tl.constexpr, + HAS_ATTENTION: tl.constexpr, + HAS_MISS_MASK: tl.constexpr, + HAS_COUNTS: tl.constexpr, + BK: tl.constexpr, +): + row = tl.program_id(0) + i = tl.arange(0, BK) + first = tl.load(metadata + row * 8 + 2) + failed = tl.load(errors + first * 8) != 0 + reference = tl.load(first_refs + row * TOP_K + i, mask=i < TOP_K, other=-1) + own = tl.load(hot_indices + row * TOP_K + i, mask=i < TOP_K, other=-1) + physical = tl.load(hot_indices + reference, mask=(i < TOP_K) & (reference >= 0), other=-1) + physical = tl.where(failed, -1, tl.where(reference >= 0, physical, own)) + tl.store(hot_indices + row * TOP_K + i, physical, mask=i < TOP_K) + if HAS_ATTENTION: + encoded = tl.where(physical >= 0, physical // HOT_BLOCK * ATTENTION_BLOCK + physical % HOT_BLOCK, -1) + tl.store(attention + row * ATTENTION_STRIDE + i, encoded, mask=i < TOP_K) + if HAS_MISS_MASK: + miss = tl.load(first_misses + row * TOP_K + i, mask=i < TOP_K, other=0) + tl.store(miss_mask + row * TOP_K + i, tl.where(failed, 0, miss), mask=i < TOP_K) + if HAS_COUNTS: + if failed: + tl.store(counts + row * COUNT_STRIDE, 0) + + +@triton.jit +def residency_stats_kernel(row_stats, stats, ROWS: tl.constexpr, BLOCK: tl.constexpr): + i = tl.arange(0, BLOCK) + hit = tl.load(row_stats + i * 4, mask=i < ROWS, other=0) + miss = tl.load(row_stats + i * 4 + 1, mask=i < ROWS, other=0) + tl.store(stats, tl.load(stats) + tl.sum(hit, 0)) + tl.store(stats + 1, tl.load(stats + 1) + tl.sum(miss, 0)) + + +@triton.jit +def gather_rows_kernel( + host_address, + hot, + source, + destination, + counts, + mask_ptr, + attention, + state, + HOST_ROWS: tl.constexpr, + HOT_ROWS: tl.constexpr, + ROW_BYTES: tl.constexpr, + HOT_BLOCK: tl.constexpr, + HOT_STRIDE_BYTES: tl.constexpr, + TOP_K: tl.constexpr, + ATTENTION_BLOCK: tl.constexpr, + COMPACT: tl.constexpr, + HAS_ATTENTION: tl.constexpr, + HAS_STATE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + col = tl.program_id(1) + offset = tl.arange(0, BLOCK) + index = row * TOP_K + col + dst = tl.load(destination + index) + padding = tl.load(state + row) < 0 if HAS_STATE else False + if HAS_ATTENTION: + value = tl.where((dst >= 0) & ~padding, dst // HOT_BLOCK * ATTENTION_BLOCK + dst % HOT_BLOCK, -1) + tl.store(attention + index, value) + if COMPACT: + copy = col < tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) + else: + copy = tl.load(mask_ptr + index) != 0 + if copy & ~padding: + src = tl.load(source + index) + valid = (src >= 0) & (dst >= 0) & (dst < HOT_ROWS) + if valid: + host = host_address.to(tl.pointer_type(tl.uint8)) + destination_offset = dst.to(tl.int64) // HOT_BLOCK * HOT_STRIDE_BYTES + dst % HOT_BLOCK * ROW_BYTES + for start in range(0, ROW_BYTES, BLOCK): + payload = tl.load( + host + src.to(tl.int64) * ROW_BYTES + start + offset, + mask=(src < HOST_ROWS) & (start + offset < ROW_BYTES), + other=0, + ) + tl.store(hot + destination_offset + start + offset, payload, mask=start + offset < ROW_BYTES) + + +@triton.jit +def invalidate_kernel( + dgi, + states, + requests, + written, + REQUESTS: tl.constexpr, + STATE_ROWS: tl.constexpr, + REGION: tl.constexpr, + BLOCK: tl.constexpr, + GROUPED: tl.constexpr = False, +): + if GROUPED: + dgi = tl.load(dgi + tl.program_id(2)).to(tl.pointer_type(tl.int32)) + token = tl.program_id(0) + offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + request = tl.load(requests + token) + state = tl.load(states + request, mask=(request >= 0) & (request < REQUESTS), other=-1) + value = tl.load(written + token) + valid = (state >= 0) & (state < STATE_ROWS) & (value >= 0) & (offset < REGION) + ptr = dgi + state.to(tl.int64) * REGION + offset + cached = tl.load(ptr, mask=valid, other=-1) + tl.store(ptr, -1, mask=valid & (cached == value)) + + +@triton.jit +def pack_attention_rows_kernel( + cache, + indices, + latent, + rope, + K: tl.constexpr, + CACHE_ROWS: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + WIDTH: tl.constexpr, + LATENT: tl.constexpr, + INDEX_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + col = tl.program_id(1) + i = tl.arange(0, BLOCK) + physical = tl.load(indices + row * INDEX_STRIDE + col) + valid = (physical >= 0) & (physical < CACHE_ROWS) + base = physical.to(tl.int64) // CACHE_BLOCK * CACHE_STRIDE + physical % CACHE_BLOCK * WIDTH + value = tl.load(cache + base + i, mask=valid & (i < WIDTH), other=0) + tl.store(latent + (row * K + col).to(tl.int64) * LATENT + i, value, mask=i < LATENT) + tl.store( + rope + (row * K + col).to(tl.int64) * (WIDTH - LATENT) + i - LATENT, value, mask=(i >= LATENT) & (i < WIDTH) + ) + + +@triton.jit +def scatter_packed_kv_kernel( + latent, + rope, + cache, + slots, + valid_rows, + CACHE_ROWS: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + WIDTH: tl.constexpr, + LATENT: tl.constexpr, + LATENT_STRIDE: tl.constexpr, + ROPE_STRIDE: tl.constexpr, + HAS_VALID: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + i = tl.arange(0, BLOCK) + slot = tl.load(slots + row) + valid = (slot >= 0) & (slot < CACHE_ROWS) + if HAS_VALID: + valid = valid & (row < tl.load(valid_rows)) + a = tl.load(latent + row.to(tl.int64) * LATENT_STRIDE + i, mask=valid & (i < LATENT), other=0) + b = tl.load(rope + row.to(tl.int64) * ROPE_STRIDE + i - LATENT, mask=valid & (i >= LATENT) & (i < WIDTH), other=0) + base = slot.to(tl.int64) // CACHE_BLOCK * CACHE_STRIDE + slot % CACHE_BLOCK * WIDTH + tl.store(cache + base + i, tl.where(i < LATENT, a, b), mask=valid & (i < WIDTH)) diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 52d40b46ac7e..aa95a48262c3 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1861,3 +1861,14 @@ # Remove this patch when upstream provides a backend capability hook for # enabling MoE sequence parallelism with DP=1. # + +# ** File: worker/patch_hisparse_streams.py ** +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# Target: vllm.v1.hisparse.runtime and MLA index_group module-local Torch bindings. +# Why: PTA 2.10 generic event waits do not propagate graph capture to CANN +# workspace allocation in the same way as native NPU stream/event bindings. +# How: Install a module-local facade before HiSparse cache/graph construction. +# Preserve existing event points, algorithms and actual global Torch state. +# Related PR: HiSparse foundation split from #19109; no Core modification. +# Future Plan: Remove after PTA's generic bindings satisfy the same changing +# capture/replay and workspace contract, or Core exposes stream factories. diff --git a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py index 588178f62352..7dcfa5a78ec7 100644 --- a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py +++ b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py @@ -29,6 +29,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, + KVCacheGroupRole, KVCacheSpec, MambaSpec, UniformTypeKVCacheSpecs, @@ -555,6 +556,11 @@ def get_kv_cache_coordinator( # type: ignore[misc] ) # vLLM main (#54736) added allow_partial_hash_hits. hybrid_kwargs["allow_partial_hash_hits"] = allow_partial_hash_hits + if any(group.role == KVCacheGroupRole.HISPARSE_SOURCE for group in kv_cache_config.kv_cache_groups): + # Core owns HiSparse admission, HMA roles, eviction and prefix state. + return _orig_get_kv_cache_coordinator( + **{name: value for name, value in hybrid_kwargs.items() if name != "eagle_attn_layer_names"} + ) # Private state blocks are excluded by the block-size resolver. Keep # their independent manager geometry when prefix caching is disabled too. num_cacheable_groups = sum(is_prefix_cacheable(group.kv_cache_spec) for group in kv_cache_config.kv_cache_groups) diff --git a/vllm_ascend/patch/worker/patch_deepseek_v2.py b/vllm_ascend/patch/worker/patch_deepseek_v2.py index e5e158568dc8..5cf2bc1f325d 100644 --- a/vllm_ascend/patch/worker/patch_deepseek_v2.py +++ b/vllm_ascend/patch/worker/patch_deepseek_v2.py @@ -89,8 +89,6 @@ def _deepseek_v2_mla_attention_init( ) -> None: # Zero-argument super() is unavailable because this function is defined # outside the class and assigned to DeepseekV2MLAAttention.__init__. - # index_group_builder is a GPU HiSparse feature; the Ascend MLA stack - # does not consume it, so it is accepted for signature compatibility. nn.Module.__init__(self) self.hidden_size = hidden_size @@ -286,6 +284,7 @@ def _deepseek_v2_mla_attention_init( quant_config, prefix, skip_topk=_skip_topk and not mtp_layer, + index_group_builder=index_group_builder, non_causal_multi_token_decode=non_causal_multi_token_decode, ) diff --git a/vllm_ascend/patch/worker/patch_hisparse_streams.py b/vllm_ascend/patch/worker/patch_hisparse_streams.py new file mode 100644 index 000000000000..f42b43c56b3f --- /dev/null +++ b/vllm_ascend/patch/worker/patch_hisparse_streams.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Use PTA native stream/event bindings only inside common HiSparse modules. + +PTA 2.10's generic event wait does not flush capture propagation before CANN +workspace allocation; the native event binding does. Keep common event points +and algorithms intact, using the platform bindings through a module-local +facade. Remove this adapter when PTA's generic bindings have equivalent graph +capture behavior. The actual torch module and model-runner timing are untouched. +""" + +from __future__ import annotations + +import importlib +import threading + +import torch as _torch +from vllm.platforms import current_platform + +_COMMON_MODULES = ("vllm.v1.hisparse.runtime", "vllm.v1.attention.backends.mla.index_group") +_facade = None +_original_group_current_stream = None + + +def _is_npu(device=None) -> bool: + if device is None or isinstance(device, int): + return current_platform.device_type == "npu" + return _torch.device(device).type == "npu" + + +def _native_stream(stream): + if stream is None or isinstance(stream, _torch.npu.Stream): + return stream + if not isinstance(stream, _torch.Stream) or stream.device.type != "npu": + raise TypeError("A HiSparse native NPU event requires an NPU stream") + # Reconstruct the same pooled c10 stream, never an ExternalStream. + return _torch.npu.Stream( + stream_id=stream.stream_id, device_index=stream.device_index, device_type=stream.device_type + ) + + +class HiSparseStream(_torch.npu.Stream): + """Native stream with the context-manager interface used by common code.""" + + def __new__(cls, device=None, priority=0, **kwargs): + if not _is_npu(device): + return _torch.Stream(device=device, priority=priority, **kwargs) + stream = super().__new__(cls, device=device, priority=priority, **kwargs) + stream._hisparse_contexts = threading.local() + return stream + + def __enter__(self): + context = _torch.npu.stream(self) + context.__enter__() + contexts = getattr(self._hisparse_contexts, "stack", None) + if contexts is None: + contexts = self._hisparse_contexts.stack = [] + contexts.append(context) + return self + + def __exit__(self, *exc): + return self._hisparse_contexts.stack.pop().__exit__(*exc) + + def wait_stream(self, stream): + return super().wait_stream(_native_stream(stream)) + + +class HiSparseEvent(_torch.npu.Event): + """Native event; never reinterpret a generic Stream as a native object.""" + + def __new__(cls, *args, **kwargs): + if not _is_npu(): + return _torch.Event(*args, **kwargs) + return super().__new__(cls, *args, **kwargs) + + def record(self, stream=None): + return super().record(_native_stream(stream)) + + def wait(self, stream=None): + return super().wait(_native_stream(stream)) + + +def native_current_stream(device=None): + if _is_npu(device): + return _torch.npu.current_stream(device) + return _torch.accelerator.current_stream(device) + + +def create_event(*args, **kwargs): + return HiSparseEvent(*args, **kwargs) + + +def _group_current_stream(): + if _is_npu(): + return native_current_stream() + return _original_group_current_stream() + + +class _AcceleratorFacade: + current_stream = staticmethod(native_current_stream) + + def __getattr__(self, name): + return getattr(_torch.accelerator, name) + + +class _TorchFacade: + Stream = HiSparseStream + Event = HiSparseEvent + accelerator = _AcceleratorFacade() + + def __getattr__(self, name): + return getattr(_torch, name) + + +def install() -> bool: + """Cold activation before the first index-group layer is registered.""" + global _facade, _original_group_current_stream + if not _is_npu(): + return False + modules = [importlib.import_module(name) for name in _COMMON_MODULES] + index_module = modules[1] + if _facade is not None: + if ( + not all(module.torch is _facade for module in modules) + or index_module.current_stream is not _group_current_stream + ): + raise RuntimeError("HiSparse stream facade changed after installation") + return True + # Install only before cache/graph construction. Leave the actual Torch + # module and all algorithm/event points intact; do not unwrap other owners. + if not callable(getattr(index_module, "current_stream", None)): + raise RuntimeError("HiSparse requires Core's current-stream interface") + if any(module.torch is not _torch for module in modules): + raise RuntimeError("HiSparse common torch binding was already replaced") + facade = _TorchFacade() + _original_group_current_stream = index_module.current_stream + for module in modules: + module.torch = facade + index_module.current_stream = _group_current_stream + _facade = facade + return True diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 6093d8f75f00..98058a9f6518 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -165,6 +165,10 @@ def is_cumem_allocator_available(self) -> bool: def is_pin_memory_available(cls): return True + @classmethod + def get_hisparse_backend_cls(cls) -> str: + return "vllm_ascend.distributed.kv_transfer.hisparse.backend.AscendHiSparseBackend" + @classmethod def opaque_attention_op(cls) -> bool: return True diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 2e45124ed14a..da4615c3ca33 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -273,6 +273,12 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: mamba_specs[layer_name] = spec continue + if isinstance(attn_module, MLAAttention) and vllm_config.attention_config.hisparse_config is not None: + # Keep Core's row-packed source spec and actual index-group ownership. + kv_cache_spec[layer_name] = replace( + spec, is_index_group_leader=attn_module.hisparse_cache.runtime.is_group_leader + ) + continue if isinstance(attn_module, MLAAttention): cache_sparse_sfa_c8 = False cache_sparse_sfa_turboquant = False @@ -418,6 +424,12 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: mamba_specs[layer_name] = replace(spec, page_size_padded=common_page_size) kv_cache_spec.update(mamba_specs) + if vllm_config.attention_config.hisparse_config is not None: + from vllm.v1.hisparse.binding import resolve_hisparse_specs + + from vllm_ascend.distributed.kv_transfer.hisparse.binding import align_hisparse_indexer_specs + + return align_hisparse_indexer_specs(resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers)) return kv_cache_spec @@ -563,6 +575,9 @@ def build_attn_metadata( plan = get_dsa_attn_kv_plan(dsa_builder.vllm_config, dsa_builder.compressor_ratio) formatted_slot_mappings = plan.format_dsa_slot_mapping(slot_mappings[:, :num_input_tokens], block_sizes) for i, kv_cache_spec in enumerate(kv_cache_groups): + # Resident/hot storage groups have no Attention metadata consumer. + if not attn_groups[i]: + continue block_table = block_tables[i] slot_mapping = slot_mappings[i] # Hybrid drafters can configure causality per KV cache group. diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 3ad40f92a70d..9fb3820d551d 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -777,6 +777,9 @@ def _scale_kv_cache_memory_for_multi_group(self, available_memory: int) -> int: derives a num_blocks (and block pool) small enough for the per-layer buffers to fit. """ + if self.vllm_config.attention_config.hisparse_config is not None: + # Core budgets the shared Resident/Hot/Indexer slab, not per-layer buffers. + return available_memory kv_cache_spec = self.get_kv_cache_spec() if not isinstance(kv_cache_spec, dict): return available_memory From 5c1ed8934f7a9d1979fc502df6d23cea325dfc44 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 16:44:49 +0800 Subject: [PATCH 002/120] docs(hisparse): add foundation release note Signed-off-by: zhenwenqi2024 --- release_notes/fragments/20926.yaml | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 release_notes/fragments/20926.yaml diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml new file mode 100644 index 000000000000..17d0715b4b32 --- /dev/null +++ b/release_notes/fragments/20926.yaml @@ -0,0 +1,8 @@ +pr: 20926 +entries: + - category: feature + text: >- + Added an experimental Ascend HiSparse backend and unquantized sparse + MLA/SFA data path, reusing Core allocation and lifecycle interfaces. + Device and full-model qualification of the foundation remains pending. + hardware: [A5] From 8a20e1c56092c14df4cb313453d2abe2ba919c79 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 17:52:55 +0800 Subject: [PATCH 003/120] refactor(hisparse): reuse Core generic stream and event APIs Remove the module-local Torch facade and its native stream/context wrappers. Preserve Core event placement and add a changing-graph fork/join regression with CANN allocation. Generic/native component comparisons and extracted merged-Core stream methods passed on Ascend 950; complete package and model qualification remain pending. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 23 +-- .../one_card/test_hisparse_ops.py | 72 ++++++++- vllm_ascend/attention/sfa_hisparse.py | 3 - vllm_ascend/patch/__init__.py | 11 -- .../patch/worker/patch_hisparse_streams.py | 141 ------------------ 5 files changed, 85 insertions(+), 165 deletions(-) delete mode 100644 vllm_ascend/patch/worker/patch_hisparse_streams.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 5db0f473da96..d5d15457d537 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -72,10 +72,10 @@ Use existing interfaces, rather than adding unused placeholder hooks: leader's scratch only on that group's ordered stream. Independent groups and current request/state mappings remain separate. -The stream facade is an isolated PTA compatibility adapter in the two common -HiSparse modules. It preserves event placement and does not replace the actual -global Torch module. Remove it when generic PTA capture/workspace behavior is -equivalent or Core provides the matching factories. +Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces are +used directly. The Ascend backend preserves Core's event placement and does not +replace its Torch bindings. The NPU tests cover cross-stream waits followed by +CANN workspace allocation and changing graph replays. Follow-up contributions should be separate and measurable: packed C8/indexer formats, native resolver/refill improvements, Host-writeback optimization, metadata @@ -116,8 +116,13 @@ pytest --confcutdir=tests/e2e/pull_request/one_card \ tests/e2e/pull_request/one_card/test_hisparse_selection_order.py ``` -These cases require an appropriate NPU environment and are not claimed to have -passed on this new split. Fresh full-model correctness, dataset accuracy, -eager/Graph lifecycle and same-workload performance remain prerequisites for -marking the foundation ready. Existing C8 benchmark results are intentionally -not presented as this implementation's performance. +The full modules have not been rerun on this new split. Isolated stream-boundary +checks on Ascend 950 with Torch 2.10.0 / torch-npu 2.10.0.post7 passed for generic +and native interfaces, including the merged Core methods' leader/follower and +lazy-event paths. Those checks substitute small CANN work for resolver/refill; +they do not qualify the complete HiSparse graph or model. + +Fresh full-model correctness, dataset accuracy, eager/Graph lifecycle and +same-workload performance remain prerequisites for marking the foundation +ready. Existing C8 benchmark results are intentionally not presented as this +implementation's performance. diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 396fafaff923..483887a694a8 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -6,7 +6,7 @@ import pytest import torch -import torch_npu # noqa: F401 +import torch_npu from vllm_ascend.ops.hisparse import ( HiSparseNPUOps, @@ -25,6 +25,76 @@ def _require_mapped_host_device(): torch.npu.set_device(0) +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("wait_event", [False, True]) +def test_generic_stream_waits_before_cann_workspace(capture, wait_event): + """Core's generic fork/join must propagate capture before CANN allocation.""" + initial = (torch.arange(64 * 64).reshape(64, 64) % 23 + 1).bfloat16() + x = initial.npu() + weight = torch.eye(64, dtype=x.dtype, device=x.device) + gamma = torch.ones(64, dtype=x.dtype, device=x.device) + output = torch.empty_like(x) + copy = torch.Stream(device=x.device) + ready = torch.Event() + done = None + owners = [] + + def forward(): + nonlocal done + compute = torch.accelerator.current_stream() + x.add_(1) + ready.record(compute) + if wait_event: + copy.wait_event(ready) + else: + copy.wait_stream(compute) + with copy: + # Keep allocation immediately after the wait. A synchronization or + # capture-state query here can hide a Host-queue propagation error. + product = torch.mm(x, weight) + normalized, _ = torch_npu.npu_rms_norm(product, gamma, epsilon=1e-6) + result = product + normalized + owners.extend((product, normalized, result)) + if done is None: + done = torch.Event() + done.record(copy) + compute.wait_event(done) + consumed = torch.mm(result, weight) + owners.append(consumed) + output.copy_(consumed) + + for _ in range(3): + forward() + torch.npu.synchronize() + owners.clear() + done = None # Core also creates per-layer events lazily during capture. + graph = None + try: + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + forward() + for increment in (0, 3, 7, 11, 19): + values = initial + increment + x.copy_(values) + torch.npu.synchronize() + if graph is None: + forward() + else: + graph.replay() + torch.npu.synchronize() + product = values.float() + 1 + normalized = (product * torch.rsqrt(product.square().mean(-1, keepdim=True) + 1e-6)).bfloat16() + expected = product.bfloat16() + normalized + torch.testing.assert_close(output.cpu(), expected, atol=0.25, rtol=0) + if graph is None: + owners.clear() + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + + def reference(topk, requests, states, source, resident, hot_table, dgi, lru, block): dgi = dgi.clone() lru = lru.clone() diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 56eec31e7555..d42f8d8fccea 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -101,9 +101,6 @@ def __init__(self, *args, **kwargs): if self.topk_indices_buffer is None: raise ValueError("Ascend HiSparse requires a model sparse-index group") builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) - from vllm_ascend.patch.worker.patch_hisparse_streams import install - - install() self.index_group, self.index_group_index = builder.register_layer( self.has_indexer and not self.skip_topk, self.vllm_config, diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index aa95a48262c3..52d40b46ac7e 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1861,14 +1861,3 @@ # Remove this patch when upstream provides a backend capability hook for # enabling MoE sequence parallelism with DP=1. # - -# ** File: worker/patch_hisparse_streams.py ** -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# Target: vllm.v1.hisparse.runtime and MLA index_group module-local Torch bindings. -# Why: PTA 2.10 generic event waits do not propagate graph capture to CANN -# workspace allocation in the same way as native NPU stream/event bindings. -# How: Install a module-local facade before HiSparse cache/graph construction. -# Preserve existing event points, algorithms and actual global Torch state. -# Related PR: HiSparse foundation split from #19109; no Core modification. -# Future Plan: Remove after PTA's generic bindings satisfy the same changing -# capture/replay and workspace contract, or Core exposes stream factories. diff --git a/vllm_ascend/patch/worker/patch_hisparse_streams.py b/vllm_ascend/patch/worker/patch_hisparse_streams.py deleted file mode 100644 index f42b43c56b3f..000000000000 --- a/vllm_ascend/patch/worker/patch_hisparse_streams.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Use PTA native stream/event bindings only inside common HiSparse modules. - -PTA 2.10's generic event wait does not flush capture propagation before CANN -workspace allocation; the native event binding does. Keep common event points -and algorithms intact, using the platform bindings through a module-local -facade. Remove this adapter when PTA's generic bindings have equivalent graph -capture behavior. The actual torch module and model-runner timing are untouched. -""" - -from __future__ import annotations - -import importlib -import threading - -import torch as _torch -from vllm.platforms import current_platform - -_COMMON_MODULES = ("vllm.v1.hisparse.runtime", "vllm.v1.attention.backends.mla.index_group") -_facade = None -_original_group_current_stream = None - - -def _is_npu(device=None) -> bool: - if device is None or isinstance(device, int): - return current_platform.device_type == "npu" - return _torch.device(device).type == "npu" - - -def _native_stream(stream): - if stream is None or isinstance(stream, _torch.npu.Stream): - return stream - if not isinstance(stream, _torch.Stream) or stream.device.type != "npu": - raise TypeError("A HiSparse native NPU event requires an NPU stream") - # Reconstruct the same pooled c10 stream, never an ExternalStream. - return _torch.npu.Stream( - stream_id=stream.stream_id, device_index=stream.device_index, device_type=stream.device_type - ) - - -class HiSparseStream(_torch.npu.Stream): - """Native stream with the context-manager interface used by common code.""" - - def __new__(cls, device=None, priority=0, **kwargs): - if not _is_npu(device): - return _torch.Stream(device=device, priority=priority, **kwargs) - stream = super().__new__(cls, device=device, priority=priority, **kwargs) - stream._hisparse_contexts = threading.local() - return stream - - def __enter__(self): - context = _torch.npu.stream(self) - context.__enter__() - contexts = getattr(self._hisparse_contexts, "stack", None) - if contexts is None: - contexts = self._hisparse_contexts.stack = [] - contexts.append(context) - return self - - def __exit__(self, *exc): - return self._hisparse_contexts.stack.pop().__exit__(*exc) - - def wait_stream(self, stream): - return super().wait_stream(_native_stream(stream)) - - -class HiSparseEvent(_torch.npu.Event): - """Native event; never reinterpret a generic Stream as a native object.""" - - def __new__(cls, *args, **kwargs): - if not _is_npu(): - return _torch.Event(*args, **kwargs) - return super().__new__(cls, *args, **kwargs) - - def record(self, stream=None): - return super().record(_native_stream(stream)) - - def wait(self, stream=None): - return super().wait(_native_stream(stream)) - - -def native_current_stream(device=None): - if _is_npu(device): - return _torch.npu.current_stream(device) - return _torch.accelerator.current_stream(device) - - -def create_event(*args, **kwargs): - return HiSparseEvent(*args, **kwargs) - - -def _group_current_stream(): - if _is_npu(): - return native_current_stream() - return _original_group_current_stream() - - -class _AcceleratorFacade: - current_stream = staticmethod(native_current_stream) - - def __getattr__(self, name): - return getattr(_torch.accelerator, name) - - -class _TorchFacade: - Stream = HiSparseStream - Event = HiSparseEvent - accelerator = _AcceleratorFacade() - - def __getattr__(self, name): - return getattr(_torch, name) - - -def install() -> bool: - """Cold activation before the first index-group layer is registered.""" - global _facade, _original_group_current_stream - if not _is_npu(): - return False - modules = [importlib.import_module(name) for name in _COMMON_MODULES] - index_module = modules[1] - if _facade is not None: - if ( - not all(module.torch is _facade for module in modules) - or index_module.current_stream is not _group_current_stream - ): - raise RuntimeError("HiSparse stream facade changed after installation") - return True - # Install only before cache/graph construction. Leave the actual Torch - # module and all algorithm/event points intact; do not unwrap other owners. - if not callable(getattr(index_module, "current_stream", None)): - raise RuntimeError("HiSparse requires Core's current-stream interface") - if any(module.torch is not _torch for module in modules): - raise RuntimeError("HiSparse common torch binding was already replaced") - facade = _TorchFacade() - _original_group_current_stream = index_module.current_stream - for module in modules: - module.torch = facade - index_module.current_stream = _group_current_stream - _facade = facade - return True From c500c536a18e54f44fb319009b90513baac93d1c Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 18:06:29 +0800 Subject: [PATCH 004/120] refactor(hisparse): pass index groups through Core MLA modules Use MLAModules.index_group_builder for the model-owned group builder and remove the extra Ascend wrapper keyword. Preserve existing checkpoint, DCP, PP and post-load behavior. Source-isolated checks use actual merged Core modules/group builders and both constructor bodies; 16 configurations and 26 CPU regressions pass. This is not NPU or full-model qualification. Signed-off-by: zhenwenqi2024 --- docs/source/user_guide/feature_guide/hisparse.md | 3 +++ vllm_ascend/ops/mla.py | 3 +-- vllm_ascend/patch/worker/patch_deepseek_v2.py | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index d5d15457d537..7ad689105f49 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -77,6 +77,9 @@ used directly. The Ascend backend preserves Core's event placement and does not replace its Torch bindings. The NPU tests cover cross-stream waits followed by CANN workspace allocation and changing graph replays. +The model-owned index-group builder is carried in Core's existing +`MLAModules.index_group_builder` field, preserving the common MLA wrapper ABI. + Follow-up contributions should be separate and measurable: packed C8/indexer formats, native resolver/refill improvements, Host-writeback optimization, metadata reuse, additional hardware and P/D transport. PR19109 retains these experiments. diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index ca27b31b0943..40cea69841b7 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -179,7 +179,6 @@ def __init__( # both layernorms as extra args, so this flag has no effect here and is # accepted for signature compatibility only. fuse_qkv_rmsnorm: bool = False, - index_group_builder=None, ) -> None: nn.Module.__init__(self) self.hidden_size = hidden_size @@ -226,7 +225,7 @@ def __init__( dcp_q_replicate=self.dcp_q_replicate_enabled, use_sparse=mla_modules.is_sparse, indexer=ascend_indexer, - index_group_builder=index_group_builder, + index_group_builder=mla_modules.index_group_builder, skip_topk=skip_topk, topk_indices_buffer=getattr(mla_modules, "topk_indices_buffer", None), non_causal_multi_token_decode=non_causal_multi_token_decode, diff --git a/vllm_ascend/patch/worker/patch_deepseek_v2.py b/vllm_ascend/patch/worker/patch_deepseek_v2.py index 5cf2bc1f325d..8f67c17f18ce 100644 --- a/vllm_ascend/patch/worker/patch_deepseek_v2.py +++ b/vllm_ascend/patch/worker/patch_deepseek_v2.py @@ -268,6 +268,7 @@ def _deepseek_v2_mla_attention_init( indexer_rotary_emb=self.indexer_rope_emb, is_sparse=self.is_v32, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) self.mla_attn = MultiHeadLatentAttentionWrapper( @@ -284,7 +285,6 @@ def _deepseek_v2_mla_attention_init( quant_config, prefix, skip_topk=_skip_topk and not mtp_layer, - index_group_builder=index_group_builder, non_causal_multi_token_decode=non_causal_multi_token_decode, ) From e2f0e2dd98720ca343a55ba59955e5200daf6d6b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 18:42:34 +0800 Subject: [PATCH 005/120] refactor(hisparse): trim compatibility adapters and consolidate tests Reuse main explicit-dtype weight loading and Core private Host-pool defaults. Keep the necessary Core coordinator, shared-budget and view bridges. Remove synthetic loader/signature checks, consolidate capture and ordering coverage, and exercise real CPU workspace ownership. Sixteen CPU cases pass and all forty declared NPU combinations remain; full device/model qualification and full CI are still pending. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 18 ++-- .../one_card/test_hisparse_attention.py | 68 +++++++++++++ .../one_card/test_hisparse_prefill_staging.py | 5 +- .../one_card/test_hisparse_selection_order.py | 85 ---------------- .../test_hisparse_foundation_cpu.py | 98 +++++++++++-------- tests/ut/ops/test_mla_weight_hook_cpu.py | 54 ---------- .../test_hisparse_capture_mirror_cpu.py | 70 ------------- .../kv_transfer/hisparse/backend.py | 9 +- vllm_ascend/ops/mla.py | 16 +-- .../platform/patch_kv_cache_coordinator.py | 1 + 10 files changed, 147 insertions(+), 277 deletions(-) delete mode 100644 tests/e2e/pull_request/one_card/test_hisparse_selection_order.py delete mode 100644 tests/ut/ops/test_mla_weight_hook_cpu.py delete mode 100644 tests/ut/worker/test_hisparse_capture_mirror_cpu.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 7ad689105f49..f64d79792885 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -80,6 +80,11 @@ CANN workspace allocation and changing graph replays. The model-owned index-group builder is carried in Core's existing `MLAModules.index_group_builder` field, preserving the common MLA wrapper ABI. +HiSparse groups use Core's coordinator and shared-slab memory budget. Ascend's +existing private-buffer budget correction and hybrid coordinator are bypassed +for those groups. V2 retains Core's source/group ownership and skips metadata +construction for Resident/Hot groups that have no Attention consumer. + Follow-up contributions should be separate and measurable: packed C8/indexer formats, native resolver/refill improvements, Host-writeback optimization, metadata reuse, additional hardware and P/D transport. PR19109 retains these experiments. @@ -96,16 +101,18 @@ Portable, source-isolated CPU boundary checks: ```bash python -m pytest -q --noconftest \ - tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py \ - tests/ut/ops/test_mla_weight_hook_cpu.py \ - tests/ut/worker/test_hisparse_capture_mirror_cpu.py + tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py ``` These check current Host offsets, failed-unregister ownership, group scratch isolation, strided indexer aliases, planner alignment, unsupported formats, -Core coordinator forwarding, loader ABI and capture-only mirror suppression. +Core coordinator forwarding and capture-only mirror suppression. They do not exercise an installed NPU package or qualify device execution. +The weight-loading wrapper keeps Ascend main's explicit activation-dtype ABI, +already covered by `tests/ut/ops/test_mla_weight_processing.py`; this foundation +does not add another calling convention or a duplicate test module. + Real-NPU cases are included for resolver/refill/overwrite invalidation, independent FP32 Attention references, selection order/duplicates, native BF16 preprocessing and Resident/Host prefill staging, in eager and changing-graph cases: @@ -115,8 +122,7 @@ pytest --confcutdir=tests/e2e/pull_request/one_card \ tests/e2e/pull_request/one_card/test_hisparse_ops.py \ tests/e2e/pull_request/one_card/test_hisparse_attention.py \ tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py \ - tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py \ - tests/e2e/pull_request/one_card/test_hisparse_selection_order.py + tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py ``` The full modules have not been rerun on this new split. Isolated stream-boundary diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 192e73e8f73e..4123b85eaa92 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -103,3 +103,71 @@ def test_sparse_attention_reference(num_queries, strided, capture): torch.npu.synchronize() if graph is not None: graph.reset() + + +@pytest.mark.parametrize("capture", [False, True]) +def test_valid_order_duplicates_and_interleaved_padding(monkeypatch, capture): + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.ops import hisparse + + torch.manual_seed(61) + device, topk, latent, rope, capacity = "npu:0", 2048, 512, 64, 32 + workspace = ( + torch.empty((capacity, topk, 1, latent), dtype=torch.bfloat16, device=device), + torch.empty((capacity, topk, 1, rope), dtype=torch.bfloat16, device=device), + torch.arange(capacity * topk // 128, dtype=torch.int32, device=device).view(capacity, -1), + torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), + torch.arange(1, capacity + 1, dtype=torch.int32, device=device), + torch.full((capacity,), topk, dtype=torch.int32, device=device), + ) + impl = SimpleNamespace( + topk=topk, kv_lora_rank=latent, qk_rope_head_dim=rope, scale=0.0625, _workspace=lambda _: workspace + ) + query = torch.randn(3, 8, latent, dtype=torch.bfloat16, device=device) + query_rope = torch.randn(3, 8, rope, dtype=torch.bfloat16, device=device) + cache = torch.randn(1, 128, latent + rope, dtype=torch.bfloat16, device=device) + indices = torch.full((3, topk), -1, dtype=torch.int32) + indices[0, :6] = torch.tensor([7, -1, 2, 7, -1, 0]) + indices[2, :6] = torch.tensor([0, 5, -1, 3, 5, 1]) + device_indices = indices.to(device) + observed = torch.empty_like(device_indices) + original_pack = hisparse.pack_attention_rows + + def observe(cache, physical, selected_latent, selected_rope, latent_dim): + observed.copy_(physical) + return original_pack(cache, physical, selected_latent, selected_rope, latent_dim) + + monkeypatch.setattr(hisparse, "pack_attention_rows", observe) + + def forward(): + return AscendHiSparseImpl._attention_physical(impl, query, query_rope, cache, device_indices) + + graph = None + try: + if capture: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + result = forward() + indices[0, :6] = torch.tensor([-1, 9, 1, 9, 4, -1]) + device_indices.copy_(indices) + for _ in range(3): + graph.replay() + else: + result = forward() + expected = torch.full_like(indices, -1) + for row in range(3): + valid = indices[row][indices[row] >= 0] + expected[row, : valid.numel()] = valid + torch.testing.assert_close(observed.cpu(), expected, rtol=0, atol=0) + assert torch.isfinite(result.cpu()).all() + assert torch.count_nonzero(result[1].cpu()) == 0 + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 8a69b1ad1daa..079eda15eafd 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -88,12 +88,11 @@ def test_resident_prefill_plan_shared_layers(resident_block_size, host_pages): host._hisparse_registration.close() -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_backend_resident_payload_graph_replay(dtype): +def test_backend_resident_payload_graph_replay(): if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("HiSparse resident gather currently targets Ascend 950") torch.npu.set_device(0) - width = 656 if dtype != torch.bfloat16 else 576 + dtype, width = torch.bfloat16, 576 stride = 3 * width raw = torch.arange(7 * width * dtype.itemsize, dtype=torch.int64).to(torch.uint8).to("npu:0") source = raw.view(dtype).as_strided((2, 2, width), (stride, width, 1), storage_offset=width) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py b/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py deleted file mode 100644 index ea224efab1ba..000000000000 --- a/tests/e2e/pull_request/one_card/test_hisparse_selection_order.py +++ /dev/null @@ -1,85 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The SFA bridge must preserve valid selection order across graph replay.""" - -from types import SimpleNamespace - -import pytest -import torch -import torch_npu # noqa: F401 - - -@pytest.mark.parametrize("capture", [False, True]) -def test_valid_order_duplicates_and_interleaved_padding(monkeypatch, capture): - if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): - pytest.skip("HiSparse SFA integration currently targets Ascend 950") - torch.npu.set_device(0) - from vllm_ascend.utils import bootstrap_custom_op_env - - bootstrap_custom_op_env() - import vllm_ascend.vllm_ascend_C # noqa: F401 - - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - from vllm_ascend.ops import hisparse - - torch.manual_seed(61) - device, topk, latent, rope, capacity = "npu:0", 2048, 512, 64, 32 - workspace = ( - torch.empty((capacity, topk, 1, latent), dtype=torch.bfloat16, device=device), - torch.empty((capacity, topk, 1, rope), dtype=torch.bfloat16, device=device), - torch.arange(capacity * topk // 128, dtype=torch.int32, device=device).view(capacity, -1), - torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), - torch.arange(1, capacity + 1, dtype=torch.int32, device=device), - torch.full((capacity,), topk, dtype=torch.int32, device=device), - ) - impl = SimpleNamespace( - topk=topk, kv_lora_rank=latent, qk_rope_head_dim=rope, scale=0.0625, _workspace=lambda _: workspace - ) - query = torch.randn(3, 8, latent, dtype=torch.bfloat16, device=device) - query_rope = torch.randn(3, 8, rope, dtype=torch.bfloat16, device=device) - cache = torch.randn(1, 128, latent + rope, dtype=torch.bfloat16, device=device) - indices = torch.full((3, topk), -1, dtype=torch.int32) - indices[0, :6] = torch.tensor([7, -1, 2, 7, -1, 0]) - indices[2, :6] = torch.tensor([0, 5, -1, 3, 5, 1]) - device_indices = indices.to(device) - observed = torch.empty_like(device_indices) - original_pack = hisparse.pack_attention_rows - - def observe(cache, physical, selected_latent, selected_rope, latent_dim): - observed.copy_(physical) - return original_pack(cache, physical, selected_latent, selected_rope, latent_dim) - - monkeypatch.setattr(hisparse, "pack_attention_rows", observe) - - def forward(): - return AscendHiSparseImpl._attention_physical(impl, query, query_rope, cache, device_indices) - - graph = None - try: - if capture: - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): - forward() - forward() - torch.npu.synchronize() - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph, stream=stream): - result = forward() - indices[0, :6] = torch.tensor([-1, 9, 1, 9, 4, -1]) - device_indices.copy_(indices) - for _ in range(3): - graph.replay() - else: - result = forward() - expected = torch.full_like(indices, -1) - for row in range(3): - valid = indices[row][indices[row] >= 0] - expected[row, : valid.numel()] = valid - torch.testing.assert_close(observed.cpu(), expected, rtol=0, atol=0) - assert torch.isfinite(result.cpu()).all() - assert torch.count_nonzero(result[1].cpu()) == 0 - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index ad6949406559..476ce6ac4314 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -4,7 +4,6 @@ import ast import ctypes -import inspect import math import sys from dataclasses import dataclass, replace @@ -35,24 +34,6 @@ def backend_type(): )["AscendHiSparseBackend"] -def test_merged_core_resolver_call_keeps_optional_outputs_in_their_current_positions(): - operation = definitions("ops/hisparse.py", {"HiSparseNPUOps"})["HiSparseNPUOps"].hisparse_resolve_residency - # Argument positions from the reviewed merged Core Runtime.swap_in call. - # The former pre-merge method rejected this call or shifted output buffers. - arguments = [object() for _ in range(26)] - bound = inspect.signature(operation).bind(object(), *arguments).arguments - for name, position in { - "miss_mask": 10, - "stats": 11, - "source_block_size": 16, - "resolved_global_indices": 17, - "valid_counts": 18, - "resident_null_block": 24, - "num_valid_rows": 25, - }.items(): - assert bound[name] is arguments[position] - - def test_failed_unregister_retains_host_owner_for_safe_retry(monkeypatch): tensor = torch.empty(64, dtype=torch.uint8) events = [] @@ -102,18 +83,8 @@ def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( monkeypatch, max_rows, hot_rows, union_rows ): - allocated = [] - - class Workspace: - def __init__(self, **geometry): - self.geometry = geometry - allocated.append(self) - - class Ops: - def __init__(self, mapping, workspace): - self.mapping, self.workspace = mapping, workspace - - monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(ResolverWorkspace=Workspace, HiSparseNPUOps=Ops)) + operations = definitions("ops/hisparse.py", {"ResolverWorkspace", "HiSparseNPUOps"}) + monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(**operations)) backend = backend_type()() def group(): @@ -134,12 +105,16 @@ def group(): first, follower = group() second, _ = group() assert follower.workspace is first.cache_ops.workspace - assert follower.mapping is first.registered_host_pool._hisparse_registration + assert follower.host_mapping is first.registered_host_pool._hisparse_registration assert second.cache_ops.workspace is not first.cache_ops.workspace - assert len(allocated) == 2 - assert allocated[0].geometry == dict( - max_rows=max_rows, hot_size=hot_rows, max_union_rows=union_rows, device=torch.device("cpu") - ) + first_workspace, second_workspace = first.cache_ops.workspace, second.cache_ops.workspace + assert first_workspace.row_keys.shape == (max_rows * hot_rows,) + assert first_workspace.cache_keys.shape == (max_rows, hot_rows) + assert first_workspace.max_union_rows == union_rows + first_workspace.row_keys.fill_(13) + second_workspace.row_keys.zero_() + assert torch.all(follower.workspace.row_keys == 13) + assert torch.count_nonzero(second_workspace.row_keys) == 0 @dataclass(frozen=True) @@ -211,8 +186,8 @@ def test_unsupported_indexer_formats_are_rejected_before_pool_binding(field): functions["get_hisparse_indexer_specs"](config) -@pytest.mark.parametrize("hash_lookahead", [0, 7]) -def test_hisparse_roles_preserve_core_admission_and_explicit_hash_budget(hash_lookahead): +@pytest.mark.parametrize("enable_caching", [False, True]) +def test_hisparse_roles_preserve_core_admission_and_explicit_hash_budget(enable_caching): calls = [] marker = object() functions = definitions( @@ -227,12 +202,51 @@ def test_hisparse_roles_preserve_core_admission_and_explicit_hash_budget(hash_lo config, max_model_len=65536, max_in_flight_tokens=129, - num_hash_lookahead_tokens=hash_lookahead, - enable_caching=False, + num_hash_lookahead_tokens=7, + enable_caching=enable_caching, ) assert result is marker and len(calls) == 1 assert calls[0]["kv_cache_config"] is config assert calls[0]["max_in_flight_tokens"] == 129 - assert calls[0]["num_hash_lookahead_tokens"] == hash_lookahead - assert calls[0]["enable_caching"] is False + assert calls[0]["num_hash_lookahead_tokens"] == 7 + assert calls[0]["enable_caching"] is enable_caching assert "eagle_attn_layer_names" not in calls[0] + + +def test_capture_metadata_prevents_warmup_host_mirrors(): + class Parent: + def build_for_cudagraph_capture(self, common, **kwargs): + assert common is common_input and kwargs["draft_index"] is draft + return self.build() + + def build(self): + return NS(hisparse_for_capture=False) + + methods = definitions( + "attention/sfa_hisparse.py", + {"AscendHiSparseMetadataBuilder", "finish_hisparse_kv_update"}, + AscendSFAMetadataBuilder=Parent, + ) + cls = methods["AscendHiSparseMetadataBuilder"] + builder = cls.__new__(cls) + common_input, draft = object(), object() + capture_metadata = builder.build_for_cudagraph_capture(common_input, draft_index=draft) + real_metadata = builder.build() + finish = methods["finish_hisparse_kv_update"] + submitted = [] + cache = NS(finish_kv_update=lambda: submitted.append(True)) + for _ in range(5): + finish(cache, capture_metadata) + assert submitted == [] + finish(cache, real_metadata) + assert submitted == [True] + + +def test_real_forward_mirror_errors_remain_visible(): + finish = definitions("attention/sfa_hisparse.py", {"finish_hisparse_kv_update"})["finish_hisparse_kv_update"] + + def duplicate(): + raise RuntimeError("HiSparse layer 0 mirrored twice.") + + with pytest.raises(RuntimeError, match="mirrored twice"): + finish(NS(finish_kv_update=duplicate), NS(hisparse_for_capture=False)) diff --git a/tests/ut/ops/test_mla_weight_hook_cpu.py b/tests/ut/ops/test_mla_weight_hook_cpu.py deleted file mode 100644 index 5c1f9a450d53..000000000000 --- a/tests/ut/ops/test_mla_weight_hook_cpu.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Exercise the real weight-hook wrapper with both loader call signatures.""" - -import ast -import sys -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch - - -@pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize("explicit_dtype", [False, True]) -@pytest.mark.parametrize("replicated", [False, True]) -def test_post_load_hook_preserves_activation_dtype(monkeypatch, sparse, explicit_dtype, replicated): - source = Path(__file__).resolve().parents[3] / "vllm_ascend/ops/mla.py" - tree = ast.parse(source.read_bytes()) - wrapper = next( - node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "wrapped_process_weights" - ) - calls = [] - - class SFA: - def process_weights_after_loading(self, dtype): - calls.append(("SFA", dtype)) - - class Dense: - def process_weights_after_loading(self, dtype): - calls.append(("replicated", dtype)) - - def old_hook(dtype): - calls.append(("original", dtype)) - - def current_hook(): - calls.append(("original", torch.bfloat16)) - - monkeypatch.setitem(sys.modules, "vllm_ascend.attention.sfa_v1", SimpleNamespace(AscendSFAImpl=SFA)) - namespace = { - "torch": torch, - "self": SimpleNamespace( - mla_attn=SimpleNamespace(impl=SFA() if sparse else Dense(), act_dtype=torch.bfloat16), - dcp_q_replicate_enabled=replicated, - ), - "original_process_weights": old_hook if explicit_dtype else current_hook, - } - exec(compile(ast.Module(body=[wrapper], type_ignores=[]), str(source), "exec"), namespace) - if explicit_dtype: - namespace["wrapped_process_weights"](torch.float16) - else: - namespace["wrapped_process_weights"]() - expected = torch.float16 if explicit_dtype else torch.bfloat16 - assert calls == [("SFA" if sparse else "replicated" if replicated else "original", expected)] diff --git a/tests/ut/worker/test_hisparse_capture_mirror_cpu.py b/tests/ut/worker/test_hisparse_capture_mirror_cpu.py deleted file mode 100644 index 3ab15b32f19a..000000000000 --- a/tests/ut/worker/test_hisparse_capture_mirror_cpu.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Capture warmups must not publish a real connector step's Host mirrors.""" - -import ast -from pathlib import Path -from types import SimpleNamespace - -import pytest - -ROOT = Path(__file__).resolve().parents[3] -BASE = ROOT / "vllm_ascend/attention/sfa_hisparse.py" - - -def load_helper(): - tree = ast.parse(BASE.read_text()) - node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "finish_hisparse_kv_update") - namespace = {} - exec(compile(ast.Module(body=[node], type_ignores=[]), str(BASE), "exec"), namespace) - return namespace[node.name] - - -@pytest.mark.parametrize("filename", ["sfa_hisparse.py"]) -def test_capture_builder_marks_only_its_fresh_metadata_and_forwards_arguments(filename): - path = BASE.with_name(filename) - tree = ast.parse(path.read_text()) - cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "AscendHiSparseMetadataBuilder") - method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "build_for_cudagraph_capture") - calls = [] - - class Parent: - def build_for_cudagraph_capture(self, common, **kwargs): - calls.append((common, kwargs)) - return self.build() - - def build(self): - return SimpleNamespace(hisparse_for_capture=False) - - cls = ast.ClassDef( - name=cls.name, bases=[ast.Name(id="Parent", ctx=ast.Load())], keywords=[], body=[method], decorator_list=[] - ) - namespace = {"Parent": Parent} - exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), str(path), "exec"), namespace) - builder = namespace[cls.name]() - common, draft = object(), object() - capture = builder.build_for_cudagraph_capture(common, draft_index=draft) - assert calls == [(common, {"draft_index": draft})] - assert capture.hisparse_for_capture is True - assert builder.build().hisparse_for_capture is False - - -def test_repeated_capture_forwards_do_not_submit_host_mirrors(): - finish = load_helper() - calls = [] - cache = SimpleNamespace(finish_kv_update=lambda: calls.append(True)) - for _ in range(5): - finish(cache, SimpleNamespace(hisparse_for_capture=True)) - assert calls == [] - finish(cache, SimpleNamespace(hisparse_for_capture=False)) - finish(cache, SimpleNamespace()) - assert calls == [True, True] - - -def test_real_forward_duplicate_errors_remain_visible(): - finish = load_helper() - - def duplicate(): - raise RuntimeError("HiSparse layer 0 mirrored twice.") - - with pytest.raises(RuntimeError, match="mirrored twice"): - finish(SimpleNamespace(finish_kv_update=duplicate), SimpleNamespace(hisparse_for_capture=False)) diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py index 5a4950b2ee06..8882c2b8e067 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py +++ b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py @@ -16,15 +16,14 @@ class AscendHiSparseBackend(HiSparseBackend): - # NPU event IPC is not implemented here. Keep private pools per TP rank; - # the common connector still owns transfer completion and TP coordination. - supports_shared_host_pool = False + # Inherit Core's private Host-pool policy; NPU event IPC is not implemented. stats_dtype = torch.int64 def check_available(self) -> None: - from vllm_ascend.ops.hisparse import HiSparseNPUOps from vllm_ascend.utils import bootstrap_custom_op_env + # Import the device kernels before the connector allocates Host storage. + importlib.import_module("vllm_ascend.ops.hisparse") if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): # A5 disables the global custom-op selection policy. HiSparse still # needs its explicit DMA dependency; load the extension lazily here @@ -36,8 +35,6 @@ def check_available(self) -> None: raise RuntimeError("Ascend HiSparse requires the compiled Ascend DMA extension.") from error if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") - if not callable(HiSparseNPUOps): - raise RuntimeError("Ascend HiSparse device kernels are unavailable.") def pin_tensor(self, tensor: torch.Tensor) -> None: from vllm_ascend.ops.hisparse import MappedHostTensor diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index 40cea69841b7..402aa56ad3a3 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -250,23 +250,17 @@ def __init__( original_process_weights = self.mla_attn.process_weights_after_loading - def wrapped_process_weights(act_dtype: torch.dtype | None = None): + def wrapped_process_weights(act_dtype: torch.dtype): from vllm_ascend.attention.sfa_v1 import AscendSFAImpl if self.dcp_q_replicate_enabled: - self.mla_attn.impl.process_weights_after_loading( - self.mla_attn.act_dtype if act_dtype is None else act_dtype - ) + self.mla_attn.impl.process_weights_after_loading(act_dtype) elif not isinstance(self.mla_attn.impl, AscendSFAImpl): - if act_dtype is None: - original_process_weights() - else: - original_process_weights(act_dtype) + # Both supported vLLM versions dispatch to the impl here. + original_process_weights(act_dtype) else: # SFA disposes kv_b_proj, so bypass upstream's dense packing. - self.mla_attn.impl.process_weights_after_loading( - self.mla_attn.act_dtype if act_dtype is None else act_dtype - ) + self.mla_attn.impl.process_weights_after_loading(act_dtype) self.mla_attn.process_weights_after_loading = wrapped_process_weights diff --git a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py index 7dcfa5a78ec7..510a7fa85e03 100644 --- a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py +++ b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py @@ -558,6 +558,7 @@ def get_kv_cache_coordinator( # type: ignore[misc] hybrid_kwargs["allow_partial_hash_hits"] = allow_partial_hash_hits if any(group.role == KVCacheGroupRole.HISPARSE_SOURCE for group in kv_cache_config.kv_cache_groups): # Core owns HiSparse admission, HMA roles, eviction and prefix state. + # Remove this bypass when the Ascend coordinator factory patch retires. return _orig_get_kv_cache_coordinator( **{name: value for name, value in hybrid_kwargs.items() if name != "eagle_attn_layer_names"} ) From 153497bb424ca7d6d73d4823170b77f5de1045cf Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 19:38:47 +0800 Subject: [PATCH 006/120] feat(hisparse): restore native C8 foundation Reuse native SFA quantization and Attention, byte-capable scatter, and Core HMA/lifecycle ownership. Bind packed main rows and page-strided Indexer keys/scales without optional optimization profiles. Add payload, addressing, alignment and eager/Graph regression coverage. Source-isolated CPU checks pass (23 cases); new-head NPU and full-model qualification remain pending. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 42 ++++-- .../one_card/test_hisparse_attention.py | 77 +++++++++- .../one_card/test_hisparse_kv_preprocess.py | 72 +++++++++ .../one_card/test_hisparse_ops.py | 130 ++++++++++++++++ .../one_card/test_hisparse_prefill_staging.py | 4 +- .../test_hisparse_foundation_cpu.py | 139 +++++++++++++++++- vllm_ascend/attention/indexer.py | 14 ++ vllm_ascend/attention/sfa_hisparse.py | 72 ++++++++- .../kv_transfer/hisparse/backend.py | 10 ++ .../kv_transfer/hisparse/binding.py | 126 ++++++++++++++-- vllm_ascend/ops/hisparse_c8.py | 139 ++++++++++++++++++ vllm_ascend/worker/v2/attn_utils.py | 16 +- 12 files changed, 807 insertions(+), 34 deletions(-) create mode 100644 vllm_ascend/ops/hisparse_c8.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index f64d79792885..149db68d3d17 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -21,8 +21,8 @@ Core sparse index group -> logical TopK -> NPU residency/refill plan packed KV -> existing native SFA operation ``` -This first data path covers unquantized main/indexer KV with the RoPE SFA -layout and Model Runner V2. Model-weight quantization is independent. Private +The data path covers unquantized KV and the native Ascend 950 C8 format with +the RoPE SFA layout and Model Runner V2. Model-weight quantization is independent. Private Host pools are used per rank because NPU IPC events are not implemented. The mapped-Host capability is checked at registration; hardware evidence for the original primitives is from Ascend 950. A3 and the new source pair still @@ -55,6 +55,27 @@ The example reserves 4 GiB Host storage per rank. Size it for the model and cont budget; it is not a model capacity claim. Keep Core's eager mirror policy for this foundation. External P/D transfer is not provided by this PR. +For native A5 C8, replace the two dtype options above with: + +```bash +--kv-cache-dtype fp8 \ +--attention-config '{"indexer_kv_dtype":"fp8","hisparse_config":{"eager_host_mirror":true}}' +``` + +The existing SFA quantizer produces each main KV row: 512 E4M3FN NoPE bytes, +64 BF16 RoPE values and four FP32 block scales, totaling 656 bytes. The adapter +copies these components as opaque bytes and calls the existing quantized SFA +Attention operation. Indexer C8 uses 128 E4M3FN keys and one FP32 scale per token. +Main and Indexer dtypes remain independently selectable. C4, TurboQuant, +INT8-format C8 and other payload geometries are not implemented by this bridge. + +Core's planner receives the complete key/scale page size and packed-row stride +alignment before choosing the shared-slab capacity. Strided views preserve the +original allocation and layer offset. The existing scatter handles packed bytes; +prefill Resident gathers reinterpret FP8 as bytes before using Core's gather. +HiSparse writes through its own current destination once; the ordinary SFA C8 +source-slot store is skipped. No new fused scatter or Attention kernel is added. + ## Extension points for contributors Use existing interfaces, rather than adding unused placeholder hooks: @@ -63,9 +84,9 @@ Use existing interfaces, rather than adding unused placeholder hooks: | --- | --- | --- | | `NPUPlatform.get_hisparse_backend_cls` | Core selects the device backend only for HiSparse | Backend capability/version support | | `HiSparseBackend.create_cache_ops` | Runtime invokes resolver, refill and invalidation operations | Faster resolver/refill, vectorization and balancing | -| `validate_kv_cache_config` / `prepare_kv_cache_views` | Core validates before allocation and binds views inside rollback scope | Packed/quantized KV formats with layout tests | +| `validate_kv_cache_config` / `prepare_kv_cache_views` | Core validates BF16/C8 before allocation and binds views inside rollback scope | Additional payload formats with layout tests | | Backend `pin_tensor` / `unpin_tensor` / `swap_blocks_batch` | Core owns completion, explicit registration release and transfer lifetime | DMA implementation and supported memory capabilities | -| `AscendHiSparseMetadataBuilder` / `AscendHiSparseImpl` | SFA factories select metadata and the layout bridge | C8 and a smaller layout-conversion cost | +| `AscendHiSparseMetadataBuilder` / `AscendHiSparseImpl` | SFA factories select metadata and the BF16/C8 layout bridge | A smaller layout-conversion cost | `ops/hisparse.py` owns operation validation and persistent per-group workspace; `ops/triton/hisparse.py` implements the reference NPU path. Followers share their @@ -85,8 +106,8 @@ existing private-buffer budget correction and hybrid coordinator are bypassed for those groups. V2 retains Core's source/group ownership and skips metadata construction for Resident/Hot groups that have no Attention consumer. -Follow-up contributions should be separate and measurable: packed C8/indexer -formats, native resolver/refill improvements, Host-writeback optimization, metadata +Follow-up contributions should be separate and measurable: other packed formats, +native resolver/refill improvements, Host-writeback optimization, metadata reuse, additional hardware and P/D transport. PR19109 retains these experiments. AVSG, lazy-MTP policy, native optimized kernels, private Core adapters and machine-specific benchmarking tools are excluded here. @@ -106,7 +127,9 @@ python -m pytest -q --noconftest \ These check current Host offsets, failed-unregister ownership, group scratch isolation, strided indexer aliases, planner alignment, unsupported formats, -Core coordinator forwarding and capture-only mirror suppression. +Core coordinator forwarding and capture-only mirror suppression. C8 checks also +cover opaque component bytes, nonzero slab offsets, stable duplicate/padding +ordering, Indexer key/scale writes and planner alignment before allocation. They do not exercise an installed NPU package or qualify device execution. The weight-loading wrapper keeps Ascend main's explicit activation-dtype ABI, @@ -114,8 +137,9 @@ already covered by `tests/ut/ops/test_mla_weight_processing.py`; this foundation does not add another calling convention or a duplicate test module. Real-NPU cases are included for resolver/refill/overwrite invalidation, independent -FP32 Attention references, selection order/duplicates, native BF16 preprocessing -and Resident/Host prefill staging, in eager and changing-graph cases: +FP32 Attention references, selection order/duplicates, native BF16/C8 preprocessing, +C8 key/scale writes, packed-byte scatter and Resident/Host prefill staging, in +eager and changing-graph cases: ```bash pytest --confcutdir=tests/e2e/pull_request/one_card \ diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 4123b85eaa92..ca1896280a4a 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -43,7 +43,7 @@ def test_sparse_attention_reference(num_queries, strided, capture): torch.arange(1, capacity + 1, dtype=torch.int32, device=device), torch.full((capacity,), topk, dtype=torch.int32, device=device), ) - impl = SimpleNamespace(kv_lora_rank=latent_dim, qk_rope_head_dim=rope_dim, scale=scale) + impl = SimpleNamespace(kv_lora_rank=latent_dim, qk_rope_head_dim=rope_dim, scale=scale, enable_sparse_sfa_c8=False) impl._workspace = lambda dtype: workspace # Quantize before taking the FP32 reference so both paths see the same inputs. source = torch.randn((blocks, block_size, latent_dim + rope_dim)).to(torch.bfloat16) @@ -121,7 +121,12 @@ def test_valid_order_duplicates_and_interleaved_padding(monkeypatch, capture): torch.full((capacity,), topk, dtype=torch.int32, device=device), ) impl = SimpleNamespace( - topk=topk, kv_lora_rank=latent, qk_rope_head_dim=rope, scale=0.0625, _workspace=lambda _: workspace + topk=topk, + kv_lora_rank=latent, + qk_rope_head_dim=rope, + scale=0.0625, + _workspace=lambda _: workspace, + enable_sparse_sfa_c8=False, ) query = torch.randn(3, 8, latent, dtype=torch.bfloat16, device=device) query_rope = torch.randn(3, 8, rope, dtype=torch.bfloat16, device=device) @@ -171,3 +176,71 @@ def forward(): torch.npu.synchronize() if graph is not None: graph.reset() + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_c8_attention_uses_existing_native_operator(strided, capture): + from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope + from vllm_ascend.ops.hisparse_c8 import c8_physical_attention, pack_c8_components + + torch.manual_seed(71) + rows, heads, blocks, block, topk = 3, 8, 2, 128, 128 + device = "npu:0" + values = torch.randn(blocks * block, 1, 1, 576, dtype=torch.bfloat16, device=device) + gamma = torch.ones(512, dtype=torch.bfloat16, device=device) + cos = torch.ones(blocks * block, 1, 1, 64, dtype=torch.bfloat16, device=device) + sin = torch.zeros_like(cos) + rope, nope, scales = custom_kv_rmsnorm_rope( + values, gamma, cos, sin, 512, 64, epsilon=1e-5, dst_type=torch.float8_e4m3fn, tile_size=128 + ) + packed = pack_c8_components(nope.contiguous(), rope.contiguous(), scales.contiguous()) + backing = torch.empty((blocks, 2 if strided else 1, block, 656), dtype=torch.float8_e4m3fn, device=device) + cache = backing[:, -1] + cache.view(torch.uint8).copy_(packed.reshape(blocks, block, 656)) + q = torch.randn(rows, heads, 512, dtype=torch.bfloat16, device=device) * 0.25 + qr = torch.randn(rows, heads, 64, dtype=torch.bfloat16, device=device) * 0.25 + indices = torch.randint(0, blocks * block, (rows, topk), dtype=torch.int32) + indices[:, 3::5] = -1 + indices[0] = -1 + indices[0, :3] = torch.tensor([3, 3, 7], dtype=torch.int32) + indices[-1] = -1 + selected = indices.to(device) + query_end, key_length = (torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) + + def forward(): + return c8_physical_attention(q, qr, cache, selected, 1 / math.sqrt(576), query_end, key_length) + + graph = None + try: + if capture: + forward() + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + output = forward() + indices[0, :4] = torch.tensor([9, 2, 9, 0], dtype=torch.int32) + selected.copy_(indices) + q.neg_() + graph.replay() + else: + output = forward() + torch.npu.synchronize() + scale = scales.contiguous().view(torch.float32).cpu().reshape(-1, 4).repeat_interleave(128, dim=-1) + latent = nope.float().cpu().reshape(-1, 512) * scale + rope_cpu = rope.contiguous().view(torch.bfloat16).cpu().reshape(-1, 64).float() + keys = latent[indices.clamp_min(0).long()] + rope_keys = rope_cpu[indices.clamp_min(0).long()] + logits = torch.einsum("thd,tkd->thk", q.float().cpu(), keys) + logits += torch.einsum("thd,tkd->thk", qr.float().cpu(), rope_keys) + logits *= 1 / math.sqrt(576) + logits.masked_fill_(~(indices >= 0)[:, None, :], -torch.inf) + probability = logits.softmax(-1).nan_to_num(0) + expected = torch.einsum("thk,tkd->thd", probability, keys) + torch.testing.assert_close(output.float().cpu(), expected, atol=0.01, rtol=0.05) + assert torch.count_nonzero(output[-1].cpu()) == 0 + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index 6433402b71b2..b382532deed3 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -39,6 +39,76 @@ def _setup(): torch.npu.set_device(0) +@pytest.mark.parametrize("capture", [False, True]) +def test_c8_preprocess_writes_hisparse_mapping_once(capture): + _setup() + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope + + rows, block = 5, 128 + values = torch.randn(rows, 576, dtype=torch.bfloat16, device="npu") + gamma = torch.ones(512, dtype=torch.bfloat16, device="npu") + cos = torch.ones(rows, 1, 1, 64, dtype=torch.bfloat16, device="npu") + sin = torch.zeros_like(cos) + raw = torch.full((3, 2, block, 656), 90, dtype=torch.uint8, device="npu") + destination = raw[:, 1].view(torch.float8_e4m3fn) + logical = [-1, 129, 0, 257, 128] + slots = torch.tensor(logical, dtype=torch.int64, device="npu") + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="npu"), hisparse_for_capture=True + ) + impl = SimpleNamespace( + enable_sparse_sfa_c8=True, + c8_cache_dtype=torch.float8_e4m3fn, + sfa_qsfa_tile_size=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + num_kv_heads=1, + kv_a_layernorm=SimpleNamespace(weight=gamma, variance_epsilon=1e-5), + hisparse_cache=SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)), + ) + # The ordinary source cache has no Resident/Hot addressing. It must never + # receive the parent's second C8 store after exec_kv has used write_target. + source_cache = (torch.empty(0, device="npu"),) + + def forward(): + rope, latent, scales = AscendHiSparseImpl.exec_kv(impl, values, cos, sin, source_cache, slots, metadata) + AscendHiSparseImpl._store_parallel_kv(impl, rope, latent, scales, source_cache, slots, metadata) + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + forward() + for count in (4, 2): + values.neg_() + metadata.query_start_loc[-1:].fill_(count) + raw.fill_(90) + if graph is None: + forward() + else: + graph.replay() + torch.npu.synchronize() + rope, latent, scales = custom_kv_rmsnorm_rope( + values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 + ) + payload = torch.cat( + [part.contiguous().view(torch.uint8).reshape(rows, -1).cpu() for part in (latent, rope, scales)], -1 + ) + expected = torch.full(raw.shape, 90, dtype=torch.uint8) + for row, slot in enumerate(logical): + if row < count and slot >= 0: + expected[slot // block, 1, slot % block] = payload[row] + torch.testing.assert_close(raw.cpu(), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + + @pytest.mark.parametrize("capture", [False, True]) def test_fused_preprocess_preserves_native_rounding_and_packed_writes(monkeypatch, capture): _setup() @@ -100,6 +170,7 @@ def persistent(key, allocate): finish_kv_update=lambda: mirrored.append(True), ) impl = SimpleNamespace( + enable_sparse_sfa_c8=False, kv_lora_rank=512, qk_rope_head_dim=64, num_kv_heads=1, @@ -181,6 +252,7 @@ def unexpected(*args, **kwargs): monkeypatch.setattr(torch_npu, "npu_kv_rmsnorm_rope_cache", unexpected) impl = SimpleNamespace( + enable_sparse_sfa_c8=False, kv_lora_rank=512, qk_rope_head_dim=0, kv_a_layernorm=norm, diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 483887a694a8..c60c9c102661 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -95,6 +95,49 @@ def forward(): graph.reset() +@pytest.mark.parametrize("capture", [False, True]) +def test_c8_bytes_reuse_existing_scatter_with_padding_and_page_gaps(capture): + rows, pages, block, width = 5, 4, 128, 656 + raw = torch.full((pages, 2, block, width), 165, dtype=torch.uint8, device="npu") + cache = raw[:, 1] + values = torch.arange(rows * width, dtype=torch.int64).to(torch.uint8).reshape(rows, width).npu() + slots_cpu = torch.tensor([257, -1, 0, 128, 511], dtype=torch.int64) + slots = slots_cpu.npu() + valid = torch.tensor([4], dtype=torch.int32, device="npu") + + def forward(): + scatter_packed_kv(values, values[:, :0], cache, slots, valid) + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + forward() + for count in (4, 2): + raw.fill_(165) + values.bitwise_xor_(1) + valid.fill_(count) + if graph is None: + forward() + else: + graph.replay() + torch.npu.synchronize() + expected = torch.full_like(raw).cpu() + expected.fill_(165) + payload = values.cpu() + for row, slot in enumerate(slots_cpu.tolist()): + if row < count and slot >= 0: + expected[slot // block, 1, slot % block] = payload[row] + torch.testing.assert_close(raw.cpu(), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + + def reference(topk, requests, states, source, resident, hot_table, dgi, lru, block): dgi = dgi.clone() lru = lru.clone() @@ -694,3 +737,90 @@ def forward(): finally: torch.npu.synchronize() graph.reset() + + +@pytest.mark.parametrize("capture", [False, True]) +def test_page_strided_indexer_write_preserves_bits_and_padding(capture): + from types import SimpleNamespace + + from vllm_ascend.attention.indexer import AscendSFAIndexerBackend + from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec + from vllm_ascend.distributed.kv_transfer.hisparse.binding import _quantized_indexer_views + + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("The native E4M3FN HiSparse format currently targets A5") + torch.npu.set_device(0) + pages, block, width, page_stride = 3, 128, 128, 419840 + size = (pages - 1) * page_stride + block * width + block * 4 + raw = torch.full((size,), 90, dtype=torch.uint8, device="npu") + packed = torch.as_strided( + raw.view(torch.float8_e4m3fn), (pages, 1, block, width), (page_stride, block * width, width, 1) + ) + spec = AscendSFAIndexerCacheSpec( + block_size=block, + num_kv_heads=1, + head_size=width, + dtype=torch.float8_e4m3fn, + cache_dtype_str="fp8", + scale_dim=1, + scale_dtype=torch.float32, + cache_sparse_li_c8=True, + cache_sparse_li_c4=False, + sfa_dcp_replicated_indexer_size=1, + ) + key, scale = _quantized_indexer_views(packed, spec) + with pytest.raises(RuntimeError, match="view size"): + key.view(-1, width) + writer = AscendSFAIndexerBackend.__new__(AscendSFAIndexerBackend) + torch.nn.Module.__init__(writer) + writer.k_cache = SimpleNamespace( + kv_cache=(key, scale), hisparse_indexer_payload_format="ascend_indexer_e4m3fn_key_fp32_scale_v1" + ) + writer.enable_sparse_li_c8 = True + writer.enable_sparse_li_c4 = False + values = torch.zeros((5, width), dtype=torch.uint8, device="npu").view(torch.float8_e4m3fn) + scales = torch.zeros(5, dtype=torch.float32, device="npu") + slots = torch.full((5,), -1, dtype=torch.int32, device="npu") + + def write(): + writer.write_cache(values, scales, slots) + + write() + torch.npu.synchronize() + graph = None + if capture: + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + write() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, stream=stream): + write() + try: + for iteration, logical in enumerate([[-1, 0, 127, 128, 383], [0, 129, -1, 382, 255], [128, 256, 1, -1, 127]]): + payload = (torch.arange(5 * width, dtype=torch.int64) + iteration * 17).to(torch.uint8).view(5, width) + scale_bits = torch.tensor( + [0x3F800000, 0x80000000, 0x7FC00123, 0x7F800000, iteration], dtype=torch.int64 + ).to(torch.int32) + values.view(torch.uint8).copy_(payload) + scales.view(torch.int32).copy_(scale_bits) + slots.copy_(torch.tensor(logical, dtype=torch.int32)) + raw.fill_(90) + torch.npu.synchronize() + graph.replay() if graph is not None else write() + torch.npu.synchronize() + expected = torch.full((size,), 90, dtype=torch.uint8) + scale_bytes = scale_bits.view(torch.uint8).view(5, 4) + for row, slot in enumerate(logical): + if slot < 0: + continue + page, offset = divmod(slot, block) + start = page * page_stride + offset * width + expected[start : start + width] = payload[row] + start = page * page_stride + block * width + offset * 4 + expected[start : start + 4] = scale_bytes[row] + torch.testing.assert_close(raw.cpu(), expected, rtol=0, atol=0) + finally: + if graph is not None: + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 079eda15eafd..16268196e9ac 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -88,11 +88,11 @@ def test_resident_prefill_plan_shared_layers(resident_block_size, host_pages): host._hisparse_registration.close() -def test_backend_resident_payload_graph_replay(): +@pytest.mark.parametrize("dtype,width", [(torch.bfloat16, 576), (torch.float8_e4m3fn, 656)]) +def test_backend_resident_payload_graph_replay(dtype, width): if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("HiSparse resident gather currently targets Ascend 950") torch.npu.set_device(0) - dtype, width = torch.bfloat16, 576 stride = 3 * width raw = torch.arange(7 * width * dtype.itemsize, dtype=torch.int64).to(torch.uint8).to("npu:0") source = raw.view(dtype).as_strided((2, 2, width), (stride, width, 1), storage_offset=width) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 476ce6ac4314..112f529a981b 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -4,6 +4,7 @@ import ast import ctypes +import importlib.util import math import sys from dataclasses import dataclass, replace @@ -126,23 +127,39 @@ class IndexerSpec: scale_dim: int = 0 cache_sparse_li_c8: bool = False cache_sparse_li_c4: bool = False + dtype: torch.dtype = torch.float32 + scale_dtype: torch.dtype = torch.float32 + sfa_dcp_replicated_indexer_size: int = 1 + page_size_padded: int | None = None @property def page_size_bytes(self): - return self.block_size * self.num_kv_heads * self.head_size * torch.float32.itemsize + return ( + self.block_size + * self.num_kv_heads + * (self.head_size * self.dtype.itemsize + self.scale_dim * self.scale_dtype.itemsize) + ) def get_block_stride_alignment(self): return self.block_stride_alignment -def binding_functions(): +def binding_functions(**bindings): return definitions( "distributed/kv_transfer/hisparse/binding.py", - {"align_hisparse_indexer_specs", "get_hisparse_indexer_specs", "_bind_indexer_views"}, + { + "align_hisparse_indexer_specs", + "get_hisparse_indexer_specs", + "_bind_indexer_views", + "_validate_indexer_spec", + "_quantized_indexer_views", + "write_hisparse_indexer_c8", + }, AscendSFAIndexerCacheSpec=IndexerSpec, UniformTypeKVCacheSpecs=type("Uniform", (), {}), math=math, replace=replace, + **bindings, ) @@ -169,7 +186,7 @@ def test_expanded_indexer_alias_preserves_shared_slab_offsets(blocks, stride): def test_alignment_updates_only_indexer_specs_without_changing_inputs(): functions = binding_functions() - source, indexer = object(), IndexerSpec() + source, indexer = NS(get_block_stride_alignment=lambda: 1), IndexerSpec() values = {"source": source, "indexer": indexer} result = functions["align_hisparse_indexer_specs"](values) assert result["source"] is source @@ -177,12 +194,34 @@ def test_alignment_updates_only_indexer_specs_without_changing_inputs(): assert values["indexer"] is indexer and indexer.block_stride_alignment == 3 +def test_c8_page_alignment_reserves_scale_bytes_and_whole_main_rows(): + @dataclass(frozen=True) + class StorageSpec: + block_stride_alignment: int = 128 + + def get_block_stride_alignment(self): + return self.block_stride_alignment + + functions = binding_functions(HiSparseResidentSpec=StorageSpec, HiSparseHotSpec=StorageSpec) + storage = StorageSpec() + indexer = IndexerSpec( + block_size=128, head_size=128, dtype=torch.float8_e4m3fn, scale_dim=1, cache_sparse_li_c8=True + ) + result = functions["align_hisparse_indexer_specs"]( + {"resident": storage, "hot": storage, "indexer": indexer}, packed_row_bytes=656 + ) + stride_alignment = math.lcm(*(spec.get_block_stride_alignment() for spec in result.values())) + assert stride_alignment % 656 == 0 and stride_alignment % 128 == 0 + assert result["indexer"].page_size_padded == 128 * (128 + 4) + assert storage.block_stride_alignment == 128 and indexer.page_size_padded is None + + @pytest.mark.parametrize("field", ["scale_dim", "cache_sparse_li_c8", "cache_sparse_li_c4"]) def test_unsupported_indexer_formats_are_rejected_before_pool_binding(field): functions = binding_functions() spec = replace(IndexerSpec(), **{field: 1}) config = NS(kv_cache_groups=[NS(layer_names=["indexer"], kv_cache_spec=spec)]) - with pytest.raises(ValueError, match="unquantized"): + with pytest.raises(ValueError, match="C4|C8|scales"): functions["get_hisparse_indexer_specs"](config) @@ -250,3 +289,93 @@ def duplicate(): with pytest.raises(RuntimeError, match="mirrored twice"): finish(NS(finish_kv_update=duplicate), NS(hisparse_for_capture=False)) + + +def c8_module(): + path = ROOT / "vllm_ascend/ops/hisparse_c8.py" + spec = importlib.util.spec_from_file_location("hisparse_c8_test_module", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_c8_packing_preserves_opaque_rope_scale_and_nan_payloads(): + module = c8_module() + rows = 3 + raw = [torch.arange(rows * width).to(torch.uint8).reshape(rows, width) for width in (512, 128, 16)] + components = [values.view(torch.float8_e4m3fn) for values in raw] + packed = module.pack_c8_components(*components) + assert packed.dtype == torch.uint8 and packed.shape == (rows, 656) + torch.testing.assert_close(packed, torch.cat(raw, dim=-1), rtol=0, atol=0) + + +@pytest.mark.parametrize("stride_rows", [128, 136]) +def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): + module = c8_module() + pages, block, width, offset = 3, 128, 656, 1312 + backing = torch.full((offset + pages * stride_rows * width,), 17, dtype=torch.uint8) + cache = backing.view(torch.float8_e4m3fn).as_strided( + (pages, block, width), (stride_rows * width, width, 1), storage_offset=offset + ) + view, actual_block, actual_stride, capacity = module.c8_storage_view(cache) + assert (actual_block, actual_stride, capacity) == (block, stride_rows, pages * block) + assert view.data_ptr() == cache.data_ptr() + assert view.untyped_storage().data_ptr() == backing.data_ptr() + indices = torch.tensor([[0, 127, 128, 255, 256, 383, -1, 384]], dtype=torch.int32) + mapped = module.map_c8_physical_indices(indices, actual_block, actual_stride, capacity) + expected = torch.tensor( + [[0, 127, stride_rows, stride_rows + 127, 2 * stride_rows, 2 * stride_rows + 127, -1, -1]], + dtype=torch.int32, + ) + torch.testing.assert_close(mapped, expected, rtol=0, atol=0) + view[2 * stride_rows, 0, 1:2].view(torch.uint8).fill_(123) + assert backing[offset + 2 * stride_rows * width + 1] == 123 + + +def test_c8_compaction_preserves_duplicate_order_and_empty_rows(): + module = c8_module() + indices = torch.tensor([[7, -1, 2, 7, -1, 0], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) + expected = torch.tensor([[7, 2, 7, 0, -1, -1], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) + torch.testing.assert_close(module.compact_c8_indices(indices), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("padding", [0, 256]) +def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(padding): + functions = binding_functions() + spec = IndexerSpec(block_size=128, head_size=128, dtype=torch.float8_e4m3fn, scale_dim=1, cache_sparse_li_c8=True) + pages, offset = 3, 256 + key_bytes = 128 * 128 + stride = spec.page_size_bytes + padding + raw = torch.full((offset + pages * stride + 256,), 165, dtype=torch.uint8) + cache = raw.view(torch.float8_e4m3fn).as_strided( + (pages, 1, 128, 128), (stride, key_bytes, 128, 1), storage_offset=offset + ) + layer = NS() + caches = {"indexer": cache} + functions["_bind_indexer_views"](caches, {"indexer": spec}, {"indexer": layer}) + key, scale = caches["indexer"] + assert key.untyped_storage().data_ptr() == scale.untyped_storage().data_ptr() == raw.data_ptr() + assert layer.hisparse_block_stride == 1 + assert layer.hisparse_indexer_payload_format == "ascend_indexer_e4m3fn_key_fp32_scale_v1" + slots = torch.tensor([127, 0, 128, 383, -1, 256], dtype=torch.int32) + values_raw = torch.arange(6 * 128).to(torch.uint8).reshape(6, 128) + values = values_raw.view(torch.float8_e4m3fn) + scales = torch.arange(6, dtype=torch.float32).reshape(6, 1) + 0.5 + + def scatter(destination, indexes, source): + indexes = indexes.flatten() + valid = indexes >= 0 + destination[indexes[valid].long()] = source[valid] + + expected = raw.clone() + scale_bytes = scales.view(torch.uint8).reshape(6, 4) + for i, slot in enumerate(slots.tolist()): + if slot < 0: + continue + page, row = divmod(slot, 128) + start = offset + page * stride + row * 128 + expected[start : start + 128] = values_raw[i] + start = offset + page * stride + key_bytes + row * 4 + expected[start : start + 4] = scale_bytes[i] + functions["write_hisparse_indexer_c8"](key, scale, values, scales, slots, scatter) + torch.testing.assert_close(raw, expected, rtol=0, atol=0) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index d4eafe8b589a..34cf99fd8521 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -286,6 +286,20 @@ def write_cache( with a different cache layout should override it. """ indexer_k_cache = self.k_cache.kv_cache[INDEXER_K_CACHE_SLOT] + if getattr(self.k_cache, "hisparse_indexer_payload_format", None) == "ascend_indexer_e4m3fn_key_fp32_scale_v1": + if not self.enable_sparse_li_c8 or self.enable_sparse_li_c4 or k_li_scale is None: + raise ValueError("HiSparse's page-strided Indexer requires native C8 values and FP32 scales") + from vllm_ascend.distributed.kv_transfer.hisparse.binding import write_hisparse_indexer_c8 + + write_hisparse_indexer_c8( + indexer_k_cache, + self.k_cache.kv_cache[INDEXER_SCALE_CACHE_SLOT], + k_li, + k_li_scale, + slot_mapping.reshape(-1), + DeviceOperator.scatter_cache, + ) + return # float4_e2m1fn_x2 / float8_e8m0fnu are not supported by the scatter # kernel; view as uint8 (same 1-byte layout) for the cache write. if k_li.dtype == torch_npu.float4_e2m1fn_x2: diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index d42f8d8fccea..cf7844d0268a 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -91,8 +91,15 @@ class AscendHiSparseImpl(AscendSFAImpl): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - if self.enable_sparse_sfa_c8 or self.enable_sparse_sfa_turboquant: - raise ValueError("Ascend HiSparse currently requires an unquantized main MLA KV cache") + if self.enable_sparse_sfa_turboquant: + raise ValueError("Ascend HiSparse does not support TurboQuant") + if self.enable_sparse_sfa_c8: + if ( + self.c8_cache_dtype != torch.float8_e4m3fn + or (self.kv_lora_rank, self.qk_rope_head_dim) != (512, 64) + or self.sfa_qsfa_tile_size != 128 + ): + raise ValueError("HiSparse C8 requires native E4M3FN NoPE512/RoPE64 with FP32 scales") parallel = self.vllm_config.parallel_config if parallel.prefill_context_parallel_size > 1 or parallel.decode_context_parallel_size > 1: raise ValueError("Ascend HiSparse does not support context parallelism") @@ -101,10 +108,16 @@ def __init__(self, *args, **kwargs): if self.topk_indices_buffer is None: raise ValueError("Ascend HiSparse requires a model sparse-index group") builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) + if self.enable_sparse_sfa_c8: + from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + + row_width = C8_ROW_BYTES + else: + row_width = self.kv_lora_rank + self.qk_rope_head_dim self.index_group, self.index_group_index = builder.register_layer( self.has_indexer and not self.skip_topk, self.vllm_config, - row_width=self.kv_lora_rank + self.qk_rope_head_dim, + row_width=row_width, kv_cache_dtype=self.vllm_config.cache_config.cache_dtype, ) if not isinstance(self.index_group, HiSparseMLAIndexGroup): @@ -167,9 +180,19 @@ def allocate(): return current_workspace_manager().get_persistent_resource(key, allocate) + def _c8_length_workspace(self): + def allocate(): + device = self.index_group.logical_topk_indices.device + return tuple(torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) + + return current_workspace_manager().get_persistent_resource(("ascend_hisparse_c8_lengths",), allocate) + def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=None): - self._workspace(hidden_states.dtype) - if self.qk_rope_head_dim: + if self.enable_sparse_sfa_c8: + self._c8_length_workspace() + else: + self._workspace(hidden_states.dtype) + if self.qk_rope_head_dim and not self.enable_sparse_sfa_c8: # Include preprocessing scratch in profiling and allocate it before # graph capture, including dummy forwards without attention metadata. self._kv_workspace(hidden_states.dtype) @@ -180,6 +203,34 @@ def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): from vllm_ascend.ops.hisparse import scatter_packed_kv assert self.kv_a_layernorm is not None + if self.enable_sparse_sfa_c8: + from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope + from vllm_ascend.ops.hisparse_c8 import pack_c8_components + + values = kv_no_split.reshape(-1, self.num_kv_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim) + rope, latent, scales = custom_kv_rmsnorm_rope( + values, + self.kv_a_layernorm.weight, + cos, + sin, + self.kv_lora_rank, + self.qk_rope_head_dim, + epsilon=self.kv_a_layernorm.variance_epsilon, + dst_type=self.c8_cache_dtype, + tile_size=self.sfa_qsfa_tile_size, + ) + rows = latent.numel() // self.kv_lora_rank + packed = pack_c8_components(latent.contiguous(), rope.contiguous(), scales.contiguous()) + cache = self.hisparse_cache + destination, mapping, count = cache.write_target(rows, slots.numel()) + packed = packed[:count] + # Reuse the existing byte-capable scatter; no fused C8 kernel or + # payload conversion is needed for the baseline data path. + scatter_packed_kv( + packed, packed[:, :0], destination.view(torch.uint8), mapping, attn_metadata.query_start_loc[-1:] + ) + finish_hisparse_kv_update(cache, attn_metadata) + return rope, latent, scales values = kv_no_split.reshape(-1, self.kv_lora_rank + self.qk_rope_head_dim) num_tokens = values.shape[0] if self.qk_rope_head_dim: @@ -216,7 +267,18 @@ def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): finish_hisparse_kv_update(cache, attn_metadata) return rope, latent + def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_sfa, attn_metadata): + # exec_kv has already written through the current HiSparse destination + # and row mapping. The ordinary C8 store uses source slots instead, + # which would duplicate the write and can overwrite unrelated slab rows. + return k_pe, k_nope + def _attention_physical(self, query, query_rope, packed_cache, physical_topk): + if self.enable_sparse_sfa_c8: + from vllm_ascend.ops.hisparse_c8 import c8_physical_attention + + lengths = self._c8_length_workspace() + return c8_physical_attention(query, query_rope, packed_cache, physical_topk, self.scale, *lengths) from vllm_ascend.ops.hisparse import pack_attention_rows latent, rope, block_table, indices, query_ends, seq_lens = self._workspace(query.dtype) diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py index 8882c2b8e067..e8fe79772567 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py +++ b/vllm_ascend/distributed/kv_transfer/hisparse/backend.py @@ -83,5 +83,15 @@ def prepare_kv_cache_views( def is_current_stream_capturing(self) -> bool: return torch.npu.is_current_stream_capturing() + def gather_resident_rows(self, source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor) -> None: + # NPU advanced indexing does not support FP8. Copy all packed components + # as bytes, retaining Core's addressing and completion ownership. + fp8 = (torch.float8_e4m3fn, torch.float8_e5m2) + if source.dtype in fp8: + source = source.view(torch.uint8) + if destination.dtype in fp8: + destination = destination.view(torch.uint8) + super().gather_resident_rows(source, destination, row_ids) + def swap_blocks_batch(self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor) -> None: torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1) diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py index 8dc0b2937560..014b21916321 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py +++ b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py @@ -8,6 +8,8 @@ import torch from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, KVCacheConfig, KVCacheSpec, UniformTypeKVCacheSpecs, @@ -16,18 +18,39 @@ from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec -def align_hisparse_indexer_specs(specs: dict[str, KVCacheSpec]) -> dict[str, KVCacheSpec]: +def align_hisparse_indexer_specs( + specs: dict[str, KVCacheSpec], + *, + packed_row_bytes: int | None = None, +) -> dict[str, KVCacheSpec]: """Declare dense indexer-page addressing to Core's shared-slab planner. Core includes this alignment when sizing each physical device block, before choosing the block count. Source MLA specs and their host layout stay intact. """ - return { - name: replace(spec, block_stride_alignment=math.lcm(spec.get_block_stride_alignment(), spec.page_size_bytes)) - if isinstance(spec, AscendSFAIndexerCacheSpec) - else spec - for name, spec in specs.items() - } + if packed_row_bytes is not None and packed_row_bytes <= 0: + raise ValueError("HiSparse packed row width must be positive") + aligned = {} + for name, spec in specs.items(): + alignment = spec.get_block_stride_alignment() + updates = {} + if isinstance(spec, AscendSFAIndexerCacheSpec): + _validate_indexer_spec(spec) + if spec.cache_sparse_li_c8: + # K and scale occupy separate contiguous regions within each + # page. Reserve both before Core computes the layer stride. + updates["page_size_padded"] = max(spec.page_size_bytes, spec.page_size_padded or 0) + alignment = math.lcm(alignment, spec.head_size, torch.float32.itemsize) + else: + alignment = math.lcm(alignment, spec.page_size_bytes) + elif packed_row_bytes is not None and isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)): + # Core derives these synthetic specs without the source row alignment. + # Declare it before shared-slab sizing, never round runtime addresses. + alignment = math.lcm(alignment, packed_row_bytes) + if alignment != spec.get_block_stride_alignment(): + updates["block_stride_alignment"] = alignment + aligned[name] = replace(spec, **updates) if updates else spec + return aligned def get_hisparse_indexer_specs( @@ -41,13 +64,91 @@ def get_hisparse_indexer_specs( if isinstance(spec, UniformTypeKVCacheSpecs): spec = spec.kv_cache_specs[name] if isinstance(spec, AscendSFAIndexerCacheSpec): - if spec.scale_dim or spec.cache_sparse_li_c8 or spec.cache_sparse_li_c4: - raise ValueError("Ascend HiSparse currently requires an unquantized indexer cache.") + _validate_indexer_spec(spec) indexer_specs[name] = spec return indexer_specs +def _validate_indexer_spec(spec: AscendSFAIndexerCacheSpec) -> None: + if spec.cache_sparse_li_c4: + raise ValueError("HiSparse indexer C4 is not supported by this adapter.") + if spec.cache_sparse_li_c8: + if not ( + spec.dtype == torch.float8_e4m3fn + and spec.scale_dtype == torch.float32 + and spec.scale_dim == 1 + and spec.num_kv_heads == 1 + and spec.head_size == 128 + and spec.block_size == 128 + and spec.sfa_dcp_replicated_indexer_size == 1 + ): + raise ValueError("HiSparse indexer C8 requires E4M3FN K128 with one FP32 scale per token, block128.") + elif spec.scale_dim: + raise ValueError("HiSparse indexer scales require an explicit supported quantization format.") + + +def _quantized_indexer_views(cache: torch.Tensor, spec: AscendSFAIndexerCacheSpec): + """Alias the current Core allocation; retain no additional storage owner.""" + _validate_indexer_spec(spec) + pages, heads, block, width = cache.shape + stride = cache.stride(0) + key_bytes = block * heads * width + scale_bytes = block * heads * torch.float32.itemsize + offset = cache.storage_offset() + if ( + cache.dtype != torch.float8_e4m3fn + or (heads, block, width) != (1, spec.block_size, spec.head_size) + or cache.stride()[1:] != (block * width, width, 1) + or pages <= 0 + or stride < key_bytes + scale_bytes + or stride % width + or (offset + key_bytes) % torch.float32.itemsize + or stride % torch.float32.itemsize + or offset + (pages - 1) * stride + key_bytes + scale_bytes > cache.untyped_storage().nbytes() + ): + raise ValueError("Invalid HiSparse FP8 indexer K/scale page geometry.") + key = torch.as_strided(cache, (pages, block, heads, width), (stride, width, width, 1)) + # Reinterpret only a zero-length contiguous alias before constructing the + # scale view; dtype conversion of a strided tensor would impose unrelated + # constraints on its last axis. Offsets below are absolute storage offsets. + raw = torch.as_strided(cache, (0,), (1,), storage_offset=0).view(torch.float32) + scale = torch.as_strided( + raw, (pages, block, heads, 1), (stride // 4, 1, 1, 1), storage_offset=(offset + key_bytes) // 4 + ) + return key, scale + + +def write_hisparse_indexer_c8(key, scale, values, scales, slots, scatter) -> None: + """Write byte-preserving K and FP32 scales using this layer's page stride. + + Derive aliases from the current cache tuple, so profiling-cache replacement + and teardown cannot leave a second set of persistent cache references. + """ + pages, block, heads, width = key.shape + stride = key.stride(0) + count = slots.numel() + if values.numel() != count * width or scales.numel() != count: + raise ValueError("HiSparse FP8 indexer writes require matching token/slot counts.") + if values.dtype != key.dtype or scales.dtype != torch.float32: + raise ValueError("HiSparse FP8 indexer writes require E4M3FN data and FP32 scales.") + key_rows = ((pages - 1) * stride + block * width) // width + scale_rows = (pages - 1) * scale.stride(0) + block + raw_key = torch.as_strided(key, (0,), (1,), storage_offset=0).view(torch.int8) + key_write = torch.as_strided(raw_key, (key_rows, width), (width, 1), storage_offset=key.storage_offset()) + scale_write = torch.as_strided(scale, (scale_rows, 1), (1, 1)) + # Ascend V2 supplies int32 slots. Page-strided scale addresses can exceed + # INT32_MAX even when logical slots fit; widen before multiplying strides. + logical = slots.reshape(-1).to(torch.int64) + page = logical // block + row = logical % block + # Preserve the native -1 padding sentinel; never turn it into a valid row. + key_slots = torch.where(logical >= 0, page * (stride // width) + row, -1) + scale_slots = torch.where(logical >= 0, page * scale.stride(0) + row, -1) + scatter(key_write, key_slots.view(-1, 1), values.contiguous().view(torch.int8).reshape(-1, width)) + scatter(scale_write, scale_slots.view(-1, 1), scales.reshape(-1, 1)) + + def _bind_indexer_views( kv_caches: dict[str, Any], indexer_specs: dict[str, AscendSFAIndexerCacheSpec], forward_context: dict[str, Any] ) -> None: @@ -57,6 +158,12 @@ def _bind_indexer_views( # shared slab and let the layer scale its block IDs and write slots. if cache.ndim != 4 or cache.shape[1] != spec.num_kv_heads: raise ValueError(f"Unexpected HiSparse indexer cache shape for {name}: {tuple(cache.shape)}") + if spec.cache_sparse_li_c8: + kv_caches[name] = _quantized_indexer_views(cache, spec) + forward_context[name].hisparse_block_stride = 1 + # Declare only after validating this allocation's actual K/scale views. + forward_context[name].hisparse_indexer_payload_format = "ascend_indexer_e4m3fn_key_fp32_scale_v1" + continue block_size = cache.shape[2] page_elements = block_size * spec.num_kv_heads * spec.head_size if cache.stride(0) % page_elements: @@ -69,4 +176,5 @@ def _bind_indexer_views( stride=(page_elements, spec.num_kv_heads * spec.head_size, spec.head_size, 1), ) forward_context[name].hisparse_block_stride = block_stride + forward_context[name].hisparse_indexer_payload_format = "bf16" kv_caches[name] = (indexer_cache,) diff --git a/vllm_ascend/ops/hisparse_c8.py b/vllm_ascend/ops/hisparse_c8.py new file mode 100644 index 000000000000..ce29c398edec --- /dev/null +++ b/vllm_ascend/ops/hisparse_c8.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Native C8 payload/layout bridge; quantization and Attention stay in existing ops.""" + +import math + +import torch + +_LATENT = 512 +_ROPE_BYTES = 128 +_SCALE_BYTES = 16 +C8_ROW_BYTES = _LATENT + _ROPE_BYTES + _SCALE_BYTES +_INT32_MAX = 2**31 - 1 + + +def pack_c8_components(nope: torch.Tensor, rope: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Pack existing SFA quantizer outputs using byte copies only.""" + if nope.dtype != torch.float8_e4m3fn: + raise ValueError("HiSparse C8 requires FP8 E4M3 NoPE") + components = (nope, rope, scales) + if any(t.device != nope.device or not t.is_contiguous() for t in components): + raise ValueError("HiSparse C8 components must be contiguous on one device") + rows = nope.numel() // _LATENT + if rows == 0 or nope.numel() != rows * _LATENT: + raise ValueError("HiSparse C8 requires complete 512-element NoPE rows") + if rope.numel() * rope.element_size() != rows * _ROPE_BYTES: + raise ValueError("HiSparse C8 requires 64 BF16 RoPE elements per row") + if scales.numel() * scales.element_size() != rows * _SCALE_BYTES: + raise ValueError("HiSparse C8 requires four FP32 scale values per row") + return torch.cat( + [ + t.view(torch.uint8).reshape(rows, width) + for t, width in zip(components, (_LATENT, _ROPE_BYTES, _SCALE_BYTES)) + ], + dim=-1, + ) + + +def c8_storage_view(cache: torch.Tensor) -> tuple[torch.Tensor, int, int, int]: + """Alias all addressable pages, retaining HMA gaps and the layer offset. + + The allocator must align the physical block stride to a whole packed row. + The returned TND view shares storage; no payload is gathered or dequantized. + """ + if cache.ndim != 3 or cache.dtype != torch.float8_e4m3fn or cache.shape[-1] != C8_ROW_BYTES: + raise ValueError("HiSparse C8 cache must be FP8 [pages, block_size, 656]") + pages, block, width = cache.shape + stride = cache.stride(0) + if min(pages, block) <= 0 or cache.stride(1) != width or cache.stride(2) != 1: + raise ValueError("HiSparse C8 requires nonempty pages with dense token rows") + if stride < block * width or stride % width: + raise ValueError("HiSparse C8 physical block stride must cover whole packed rows") + logical_capacity = pages * block + span = (pages - 1) * stride + block * width + addressable_rows = span // width + if max(logical_capacity, addressable_rows) > _INT32_MAX: + raise ValueError("HiSparse C8 physical row IDs must fit int32") + offset = cache.storage_offset() + if (offset + span) * cache.element_size() > cache.untyped_storage().nbytes(): + raise ValueError("HiSparse C8 view exceeds its backing allocation") + view = torch.as_strided(cache, (addressable_rows, 1, width), (width, width, 1), storage_offset=offset) + return view, block, stride // width, logical_capacity + + +def map_c8_physical_indices(indices: torch.Tensor, block: int, stride_rows: int, capacity: int) -> torch.Tensor: + """Map virtual page rows to native contiguous-row addresses on device.""" + if indices.ndim != 2 or indices.dtype != torch.int32 or block <= 0 or stride_rows < block: + raise ValueError("HiSparse C8 requires int32 [queries, top_k] indices and valid page geometry") + if capacity <= 0 or capacity % block or (capacity // block - 1) * stride_rows + block > _INT32_MAX: + raise ValueError("HiSparse C8 physical row capacity is invalid") + valid = (indices >= 0) & (indices < capacity) + safe = indices.clamp_min(0).to(torch.int64) + mapped = safe.div(block, rounding_mode="floor") * stride_rows + safe.remainder(block) + return torch.where(valid, mapped, -1).to(torch.int32) + + +def compact_c8_indices(indices: torch.Tensor) -> torch.Tensor: + """Keep valid IDs in their original order, including every duplicate.""" + if indices.ndim != 2 or indices.dtype != torch.int32: + raise ValueError("C8 compaction requires int32 [queries, top_k] indices") + width = indices.shape[1] + if width > 2**23: + raise ValueError("C8 compaction keys must remain exact FP32 integers") + positions = torch.arange(width, dtype=torch.float32, device=indices.device) + keys = positions + (indices < 0).to(torch.float32) * width + order = torch.argsort(keys, dim=-1) + return torch.gather(indices, -1, order) + + +def c8_physical_attention(query, query_rope, cache, indices, scale, query_ends, key_length): + """Adapt native packed rows to the existing quantized SFA operation. + + Keep compaction stable, including duplicates. The two current lengths are + device writes captured on the consumer stream; no Host read or stream/event + policy is introduced here. + """ + view, block, stride_rows, capacity = c8_storage_view(cache) + rows = query.shape[0] + if query.ndim != 3 or query.dtype != torch.bfloat16 or query.shape[-1] != _LATENT or rows <= 0: + raise ValueError("HiSparse C8 query must be nonempty BF16 [queries, heads, 512]") + if query_rope.shape != (*query.shape[:-1], 64) or query_rope.dtype != query.dtype: + raise ValueError("HiSparse C8 requires matching BF16 query RoPE") + if indices.ndim != 2 or indices.shape[0] != rows or indices.shape[1] <= 0 or indices.shape[1] % 128: + raise ValueError("HiSparse C8 top-k width must be positive and divisible by 128") + if query.device.type != "npu" or any( + tensor.device != query.device for tensor in (query_rope, cache, indices, query_ends, key_length) + ): + raise ValueError("HiSparse C8 inputs must share one NPU") + if any(tensor.shape != (1,) or tensor.dtype != torch.int32 for tensor in (query_ends, key_length)): + raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") + if not math.isfinite(scale) or scale <= 0: + raise ValueError("HiSparse C8 attention scale must be finite and positive") + mapped = compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) + packed_query = torch.cat((query, query_rope), dim=-1) + query_ends.fill_(rows) + key_length.fill_(view.shape[0]) + result = torch.ops._C_ascend.npu_kv_quant_sparse_flash_attention_vllm( + query=packed_query, + key=view, + value=view, + sparse_indices=mapped.unsqueeze(1), + scale_value=scale, + sparse_block_size=1, + actual_seq_lengths_query=query_ends, + actual_seq_lengths_kv=key_length, + layout_query="TND", + layout_kv="TND", + sparse_mode=0, + attention_mode=2, + quant_scale_repo_mode=1, + tile_size=128, + key_quant_mode=2, + value_quant_mode=2, + rope_head_dim=64, + return_softmax_lse=True, + ) + if isinstance(result, tuple): + result = result[0] + return torch.where((mapped >= 0).any(dim=-1)[:, None, None], result, 0) diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index da4615c3ca33..7364c427eca6 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -275,8 +275,13 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if isinstance(attn_module, MLAAttention) and vllm_config.attention_config.hisparse_config is not None: # Keep Core's row-packed source spec and actual index-group ownership. + source_layout = {} + if getattr(attn_module.impl, "enable_sparse_sfa_c8", False): + from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + + source_layout.update(head_size=C8_ROW_BYTES, dtype=attn_module.impl.c8_cache_dtype) kv_cache_spec[layer_name] = replace( - spec, is_index_group_leader=attn_module.hisparse_cache.runtime.is_group_leader + spec, is_index_group_leader=attn_module.hisparse_cache.runtime.is_group_leader, **source_layout ) continue if isinstance(attn_module, MLAAttention): @@ -429,7 +434,14 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: from vllm_ascend.distributed.kv_transfer.hisparse.binding import align_hisparse_indexer_specs - return align_hisparse_indexer_specs(resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers)) + packed_row_bytes = None + if get_ascend_config().enable_sparse_sfa_c8: + from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + + packed_row_bytes = C8_ROW_BYTES + return align_hisparse_indexer_specs( + resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers), packed_row_bytes=packed_row_bytes + ) return kv_cache_spec From b3a6bf2f631234f5eeb54697f2b5a2d90611dd32 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 20:09:16 +0800 Subject: [PATCH 007/120] test(hisparse): narrow foundation coverage to Ascend adapters Remove the embedded activation dump, generic stream probe, duplicate prefill and order checks, and mock-only Core forwarding tests. Reuse existing SFA preprocessing as the reference for one BF16/C8 destination test. Retain device resolver/refill/invalidation, packed layouts and graph coverage. Production code and Attention kernels are unchanged. Twenty source-isolated CPU cases pass; NPU qualification remains pending. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 22 +- .../one_card/test_hisparse_attention.py | 73 ------ .../one_card/test_hisparse_kv_preprocess.py | 243 +++--------------- .../one_card/test_hisparse_ops.py | 100 +------ .../test_hisparse_foundation_cpu.py | 37 --- 5 files changed, 63 insertions(+), 412 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 149db68d3d17..2d677c897030 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -95,8 +95,8 @@ current request/state mappings remain separate. Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces are used directly. The Ascend backend preserves Core's event placement and does not -replace its Torch bindings. The NPU tests cover cross-stream waits followed by -CANN workspace allocation and changing graph replays. +replace its Torch bindings. This PR tests its own NPU operations and changing +graph replays rather than duplicating generic Torch stream tests. The model-owned index-group builder is carried in Core's existing `MLAModules.index_group_builder` field, preserving the common MLA wrapper ABI. @@ -127,7 +127,7 @@ python -m pytest -q --noconftest \ These check current Host offsets, failed-unregister ownership, group scratch isolation, strided indexer aliases, planner alignment, unsupported formats, -Core coordinator forwarding and capture-only mirror suppression. C8 checks also +and capture-only mirror suppression. C8 checks also cover opaque component bytes, nonzero slab offsets, stable duplicate/padding ordering, Indexer key/scale writes and planner alignment before allocation. They do not exercise an installed NPU package or qualify device execution. @@ -136,6 +136,22 @@ The weight-loading wrapper keeps Ascend main's explicit activation-dtype ABI, already covered by `tests/ut/ops/test_mla_weight_processing.py`; this foundation does not add another calling convention or a duplicate test module. +Reuse the existing suites for behavior this PR does not implement: + +| Behavior | Existing coverage | +| --- | --- | +| Core planning and HMA budgets | Core `tests/v1/core/test_kv_cache_planning.py` and `test_kv_cache_utils.py` | +| Core Host-mirror ordering, zero-token steps and request residency | Core `tests/v1/kv_connector/unit/test_hisparse_connector.py` | +| Core prefill plan and sparse layout rules | Core `tests/v1/attention/test_sparse_mla_backends.py` and `test_sparse_mla_kv_cache_layout.py` | +| Native SFA preprocessing and operator dispatch | Ascend `tests/ut/attention/test_sfa_v1.py` | + +These suites retain their original platform/dependency requirements. In +particular, Core's CUDA-kernel cases do not qualify Ascend kernels. The NPU +modules below keep coverage for Ascend resolver/refill/invalidation, strided +Resident/Host staging and the BF16/C8 data bridge. Preprocessing uses the existing +native SFA operations as its reference; it contains no embedded activation dump +or independent-versus-fused rounding requirement. + Real-NPU cases are included for resolver/refill/overwrite invalidation, independent FP32 Attention references, selection order/duplicates, native BF16/C8 preprocessing, C8 key/scale writes, packed-byte scatter and Resident/Host prefill staging, in diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index ca1896280a4a..a294e8ae4efa 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -105,79 +105,6 @@ def test_sparse_attention_reference(num_queries, strided, capture): graph.reset() -@pytest.mark.parametrize("capture", [False, True]) -def test_valid_order_duplicates_and_interleaved_padding(monkeypatch, capture): - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - from vllm_ascend.ops import hisparse - - torch.manual_seed(61) - device, topk, latent, rope, capacity = "npu:0", 2048, 512, 64, 32 - workspace = ( - torch.empty((capacity, topk, 1, latent), dtype=torch.bfloat16, device=device), - torch.empty((capacity, topk, 1, rope), dtype=torch.bfloat16, device=device), - torch.arange(capacity * topk // 128, dtype=torch.int32, device=device).view(capacity, -1), - torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), - torch.arange(1, capacity + 1, dtype=torch.int32, device=device), - torch.full((capacity,), topk, dtype=torch.int32, device=device), - ) - impl = SimpleNamespace( - topk=topk, - kv_lora_rank=latent, - qk_rope_head_dim=rope, - scale=0.0625, - _workspace=lambda _: workspace, - enable_sparse_sfa_c8=False, - ) - query = torch.randn(3, 8, latent, dtype=torch.bfloat16, device=device) - query_rope = torch.randn(3, 8, rope, dtype=torch.bfloat16, device=device) - cache = torch.randn(1, 128, latent + rope, dtype=torch.bfloat16, device=device) - indices = torch.full((3, topk), -1, dtype=torch.int32) - indices[0, :6] = torch.tensor([7, -1, 2, 7, -1, 0]) - indices[2, :6] = torch.tensor([0, 5, -1, 3, 5, 1]) - device_indices = indices.to(device) - observed = torch.empty_like(device_indices) - original_pack = hisparse.pack_attention_rows - - def observe(cache, physical, selected_latent, selected_rope, latent_dim): - observed.copy_(physical) - return original_pack(cache, physical, selected_latent, selected_rope, latent_dim) - - monkeypatch.setattr(hisparse, "pack_attention_rows", observe) - - def forward(): - return AscendHiSparseImpl._attention_physical(impl, query, query_rope, cache, device_indices) - - graph = None - try: - if capture: - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): - forward() - forward() - torch.npu.synchronize() - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph, stream=stream): - result = forward() - indices[0, :6] = torch.tensor([-1, 9, 1, 9, 4, -1]) - device_indices.copy_(indices) - for _ in range(3): - graph.replay() - else: - result = forward() - expected = torch.full_like(indices, -1) - for row in range(3): - valid = indices[row][indices[row] >= 0] - expected[row, : valid.numel()] = valid - torch.testing.assert_close(observed.cpu(), expected, rtol=0, atol=0) - assert torch.isfinite(result.cpu()).all() - assert torch.count_nonzero(result[1].cpu()) == 0 - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() - - @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) def test_c8_attention_uses_existing_native_operator(strided, capture): diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index b382532deed3..d97b05bb7cad 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Keep HiSparse KV preprocessing numerically identical to native SFA.""" +"""Check HiSparse destinations against the existing SFA preprocessing ops.""" from types import SimpleNamespace @@ -8,30 +8,6 @@ import torch import torch_npu -# One fixed BF16 activation row with unit gamma and epsilon 1e-5. Its RMS -# normalization lands within 8e-8 of three BF16 rounding midpoints. Repeating -# this row reproduces the standalone-vs-fused difference; no model is loaded. -_ROUNDING_ROW_HEX = ( - "bc3df03d173e0d3c1bbe9cbcefbdb73822bde4ba14bd2ebda73dff3d003d1e3d1e3d053eec3da83d21be63bd8d3d40bd1cbda03bb1bd1abe04bb583c" - "563ec63bfdbdd1bd6b3d70bde73d103e343e0e3de13d1ebc323e03bd77bdec3c3abd2abca6bd043d4ebdf5bd1e3ce13da73dbebd083da73d063ec03d" - "1cbe71bced3c92bd3cbdd13d71bdc83dc13ca2bb85bcb1bd533d4a3dbabd57bdb43cc33c0bbb163e683d953dbeba65bdcb3ac73c56bd23bd86bd12bd" - "9ebd2abeb93def3ce33d7f3d783dca3cc3bc9b3d6b3d513d043d1d3dd33da03c9c3d023d59bd943dc9bd29bd2dbd8f3dcc3c0c3cba3d863d1f3ea23d" - "3b3db73d083df7bc9d3cc8bd43be19bef5bd91bd4ebe00bd65b9403e7f3dc5bcf6bb1b3d82bc883d25bd81bbe2bd553dc7bc673c0d3ce3bd74bc36bd" - "ba3d213c9d3b2dbbee3d6a3d823d8abb1dbd4abd883d383d85bdd73dacbdc5bae2bd57bdd3bd32bec9bd05be06bd6a3cd93d333dbf3c8ebcb73dc73a" - "11bd89bd213b0abc09bd663b5abccdbd0b3c623d753d603eb3bdc6bd813d9d3b2c3cb6bb9c3d3c3d853c96bb903d3fbdf33cd2ba0ebd25be05bd353d" - "edbdd3bb5b3e843d0e3dd53ca93cefbc9cbc633d463b57bd033cc8bda13dbe3dc4bd2f3e41bdeabd0f3dc9bd7a3cc13df8bd5b3c263d223ed83dcd3d" - "353d8ebdc53d673d0dbeccbda1bd1bbd4abdfbbca7bd3bbd00bd2bbc4d3c70bc74bdf13dd7bc2b3e31bd4abcd0bdb8bd9fbce4bba13b95bd2d3e51bd" - "1c3e9cbddfbae73d37bdddbce5bcaebae7bd39bda3bd0ebc0f3c83bd85bc34bbcf3cb53cc93c27bbabbdd6bd923c013ea43d3a3d2ebd643da13de33d" - "2e3c8dbd943c9ebdc63d92bb9ebd673d14bd283d8fbce43d323d883c493c5d3d66b85cbdafbdb1bd2d3c373d873d4abda5bdb9bcb8bd71bd76bd14be" - "443b083daf3cdf3d743d91bd913ddf3d3cbd2d3d53bd1cbe95bd15be88bdf53b8c3c763b9e3de33d2e3e82bc3e3d9ebd8dbd073ce3bc483d353d023e" - "4cbbb83d0b3e3c3e0abe74bd04be0fbdaabdd6bd3a3bf93a573c5f3e0c3d283dcc3d2b3d043eae3ca8bd90bd0abcb33dbebdc9bdee3cfbbdbd397dbd" - "243dbe3df13d203d84bd043dc0bdb0bca6bd06bd313dacbdf7bc05be263ec93dc03dd73c05bdad3d32bbcfbd31be5cbd84bda7bdd43d1d3d143e6a3d" - "053ead3b57bd173dd0bb83bdf73ce3bd773e323e99bbea3d793d5bbd90bce53d07bed2bc79bd98bd81bc713d41bd2dbc153d443c493e2b3e29be53bd" - "40be02bed83d13bdba3da2bc7f3d2bbded3dde3dba3cdcbd9dbc533e18beacbd54bd8fbb003db2bdc63df5bc033d763e1f3e13bcd6bc47bd33bd27be" - "ae3d323d183dc0bdbfbd50bd89bcc43cfebdacbdbabcf8bcea3cce3d22be963d093e153e2c3d2d3b2abcdc3d0cbc48bdb5bd1d3d40bdddbdb7bce23d" - "063e123b" -) - def _setup(): if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): @@ -40,30 +16,41 @@ def _setup(): @pytest.mark.parametrize("capture", [False, True]) -def test_c8_preprocess_writes_hisparse_mapping_once(capture): +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_preprocess_writes_hisparse_mapping_once(cache_dtype, capture): _setup() from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope rows, block = 5, 128 - values = torch.randn(rows, 576, dtype=torch.bfloat16, device="npu") + c8 = cache_dtype == torch.float8_e4m3fn + width = 656 if c8 else 576 + torch.manual_seed(71) + values = torch.randn(rows, 2624, dtype=torch.bfloat16, device="npu")[:, 2048:] gamma = torch.ones(512, dtype=torch.bfloat16, device="npu") cos = torch.ones(rows, 1, 1, 64, dtype=torch.bfloat16, device="npu") sin = torch.zeros_like(cos) - raw = torch.full((3, 2, block, 656), 90, dtype=torch.uint8, device="npu") - destination = raw[:, 1].view(torch.float8_e4m3fn) + raw = torch.full((3, 2, block, width * cache_dtype.itemsize), 90, dtype=torch.uint8, device="npu") + destination = raw[:, 1].view(cache_dtype) logical = [-1, 129, 0, 257, 128] slots = torch.tensor(logical, dtype=torch.int64, device="npu") metadata = SimpleNamespace( query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="npu"), hisparse_for_capture=True ) + scratch = ( + torch.empty(1, block, 1, 512, dtype=values.dtype, device="npu"), + torch.empty(1, block, 1, 64, dtype=values.dtype, device="npu"), + torch.arange(block, dtype=torch.int64, device="npu"), + ) impl = SimpleNamespace( - enable_sparse_sfa_c8=True, + enable_sparse_sfa_c8=c8, c8_cache_dtype=torch.float8_e4m3fn, sfa_qsfa_tile_size=128, kv_lora_rank=512, qk_rope_head_dim=64, num_kv_heads=1, + _kv_staging_capacity=block, + _kv_workspace=lambda dtype: scratch, kv_a_layernorm=SimpleNamespace(weight=gamma, variance_epsilon=1e-5), hisparse_cache=SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)), ) @@ -72,7 +59,8 @@ def test_c8_preprocess_writes_hisparse_mapping_once(capture): source_cache = (torch.empty(0, device="npu"),) def forward(): - rope, latent, scales = AscendHiSparseImpl.exec_kv(impl, values, cos, sin, source_cache, slots, metadata) + rope, latent, *scales = AscendHiSparseImpl.exec_kv(impl, values, cos, sin, source_cache, slots, metadata) + scales = scales[0] if scales else None AscendHiSparseImpl._store_parallel_kv(impl, rope, latent, scales, source_cache, slots, metadata) graph = None @@ -92,11 +80,29 @@ def forward(): else: graph.replay() torch.npu.synchronize() - rope, latent, scales = custom_kv_rmsnorm_rope( - values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 - ) + if c8: + rope, latent, scales = custom_kv_rmsnorm_rope( + values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 + ) + components = (latent, rope, scales) + else: + # Reuse native SFA as the oracle; test only the new destination, + # strides, row mapping and padding, not RMSNorm arithmetic. + reference_k, reference_r = torch.empty_like(scratch[0]), torch.empty_like(scratch[1]) + torch_npu.npu_kv_rmsnorm_rope_cache( + values.view(rows, 1, 1, 576), + gamma, + cos, + sin, + scratch[2][:rows], + reference_r, + reference_k, + epsilon=1e-5, + cache_mode="PA", + ) + components = (reference_k.view(-1, 512)[:rows], reference_r.view(-1, 64)[:rows]) payload = torch.cat( - [part.contiguous().view(torch.uint8).reshape(rows, -1).cpu() for part in (latent, rope, scales)], -1 + [part.contiguous().view(torch.uint8).reshape(rows, -1).cpu() for part in components], -1 ) expected = torch.full(raw.shape, 90, dtype=torch.uint8) for row, slot in enumerate(logical): @@ -107,170 +113,3 @@ def forward(): torch.npu.synchronize() if graph is not None: graph.reset() - - -@pytest.mark.parametrize("capture", [False, True]) -def test_fused_preprocess_preserves_native_rounding_and_packed_writes(monkeypatch, capture): - _setup() - from vllm_ascend.attention import sfa_hisparse - - tokens, width, capacity = 17, 576, 128 - row = torch.frombuffer(bytearray.fromhex(_ROUNDING_ROW_HEX), dtype=torch.bfloat16).clone() - assert row.shape == (512,) - values = torch.empty(tokens, 2624, dtype=torch.bfloat16, device="npu")[:, 2048:] - source = torch.cat((row, torch.arange(64).to(torch.bfloat16) / 64)).expand(tokens, -1).clone() - values.copy_(source.to("npu")) - cos = torch.ones(tokens, 1, 1, 64, dtype=values.dtype, device="npu") - sin = torch.zeros_like(cos) - gamma = torch.ones(512, dtype=values.dtype, device="npu") - local_slots = torch.arange(tokens, dtype=torch.int64, device="npu") - reference_k = torch.empty(1, 128, 1, 512, dtype=values.dtype, device="npu") - reference_r = torch.empty(1, 128, 1, 64, dtype=values.dtype, device="npu") - - def native(): - torch_npu.npu_kv_rmsnorm_rope_cache( - values.view(tokens, 1, 1, width), - gamma, - cos, - sin, - local_slots, - reference_r, - reference_k, - epsilon=1e-5, - cache_mode="PA", - ) - return torch.cat((reference_k.view(-1, 512)[:tokens], reference_r.view(-1, 64)[:tokens]), -1).cpu() - - reference = native() - old_split, _ = torch_npu.npu_rms_norm(values[:, :512], gamma, epsilon=1e-5) - # A random oracle can miss this rare rounding boundary. This fixture must - # demonstrate the old path's disagreement before checking the replacement. - assert not torch.equal(old_split.cpu(), reference[:, :512]) - - resources = {} - - def persistent(key, allocate): - if key not in resources: - resources[key] = allocate() - return resources[key] - - monkeypatch.setattr( - sfa_hisparse, "current_workspace_manager", lambda: SimpleNamespace(get_persistent_resource=persistent) - ) - raw = torch.empty(4 * 6 * 128 * width * 2, dtype=torch.int8, device="npu") - backing = raw.view(values.dtype) - target = backing.as_strided((4, 128, width), (6 * 128 * width, width, 1), 2 * 128 * width) - slots_cpu = 256 + torch.arange(tokens, dtype=torch.int64) - slots_cpu[::4] = -1 - slots = slots_cpu.to("npu") - ends = torch.tensor([0, tokens], dtype=torch.int32, device="npu") - mirrored = [] - handle = SimpleNamespace( - write_target=lambda count, slot_count: (target, slots, count), - finish_kv_update=lambda: mirrored.append(True), - ) - impl = SimpleNamespace( - enable_sparse_sfa_c8=False, - kv_lora_rank=512, - qk_rope_head_dim=64, - num_kv_heads=1, - _kv_staging_capacity=capacity, - index_group=SimpleNamespace(logical_topk_indices=slots), - kv_a_layernorm=SimpleNamespace(weight=gamma, variance_epsilon=1e-5), - hisparse_cache=handle, - ) - impl._kv_workspace = lambda dtype: sfa_hisparse.AscendHiSparseImpl._kv_workspace(impl, dtype) - first = impl._kv_workspace(values.dtype) - assert all(a is b for a, b in zip(first, impl._kv_workspace(values.dtype))) - assert first[2].dtype == torch.int64 and first[2].numel() == capacity - - def forward(): - return sfa_hisparse.AscendHiSparseImpl.exec_kv( - impl, values, cos, sin, None, slots, SimpleNamespace(query_start_loc=ends) - ) - - graph = None - try: - backing.fill_(-11) - if capture: - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): - forward() - forward() - torch.npu.synchronize() - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph, stream=stream): - rope, latent = forward() - else: - rope, latent = forward() - torch.testing.assert_close(latent.cpu(), reference[:, :512], rtol=0, atol=0) - torch.testing.assert_close(rope.cpu(), reference[:, 512:], rtol=0, atol=0) - assert mirrored - for iteration in range(3): - values.copy_((source * (iteration + 2)).to("npu")) - reference = native() - slots_cpu = 128 + torch.arange(tokens, dtype=torch.int64) - slots_cpu[1::4] = -1 - slots.copy_(slots_cpu) - valid = tokens - iteration - 1 - ends[1].fill_(valid) - backing.fill_(-11) - if capture: - graph.replay() - else: - rope, latent = forward() - expected_backing = torch.full_like(backing, -11).cpu() - expected = expected_backing.as_strided(target.shape, target.stride(), target.storage_offset()) - for row_id, slot in enumerate(slots_cpu.tolist()): - if row_id < valid and slot >= 0: - expected[slot // 128, slot % 128].copy_(reference[row_id]) - torch.testing.assert_close(backing.cpu(), expected_backing, rtol=0, atol=0) - torch.testing.assert_close(latent.cpu(), reference[:, :512], rtol=0, atol=0) - torch.testing.assert_close(rope.cpu(), reference[:, 512:], rtol=0, atol=0) - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() - - -def test_nope_preprocess_uses_native_layernorm(monkeypatch): - _setup() - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - - values = torch.arange(3 * 512, dtype=torch.float32, device="npu").view(3, 512).to(torch.bfloat16) - target = torch.full((1, 128, 512), -7, dtype=values.dtype, device="npu") - slots = torch.tensor([1, -1, 4], dtype=torch.int64, device="npu") - calls = [] - - def norm(x): - calls.append(True) - return x * 2 - - def unexpected(*args, **kwargs): - raise AssertionError("NoPE must not call the RoPE fused kernel") - - monkeypatch.setattr(torch_npu, "npu_kv_rmsnorm_rope_cache", unexpected) - impl = SimpleNamespace( - enable_sparse_sfa_c8=False, - kv_lora_rank=512, - qk_rope_head_dim=0, - kv_a_layernorm=norm, - hisparse_cache=SimpleNamespace( - write_target=lambda count, slot_count: (target, slots, count), - finish_kv_update=lambda: None, - ), - ) - rope, latent = AscendHiSparseImpl.exec_kv( - impl, - values, - None, - None, - None, - slots, - SimpleNamespace(query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device="npu")), - ) - expected = torch.full_like(target, -7).cpu() - expected[0, 1].copy_((values[0] * 2).cpu()) - torch.testing.assert_close(target.cpu(), expected, rtol=0, atol=0) - assert calls and rope.shape == (3, 0) and latent.shape == values.shape diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index c60c9c102661..8d3374239e80 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -6,7 +6,7 @@ import pytest import torch -import torch_npu +import torch_npu # noqa: F401 from vllm_ascend.ops.hisparse import ( HiSparseNPUOps, @@ -25,76 +25,6 @@ def _require_mapped_host_device(): torch.npu.set_device(0) -@pytest.mark.parametrize("capture", [False, True]) -@pytest.mark.parametrize("wait_event", [False, True]) -def test_generic_stream_waits_before_cann_workspace(capture, wait_event): - """Core's generic fork/join must propagate capture before CANN allocation.""" - initial = (torch.arange(64 * 64).reshape(64, 64) % 23 + 1).bfloat16() - x = initial.npu() - weight = torch.eye(64, dtype=x.dtype, device=x.device) - gamma = torch.ones(64, dtype=x.dtype, device=x.device) - output = torch.empty_like(x) - copy = torch.Stream(device=x.device) - ready = torch.Event() - done = None - owners = [] - - def forward(): - nonlocal done - compute = torch.accelerator.current_stream() - x.add_(1) - ready.record(compute) - if wait_event: - copy.wait_event(ready) - else: - copy.wait_stream(compute) - with copy: - # Keep allocation immediately after the wait. A synchronization or - # capture-state query here can hide a Host-queue propagation error. - product = torch.mm(x, weight) - normalized, _ = torch_npu.npu_rms_norm(product, gamma, epsilon=1e-6) - result = product + normalized - owners.extend((product, normalized, result)) - if done is None: - done = torch.Event() - done.record(copy) - compute.wait_event(done) - consumed = torch.mm(result, weight) - owners.append(consumed) - output.copy_(consumed) - - for _ in range(3): - forward() - torch.npu.synchronize() - owners.clear() - done = None # Core also creates per-layer events lazily during capture. - graph = None - try: - if capture: - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph): - forward() - for increment in (0, 3, 7, 11, 19): - values = initial + increment - x.copy_(values) - torch.npu.synchronize() - if graph is None: - forward() - else: - graph.replay() - torch.npu.synchronize() - product = values.float() + 1 - normalized = (product * torch.rsqrt(product.square().mean(-1, keepdim=True) + 1e-6)).bfloat16() - expected = product.bfloat16() + normalized - torch.testing.assert_close(output.cpu(), expected, atol=0.25, rtol=0) - if graph is None: - owners.clear() - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() - - @pytest.mark.parametrize("capture", [False, True]) def test_c8_bytes_reuse_existing_scatter_with_padding_and_page_gaps(capture): rows, pages, block, width = 5, 4, 128, 656 @@ -636,30 +566,6 @@ def test_plan_gather_and_written_slot_invalidation(): mapping.close() -def test_upstream_prefill_staging_plan(): - from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan - - host = allocate_mapped_host((32, 16), torch.float32) - mapping = host._hisparse_registration - host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) - table = torch.tensor([[2, 1, 0], [3, 0, 0]], dtype=torch.int32, device="npu") - lengths = torch.tensor([5, 4], dtype=torch.int32, device="npu") - try: - plan = build_hisparse_prefill_staging_plan(table, lengths, 4, 3) - assert plan.miss_mask.dtype == torch.int32 - hot = torch.full((plan.row_ids.size(1) // 4, 4, 16), -17.0, device="npu") - ops = HiSparseNPUOps(mapping) - ops.hisparse_gather_plan(host, hot, plan.row_ids, plan.dst_rows, plan.miss_mask, None, None, 0) - torch.npu.synchronize() - source = plan.row_ids.cpu()[0] - expected = torch.full((source.numel(), 16), -17.0) - expected[source >= 0] = host[source[source >= 0].long()] - torch.testing.assert_close(hot.cpu().view_as(expected), expected) - finally: - torch.npu.synchronize() - mapping.close() - - def test_invalidate_large_hot_capacity_graph(): dgi = torch.arange(20480, dtype=torch.int32, device="npu").view(2, 10240) initial = dgi.cpu() @@ -694,13 +600,13 @@ def forward(): graph.reset() -@pytest.mark.parametrize("groups", [1, 3, 21]) -def test_common_group_invalidation_changed_graph_replay(groups): +def test_common_group_invalidation_changed_graph_replay(): """Exercise the current common pointer-table caller with changing writes.""" from types import SimpleNamespace from vllm.v1.hisparse.runtime import HiSparseSlotInvalidator + groups = 3 initial = torch.arange(48, dtype=torch.int32).view(3, 16) tables = [(initial + group * 7).npu() for group in range(groups)] states = torch.tensor([2, 0, 1], dtype=torch.int32, device="npu") diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 112f529a981b..0c00d07a8581 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -225,33 +225,6 @@ def test_unsupported_indexer_formats_are_rejected_before_pool_binding(field): functions["get_hisparse_indexer_specs"](config) -@pytest.mark.parametrize("enable_caching", [False, True]) -def test_hisparse_roles_preserve_core_admission_and_explicit_hash_budget(enable_caching): - calls = [] - marker = object() - functions = definitions( - "patch/platform/patch_kv_cache_coordinator.py", - {"get_kv_cache_coordinator"}, - KVCacheGroupRole=NS(HISPARSE_SOURCE="source"), - _select_kv_token_budget=lambda maximum, inflight: maximum if inflight is None else inflight, - _orig_get_kv_cache_coordinator=lambda **kwargs: (calls.append(kwargs), marker)[1], - ) - config = NS(kv_cache_groups=[NS(role="source"), NS(role="resident"), NS(role="hot")]) - result = functions["get_kv_cache_coordinator"]( - config, - max_model_len=65536, - max_in_flight_tokens=129, - num_hash_lookahead_tokens=7, - enable_caching=enable_caching, - ) - assert result is marker and len(calls) == 1 - assert calls[0]["kv_cache_config"] is config - assert calls[0]["max_in_flight_tokens"] == 129 - assert calls[0]["num_hash_lookahead_tokens"] == 7 - assert calls[0]["enable_caching"] is enable_caching - assert "eagle_attn_layer_names" not in calls[0] - - def test_capture_metadata_prevents_warmup_host_mirrors(): class Parent: def build_for_cudagraph_capture(self, common, **kwargs): @@ -281,16 +254,6 @@ def build(self): assert submitted == [True] -def test_real_forward_mirror_errors_remain_visible(): - finish = definitions("attention/sfa_hisparse.py", {"finish_hisparse_kv_update"})["finish_hisparse_kv_update"] - - def duplicate(): - raise RuntimeError("HiSparse layer 0 mirrored twice.") - - with pytest.raises(RuntimeError, match="mirrored twice"): - finish(NS(finish_kv_update=duplicate), NS(hisparse_for_capture=False)) - - def c8_module(): path = ROOT / "vllm_ascend/ops/hisparse_c8.py" spec = importlib.util.spec_from_file_location("hisparse_c8_test_module", path) From 139fb5f53d61030112fa7ce9d75d228d09836239 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 20:18:11 +0800 Subject: [PATCH 008/120] test(hisparse): scope FP8 platform guard to C8 attention Allow the BF16 bridge test to use the native SFA registration on A2, A3 and A5. Keep only the native E4M3FN C8 case restricted to A5 and skip the unregistered 310P operation. No production or kernel changes. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_attention.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index a294e8ae4efa..8cd4ab3a25c1 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -12,13 +12,17 @@ @pytest.fixture(autouse=True) -def _require_ascend950_sfa(): - if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): - pytest.skip("HiSparse SFA integration currently targets Ascend 950") +def _require_npu_sfa(): + if not torch.npu.is_available(): + pytest.skip("HiSparse SFA tests require an NPU") + from vllm_ascend.device.device_config import is_310p + + if is_310p(): + pytest.skip("The native SFA operation is not registered for Ascend 310P") torch.npu.set_device(0) from vllm_ascend.utils import bootstrap_custom_op_env - # Load the required SFA extension without enabling all A5 custom paths. + # Load the required SFA extension without changing custom-op selection. bootstrap_custom_op_env() import vllm_ascend.vllm_ascend_C # noqa: F401 @@ -108,6 +112,10 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) def test_c8_attention_uses_existing_native_operator(strided, capture): + from vllm_ascend.device.device_config import is_950 + + if not is_950(): + pytest.skip("This case uses the native A5 E4M3FN C8 payload; A2/A3 use an INT8 format") from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope from vllm_ascend.ops.hisparse_c8 import c8_physical_attention, pack_c8_components From 45240271a4195e0eb18651deb5bf7f059d39b5df Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 20:38:57 +0800 Subject: [PATCH 009/120] test(hisparse): import production modules directly in UT Remove AST extraction, dynamic exec/file loading and replacement base/spec classes. Exercise the actual backend inheritance, Core/Ascend cache specs and metadata builder through normal imports and the existing CPU-UT conftest. Keep the same twenty boundary cases. Standard pytest was attempted but collection is blocked by missing local vllm dependencies; no source-isolated pass is claimed for this version. Focused syntax, Ruff, Markdown and diff checks pass. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 9 +- .../test_hisparse_foundation_cpu.py | 227 ++++++------------ 2 files changed, 85 insertions(+), 151 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 2d677c897030..556f1385fcc7 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -118,10 +118,10 @@ A same-shaped batch is not permission to reuse old lengths or request mappings. ## Validation -Portable, source-isolated CPU boundary checks: +CPU boundary checks in the normal installed development environment: ```bash -python -m pytest -q --noconftest \ +python -m pytest -q \ tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py ``` @@ -130,7 +130,10 @@ isolation, strided indexer aliases, planner alignment, unsupported formats, and capture-only mirror suppression. C8 checks also cover opaque component bytes, nonzero slab offsets, stable duplicate/padding ordering, Indexer key/scale writes and planner alignment before allocation. -They do not exercise an installed NPU package or qualify device execution. +The tests directly import the production modules and Core/Ascend cache specs, +using the repository's normal CPU-UT conftest. Install the matching Core and +Ascend test dependencies before running them. CPU checks do not qualify device +execution. The weight-loading wrapper keeps Ascend main's explicit activation-dtype ABI, already covered by `tests/ut/ops/test_mla_weight_processing.py`; this foundation diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 0c00d07a8581..b7c27c0740a0 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -1,38 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Source-isolated CPU boundary checks; these do not qualify NPU execution.""" +"""CPU boundary checks using the installed Core and Ascend implementations.""" -import ast -import ctypes -import importlib.util import math -import sys -from dataclasses import dataclass, replace -from pathlib import Path +from dataclasses import replace from types import SimpleNamespace as NS import pytest import torch - -ROOT = Path(__file__).resolve().parents[4] - - -def definitions(relative, names, **bindings): - path = ROOT / "vllm_ascend" / relative - tree = ast.parse(path.read_bytes()) - nodes = [n for n in tree.body if isinstance(n, (ast.ClassDef, ast.FunctionDef)) and n.name in names] - assert len(nodes) == len(names) - future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) - namespace = {"torch": torch, **bindings} - module = ast.fix_missing_locations(ast.Module(body=[future, *nodes], type_ignores=[])) - exec(compile(module, str(path), "exec"), namespace) - return namespace - - -def backend_type(): - return definitions( - "distributed/kv_transfer/hisparse/backend.py", {"AscendHiSparseBackend"}, HiSparseBackend=object - )["AscendHiSparseBackend"] +from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec + +from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder, finish_hisparse_kv_update +from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder +from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec +from vllm_ascend.distributed.kv_transfer.hisparse.backend import AscendHiSparseBackend +from vllm_ascend.distributed.kv_transfer.hisparse.binding import ( + _bind_indexer_views, + align_hisparse_indexer_specs, + get_hisparse_indexer_specs, + write_hisparse_indexer_c8, +) +from vllm_ascend.ops import hisparse as hisparse_ops +from vllm_ascend.ops.hisparse_c8 import ( + c8_storage_view, + compact_c8_indices, + map_c8_physical_indices, + pack_c8_components, +) def test_failed_unregister_retains_host_owner_for_safe_retry(monkeypatch): @@ -48,8 +42,8 @@ def close(self): if len(events) == 1: raise RuntimeError("driver did not retire registration") - monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(MappedHostTensor=Registration)) - backend = backend_type()() + monkeypatch.setattr(hisparse_ops, "MappedHostTensor", Registration) + backend = AscendHiSparseBackend() backend.pin_tensor(tensor) registration = tensor._hisparse_registration assert registration.tensor is tensor @@ -62,9 +56,8 @@ def close(self): def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): - cls = definitions("ops/hisparse.py", {"MappedHostTensor"}, ctypes=ctypes)["MappedHostTensor"] backing = torch.arange(128, dtype=torch.uint8) - mapping = cls.__new__(cls) + mapping = hisparse_ops.MappedHostTensor.__new__(hisparse_ops.MappedHostTensor) mapping.tensor = backing mapping.host_address = backing.data_ptr() mapping.nbytes = backing.numel() @@ -81,12 +74,8 @@ def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): @pytest.mark.parametrize("max_rows,hot_rows,union_rows", [(1, 16, 8), (9, 48, 24), (37, 128, 72)]) -def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( - monkeypatch, max_rows, hot_rows, union_rows -): - operations = definitions("ops/hisparse.py", {"ResolverWorkspace", "HiSparseNPUOps"}) - monkeypatch.setitem(sys.modules, "vllm_ascend.ops.hisparse", NS(**operations)) - backend = backend_type()() +def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups(max_rows, hot_rows, union_rows): + backend = AscendHiSparseBackend() def group(): owner = NS() @@ -118,61 +107,22 @@ def group(): assert torch.count_nonzero(second_workspace.row_keys) == 0 -@dataclass(frozen=True) -class IndexerSpec: - block_size: int = 4 - num_kv_heads: int = 1 - head_size: int = 4 - block_stride_alignment: int = 3 - scale_dim: int = 0 - cache_sparse_li_c8: bool = False - cache_sparse_li_c4: bool = False - dtype: torch.dtype = torch.float32 - scale_dtype: torch.dtype = torch.float32 - sfa_dcp_replicated_indexer_size: int = 1 - page_size_padded: int | None = None - - @property - def page_size_bytes(self): - return ( - self.block_size - * self.num_kv_heads - * (self.head_size * self.dtype.itemsize + self.scale_dim * self.scale_dtype.itemsize) - ) - - def get_block_stride_alignment(self): - return self.block_stride_alignment - - -def binding_functions(**bindings): - return definitions( - "distributed/kv_transfer/hisparse/binding.py", - { - "align_hisparse_indexer_specs", - "get_hisparse_indexer_specs", - "_bind_indexer_views", - "_validate_indexer_spec", - "_quantized_indexer_views", - "write_hisparse_indexer_c8", - }, - AscendSFAIndexerCacheSpec=IndexerSpec, - UniformTypeKVCacheSpecs=type("Uniform", (), {}), - math=math, - replace=replace, - **bindings, +@pytest.fixture +def indexer_spec(): + return AscendSFAIndexerCacheSpec( + block_size=4, num_kv_heads=1, head_size=4, dtype=torch.float32, block_stride_alignment=3 ) @pytest.mark.parametrize("blocks,stride", [(1, 1), (3, 2), (5, 3)]) -def test_expanded_indexer_alias_preserves_shared_slab_offsets(blocks, stride): - functions = binding_functions() - spec = IndexerSpec() +def test_expanded_indexer_alias_preserves_shared_slab_offsets(indexer_spec, blocks, stride): + spec = indexer_spec page = 16 backing = torch.arange(blocks * stride * page + page, dtype=torch.float32) original = backing.as_strided((blocks, 1, 4, 4), (stride * page, page, 4, 1), storage_offset=page) layer = NS() caches = {"indexer": original} - functions["_bind_indexer_views"](caches, {"indexer": spec}, {"indexer": layer}) + _bind_indexer_views(caches, {"indexer": spec}, {"indexer": layer}) view = caches["indexer"][0] assert view.data_ptr() == original.data_ptr() assert view.storage_offset() == original.storage_offset() @@ -184,108 +134,83 @@ def test_expanded_indexer_alias_preserves_shared_slab_offsets(blocks, stride): assert original[-1, 0, 1, 2] == -73 -def test_alignment_updates_only_indexer_specs_without_changing_inputs(): - functions = binding_functions() - source, indexer = NS(get_block_stride_alignment=lambda: 1), IndexerSpec() +def test_alignment_updates_only_indexer_specs_without_changing_inputs(indexer_spec): + source = MLAAttentionSpec(block_size=4, num_kv_heads=1, head_size=4, dtype=torch.float32) + indexer = indexer_spec values = {"source": source, "indexer": indexer} - result = functions["align_hisparse_indexer_specs"](values) + result = align_hisparse_indexer_specs(values) assert result["source"] is source assert result["indexer"].get_block_stride_alignment() == math.lcm(3, indexer.page_size_bytes) assert values["indexer"] is indexer and indexer.block_stride_alignment == 3 -def test_c8_page_alignment_reserves_scale_bytes_and_whole_main_rows(): - @dataclass(frozen=True) - class StorageSpec: - block_stride_alignment: int = 128 - - def get_block_stride_alignment(self): - return self.block_stride_alignment - - functions = binding_functions(HiSparseResidentSpec=StorageSpec, HiSparseHotSpec=StorageSpec) - storage = StorageSpec() - indexer = IndexerSpec( - block_size=128, head_size=128, dtype=torch.float8_e4m3fn, scale_dim=1, cache_sparse_li_c8=True - ) - result = functions["align_hisparse_indexer_specs"]( - {"resident": storage, "hot": storage, "indexer": indexer}, packed_row_bytes=656 +def test_c8_page_alignment_reserves_scale_bytes_and_whole_main_rows(indexer_spec): + resident = HiSparseResidentSpec(block_size=128, page_size=128 * 656, block_stride_alignment=128) + hot = HiSparseHotSpec(block_size=128, page_size=128 * 656, blocks_per_request=1, block_stride_alignment=128) + indexer = replace( + indexer_spec, + block_size=128, + head_size=128, + dtype=torch.float8_e4m3fn, + scale_dim=1, + scale_dtype=torch.float32, + cache_sparse_li_c8=True, ) + result = align_hisparse_indexer_specs({"resident": resident, "hot": hot, "indexer": indexer}, packed_row_bytes=656) stride_alignment = math.lcm(*(spec.get_block_stride_alignment() for spec in result.values())) assert stride_alignment % 656 == 0 and stride_alignment % 128 == 0 assert result["indexer"].page_size_padded == 128 * (128 + 4) - assert storage.block_stride_alignment == 128 and indexer.page_size_padded is None + assert resident.block_stride_alignment == hot.block_stride_alignment == 128 + assert indexer.page_size_padded is None @pytest.mark.parametrize("field", ["scale_dim", "cache_sparse_li_c8", "cache_sparse_li_c4"]) -def test_unsupported_indexer_formats_are_rejected_before_pool_binding(field): - functions = binding_functions() - spec = replace(IndexerSpec(), **{field: 1}) +def test_unsupported_indexer_formats_are_rejected_before_pool_binding(indexer_spec, field): + spec = replace(indexer_spec, **{field: 1}) config = NS(kv_cache_groups=[NS(layer_names=["indexer"], kv_cache_spec=spec)]) with pytest.raises(ValueError, match="C4|C8|scales"): - functions["get_hisparse_indexer_specs"](config) - - -def test_capture_metadata_prevents_warmup_host_mirrors(): - class Parent: - def build_for_cudagraph_capture(self, common, **kwargs): - assert common is common_input and kwargs["draft_index"] is draft - return self.build() + get_hisparse_indexer_specs(config) - def build(self): - return NS(hisparse_for_capture=False) - methods = definitions( - "attention/sfa_hisparse.py", - {"AscendHiSparseMetadataBuilder", "finish_hisparse_kv_update"}, - AscendSFAMetadataBuilder=Parent, +def test_capture_metadata_prevents_warmup_host_mirrors(monkeypatch): + capture_metadata = NS(hisparse_for_capture=False) + monkeypatch.setattr( + AscendSFAMetadataBuilder, "build_for_cudagraph_capture", lambda *args, **kwargs: capture_metadata ) - cls = methods["AscendHiSparseMetadataBuilder"] - builder = cls.__new__(cls) - common_input, draft = object(), object() - capture_metadata = builder.build_for_cudagraph_capture(common_input, draft_index=draft) - real_metadata = builder.build() - finish = methods["finish_hisparse_kv_update"] + builder = AscendHiSparseMetadataBuilder.__new__(AscendHiSparseMetadataBuilder) + builder.build_for_cudagraph_capture(NS()) + real_metadata = NS(hisparse_for_capture=False) submitted = [] cache = NS(finish_kv_update=lambda: submitted.append(True)) for _ in range(5): - finish(cache, capture_metadata) + finish_hisparse_kv_update(cache, capture_metadata) assert submitted == [] - finish(cache, real_metadata) + finish_hisparse_kv_update(cache, real_metadata) assert submitted == [True] -def c8_module(): - path = ROOT / "vllm_ascend/ops/hisparse_c8.py" - spec = importlib.util.spec_from_file_location("hisparse_c8_test_module", path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - def test_c8_packing_preserves_opaque_rope_scale_and_nan_payloads(): - module = c8_module() rows = 3 raw = [torch.arange(rows * width).to(torch.uint8).reshape(rows, width) for width in (512, 128, 16)] components = [values.view(torch.float8_e4m3fn) for values in raw] - packed = module.pack_c8_components(*components) + packed = pack_c8_components(*components) assert packed.dtype == torch.uint8 and packed.shape == (rows, 656) torch.testing.assert_close(packed, torch.cat(raw, dim=-1), rtol=0, atol=0) @pytest.mark.parametrize("stride_rows", [128, 136]) def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): - module = c8_module() pages, block, width, offset = 3, 128, 656, 1312 backing = torch.full((offset + pages * stride_rows * width,), 17, dtype=torch.uint8) cache = backing.view(torch.float8_e4m3fn).as_strided( (pages, block, width), (stride_rows * width, width, 1), storage_offset=offset ) - view, actual_block, actual_stride, capacity = module.c8_storage_view(cache) + view, actual_block, actual_stride, capacity = c8_storage_view(cache) assert (actual_block, actual_stride, capacity) == (block, stride_rows, pages * block) assert view.data_ptr() == cache.data_ptr() assert view.untyped_storage().data_ptr() == backing.data_ptr() indices = torch.tensor([[0, 127, 128, 255, 256, 383, -1, 384]], dtype=torch.int32) - mapped = module.map_c8_physical_indices(indices, actual_block, actual_stride, capacity) + mapped = map_c8_physical_indices(indices, actual_block, actual_stride, capacity) expected = torch.tensor( [[0, 127, stride_rows, stride_rows + 127, 2 * stride_rows, 2 * stride_rows + 127, -1, -1]], dtype=torch.int32, @@ -296,16 +221,22 @@ def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): def test_c8_compaction_preserves_duplicate_order_and_empty_rows(): - module = c8_module() indices = torch.tensor([[7, -1, 2, 7, -1, 0], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) expected = torch.tensor([[7, 2, 7, 0, -1, -1], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) - torch.testing.assert_close(module.compact_c8_indices(indices), expected, rtol=0, atol=0) + torch.testing.assert_close(compact_c8_indices(indices), expected, rtol=0, atol=0) @pytest.mark.parametrize("padding", [0, 256]) -def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(padding): - functions = binding_functions() - spec = IndexerSpec(block_size=128, head_size=128, dtype=torch.float8_e4m3fn, scale_dim=1, cache_sparse_li_c8=True) +def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(indexer_spec, padding): + spec = replace( + indexer_spec, + block_size=128, + head_size=128, + dtype=torch.float8_e4m3fn, + scale_dim=1, + scale_dtype=torch.float32, + cache_sparse_li_c8=True, + ) pages, offset = 3, 256 key_bytes = 128 * 128 stride = spec.page_size_bytes + padding @@ -315,7 +246,7 @@ def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(padding): ) layer = NS() caches = {"indexer": cache} - functions["_bind_indexer_views"](caches, {"indexer": spec}, {"indexer": layer}) + _bind_indexer_views(caches, {"indexer": spec}, {"indexer": layer}) key, scale = caches["indexer"] assert key.untyped_storage().data_ptr() == scale.untyped_storage().data_ptr() == raw.data_ptr() assert layer.hisparse_block_stride == 1 @@ -340,5 +271,5 @@ def scatter(destination, indexes, source): expected[start : start + 128] = values_raw[i] start = offset + page * stride + key_bytes + row * 4 expected[start : start + 4] = scale_bytes[i] - functions["write_hisparse_indexer_c8"](key, scale, values, scales, slots, scatter) + write_hisparse_indexer_c8(key, scale, values, scales, slots, scatter) torch.testing.assert_close(raw, expected, rtol=0, atol=0) From fa9704aa39c286683cffb135decfd87df96d85af Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 21:24:59 +0800 Subject: [PATCH 010/120] refactor(hisparse): select Indexer writes through the backend Use a HiSparse Indexer subclass and the existing write_cache override. Restore ordinary Indexer write/build behavior and remove payload-format attributes and string routing. Reuse the cache backend choice in the model wrapper while retaining model-specific factories. Keep the existing C8 stride helpers and native kernels unchanged. Update existing writer/wrapper tests; standard UT collection remains blocked by missing local vllm. Focused checks pass; full lint failed installing actionlint after a remote disconnect. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 7 +++ .../one_card/test_hisparse_ops.py | 8 ++-- .../test_hisparse_foundation_cpu.py | 14 ++++-- tests/ut/ops/test_mla.py | 5 ++- vllm_ascend/attention/indexer.py | 44 ++++++++++++------- .../kv_transfer/hisparse/binding.py | 3 -- vllm_ascend/ops/mla.py | 5 ++- vllm_ascend/patch/__init__.py | 3 +- vllm_ascend/patch/worker/patch_attn_utils.py | 5 ++- 9 files changed, 60 insertions(+), 34 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 556f1385fcc7..abccd2b0887a 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -87,6 +87,13 @@ Use existing interfaces, rather than adding unused placeholder hooks: | `validate_kv_cache_config` / `prepare_kv_cache_views` | Core validates BF16/C8 before allocation and binds views inside rollback scope | Additional payload formats with layout tests | | Backend `pin_tensor` / `unpin_tensor` / `swap_blocks_batch` | Core owns completion, explicit registration release and transfer lifetime | DMA implementation and supported memory capabilities | | `AscendHiSparseMetadataBuilder` / `AscendHiSparseImpl` | SFA factories select metadata and the BF16/C8 layout bridge | A smaller layout-conversion cost | +| `AscendHiSparseIndexerBackend.write_cache` | The existing cache backend interface selects HiSparse once; its override handles C8 page strides | Other Indexer payload layouts | + +The model's Indexer wrapper reuses the cache layer's `get_attn_backend()` choice, +preserving model-specific backend factories. HiSparse owns its metadata builder +and C8 write override. The ordinary Indexer write path remains unchanged, and +BF16 HiSparse writes reuse it after metadata remaps the slots. Cache binding +creates the key/scale views; it does not attach payload-format routing strings. `ops/hisparse.py` owns operation validation and persistent per-group workspace; `ops/triton/hisparse.py` implements the reference NPU path. Followers share their diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 8d3374239e80..5c35f9c8e622 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -649,7 +649,7 @@ def forward(): def test_page_strided_indexer_write_preserves_bits_and_padding(capture): from types import SimpleNamespace - from vllm_ascend.attention.indexer import AscendSFAIndexerBackend + from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.distributed.kv_transfer.hisparse.binding import _quantized_indexer_views @@ -677,11 +677,9 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): key, scale = _quantized_indexer_views(packed, spec) with pytest.raises(RuntimeError, match="view size"): key.view(-1, width) - writer = AscendSFAIndexerBackend.__new__(AscendSFAIndexerBackend) + writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) torch.nn.Module.__init__(writer) - writer.k_cache = SimpleNamespace( - kv_cache=(key, scale), hisparse_indexer_payload_format="ascend_indexer_e4m3fn_key_fp32_scale_v1" - ) + writer.k_cache = SimpleNamespace(kv_cache=(key, scale)) writer.enable_sparse_li_c8 = True writer.enable_sparse_li_c4 = False values = torch.zeros((5, width), dtype=torch.uint8, device="npu").view(torch.float8_e4m3fn) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index b7c27c0740a0..ce1383f4c1e1 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -10,15 +10,16 @@ import torch from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec +from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder, finish_hisparse_kv_update from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec +from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.distributed.kv_transfer.hisparse.backend import AscendHiSparseBackend from vllm_ascend.distributed.kv_transfer.hisparse.binding import ( _bind_indexer_views, align_hisparse_indexer_specs, get_hisparse_indexer_specs, - write_hisparse_indexer_c8, ) from vllm_ascend.ops import hisparse as hisparse_ops from vllm_ascend.ops.hisparse_c8 import ( @@ -227,7 +228,7 @@ def test_c8_compaction_preserves_duplicate_order_and_empty_rows(): @pytest.mark.parametrize("padding", [0, 256]) -def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(indexer_spec, padding): +def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(monkeypatch, indexer_spec, padding): spec = replace( indexer_spec, block_size=128, @@ -250,7 +251,6 @@ def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(indexer_spe key, scale = caches["indexer"] assert key.untyped_storage().data_ptr() == scale.untyped_storage().data_ptr() == raw.data_ptr() assert layer.hisparse_block_stride == 1 - assert layer.hisparse_indexer_payload_format == "ascend_indexer_e4m3fn_key_fp32_scale_v1" slots = torch.tensor([127, 0, 128, 383, -1, 256], dtype=torch.int32) values_raw = torch.arange(6 * 128).to(torch.uint8).reshape(6, 128) values = values_raw.view(torch.float8_e4m3fn) @@ -261,6 +261,12 @@ def scatter(destination, indexes, source): valid = indexes >= 0 destination[indexes[valid].long()] = source[valid] + monkeypatch.setattr(DeviceOperator, "scatter_cache", scatter) + writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) + torch.nn.Module.__init__(writer) + writer.k_cache = NS(kv_cache=caches["indexer"]) + writer.enable_sparse_li_c8 = True + writer.enable_sparse_li_c4 = False expected = raw.clone() scale_bytes = scales.view(torch.uint8).reshape(6, 4) for i, slot in enumerate(slots.tolist()): @@ -271,5 +277,5 @@ def scatter(destination, indexes, source): expected[start : start + 128] = values_raw[i] start = offset + page * stride + key_bytes + row * 4 expected[start : start + 4] = scale_bytes[i] - write_hisparse_indexer_c8(key, scale, values, scales, slots, scatter) + writer.write_cache(values, scales, slots) torch.testing.assert_close(raw, expected, rtol=0, atol=0) diff --git a/tests/ut/ops/test_mla.py b/tests/ut/ops/test_mla.py index 0ef11cb2cc2a..87300e451052 100644 --- a/tests/ut/ops/test_mla.py +++ b/tests/ut/ops/test_mla.py @@ -318,11 +318,12 @@ def test_write_cache_without_li_c8_writes_k_only(self, mock_scatter): class TestIndexerWrapper(TestBase): - @patch("vllm_ascend.ops.mla.AscendSFAIndexerBackend") - def test_constructs_backend_and_delegates_sfa_interface(self, mock_backend_cls): + def test_constructs_backend_and_delegates_sfa_interface(self): vllm_indexer = MagicMock(name="vllm_indexer") + mock_backend_cls = vllm_indexer.k_cache.get_attn_backend.return_value wrapper = IndexerWrapper(vllm_indexer, qk_rope_head_dim=64) + vllm_indexer.k_cache.get_attn_backend.assert_called_once_with() mock_backend_cls.assert_called_once_with(vllm_indexer, 64) self.assertIs(wrapper.impl, mock_backend_cls.return_value) self.assertIs(wrapper.k_cache, wrapper.impl.k_cache) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 34cf99fd8521..9285c92af657 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -154,8 +154,6 @@ def get_name() -> str: @staticmethod def get_builder_cls(): - if get_current_vllm_config().attention_config.hisparse_config is not None: - return AscendHiSparseIndexerMetadataBuilder return AscendSFAIndexerMetadataBuilder @staticmethod @@ -286,20 +284,6 @@ def write_cache( with a different cache layout should override it. """ indexer_k_cache = self.k_cache.kv_cache[INDEXER_K_CACHE_SLOT] - if getattr(self.k_cache, "hisparse_indexer_payload_format", None) == "ascend_indexer_e4m3fn_key_fp32_scale_v1": - if not self.enable_sparse_li_c8 or self.enable_sparse_li_c4 or k_li_scale is None: - raise ValueError("HiSparse's page-strided Indexer requires native C8 values and FP32 scales") - from vllm_ascend.distributed.kv_transfer.hisparse.binding import write_hisparse_indexer_c8 - - write_hisparse_indexer_c8( - indexer_k_cache, - self.k_cache.kv_cache[INDEXER_SCALE_CACHE_SLOT], - k_li, - k_li_scale, - slot_mapping.reshape(-1), - DeviceOperator.scatter_cache, - ) - return # float4_e2m1fn_x2 / float8_e8m0fnu are not supported by the scatter # kernel; view as uint8 (same 1-byte layout) for the cache write. if k_li.dtype == torch_npu.float4_e2m1fn_x2: @@ -538,6 +522,34 @@ def forward( ) +class AscendHiSparseIndexerBackend(AscendSFAIndexerBackend): + """Keep HiSparse shared-slab cache writes behind the Indexer interface.""" + + @staticmethod + def get_name() -> str: + return "ASCEND_HISPARSE_INDEXER" + + @staticmethod + def get_builder_cls(): + return AscendHiSparseIndexerMetadataBuilder + + def write_cache(self, k_li, k_li_scale, slot_mapping) -> None: + if not self.enable_sparse_li_c8: + return super().write_cache(k_li, k_li_scale, slot_mapping) + if k_li_scale is None: + raise ValueError("HiSparse Indexer C8 requires FP32 scales") + from vllm_ascend.distributed.kv_transfer.hisparse.binding import write_hisparse_indexer_c8 + + write_hisparse_indexer_c8( + self.k_cache.kv_cache[INDEXER_K_CACHE_SLOT], + self.k_cache.kv_cache[INDEXER_SCALE_CACHE_SLOT], + k_li, + k_li_scale, + slot_mapping.reshape(-1), + DeviceOperator.scatter_cache, + ) + + class AscendSFAIndexerMetadataBuilder(AttentionMetadataBuilder[AscendSFAIndexerMetadata]): """Builds the metadata consumed by SFA indexer forwards. diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py index 014b21916321..175c0201a813 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py +++ b/vllm_ascend/distributed/kv_transfer/hisparse/binding.py @@ -161,8 +161,6 @@ def _bind_indexer_views( if spec.cache_sparse_li_c8: kv_caches[name] = _quantized_indexer_views(cache, spec) forward_context[name].hisparse_block_stride = 1 - # Declare only after validating this allocation's actual K/scale views. - forward_context[name].hisparse_indexer_payload_format = "ascend_indexer_e4m3fn_key_fp32_scale_v1" continue block_size = cache.shape[2] page_elements = block_size * spec.num_kv_heads * spec.head_size @@ -176,5 +174,4 @@ def _bind_indexer_views( stride=(page_elements, spec.num_kv_heads * spec.head_size, spec.head_size, 1), ) forward_context[name].hisparse_block_stride = block_stride - forward_context[name].hisparse_indexer_payload_format = "bf16" kv_caches[name] = (indexer_cache,) diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index 402aa56ad3a3..03707a0d7865 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -31,7 +31,6 @@ from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata # type: ignore -from vllm_ascend.attention.indexer import AscendSFAIndexerBackend from vllm_ascend.attention.mla_v1 import AscendMLAAttention from vllm_ascend.attention.utils import mark_fused_preprocess_weights @@ -66,7 +65,9 @@ def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: self.register_parameter(name, parameter) backend_factory = getattr(type(vllm_indexer), "get_ascend_indexer_backend_cls", None) - backend_cls = backend_factory(vllm_indexer) if backend_factory is not None else AscendSFAIndexerBackend + backend_cls = ( + backend_factory(vllm_indexer) if backend_factory is not None else vllm_indexer.k_cache.get_attn_backend() + ) self.impl = backend_cls(vllm_indexer, qk_rope_head_dim) # Interface consumed by the SFA impl - delegated to the backend impl. diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 52d40b46ac7e..12113a662d21 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1528,7 +1528,8 @@ # the upstream indexer backend is GPU-specific. # How： # Route SFA indexer cache layers to the existing -# `AscendSFAIndexerBackend`. +# `AscendSFAIndexerBackend`, or `AscendHiSparseIndexerBackend` when +# HiSparse is enabled. The model wrapper uses the cache's backend choice. # Related PR (if no, explain why): # https://github.com/vllm-project/vllm-ascend/pull/13069 # Future Plan: diff --git a/vllm_ascend/patch/worker/patch_attn_utils.py b/vllm_ascend/patch/worker/patch_attn_utils.py index fed8ac8d504a..77c021560c54 100644 --- a/vllm_ascend/patch/worker/patch_attn_utils.py +++ b/vllm_ascend/patch/worker/patch_attn_utils.py @@ -4,10 +4,11 @@ import vllm import vllm.v1.worker.utils as v1_worker_utils from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, AttentionInterface +from vllm.config import get_current_vllm_config from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 -from vllm_ascend.attention.indexer import AscendSFAIndexerBackend +from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend, AscendSFAIndexerBackend from vllm_ascend.patch.worker.patch_bind_kv_cache import ( bind_kv_cache, bind_kv_cache_to_layers, @@ -23,6 +24,8 @@ def _get_ascend_sfa_indexer_backend(_self): + if get_current_vllm_config().attention_config.hisparse_config is not None: + return AscendHiSparseIndexerBackend return AscendSFAIndexerBackend From b22942dc9916d766bbba6bc800d937ac8e5a0e33 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 22:02:36 +0800 Subject: [PATCH 011/120] refactor(hisparse): group shared cache code and own Attention dispatch Place backend, binding, cache operations and C8 layout helpers in vllm_ascend/hisparse. Move the C8 native Attention invocation into AscendHiSparseImpl, which already owns its length workspace. Update the platform path and all consumers without compatibility forwarding modules. Keep data transformations, allocator behavior, kernel arguments and native/Triton kernels unchanged. Existing C8 device coverage now enters the Attention backend. Packaging, syntax, focused lint and structural equivalence checks pass; standard UT and full lint remain blocked by local dependencies/tool installation. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 20 +++++- .../one_card/test_hisparse_attention.py | 10 ++- .../one_card/test_hisparse_ops.py | 4 +- .../one_card/test_hisparse_prefill_staging.py | 4 +- .../test_hisparse_foundation_cpu.py | 8 +-- vllm_ascend/attention/indexer.py | 2 +- vllm_ascend/attention/sfa_hisparse.py | 69 ++++++++++++++++--- .../kv_transfer => }/hisparse/__init__.py | 2 +- .../kv_transfer => }/hisparse/backend.py | 8 +-- .../kv_transfer => }/hisparse/binding.py | 2 +- .../{ops/hisparse_c8.py => hisparse/c8.py} | 56 +-------------- .../hisparse.py => hisparse/cache_ops.py} | 2 +- vllm_ascend/platform.py | 2 +- vllm_ascend/worker/v2/attn_utils.py | 6 +- 14 files changed, 109 insertions(+), 86 deletions(-) rename vllm_ascend/{distributed/kv_transfer => }/hisparse/__init__.py (59%) rename vllm_ascend/{distributed/kv_transfer => }/hisparse/backend.py (93%) rename vllm_ascend/{distributed/kv_transfer => }/hisparse/binding.py (99%) rename vllm_ascend/{ops/hisparse_c8.py => hisparse/c8.py} (62%) rename vllm_ascend/{ops/hisparse.py => hisparse/cache_ops.py} (99%) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index abccd2b0887a..92ecea2fa2c9 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -95,11 +95,29 @@ and C8 write override. The ordinary Indexer write path remains unchanged, and BF16 HiSparse writes reuse it after metadata remaps the slots. Cache binding creates the key/scale views; it does not attach payload-format routing strings. -`ops/hisparse.py` owns operation validation and persistent per-group workspace; +`hisparse/cache_ops.py` owns operation validation and persistent per-group workspace; `ops/triton/hisparse.py` implements the reference NPU path. Followers share their leader's scratch only on that group's ordered stream. Independent groups and current request/state mappings remain separate. +Code ownership follows these boundaries: + +```text +vllm_ascend/ + attention/sfa_hisparse.py # Attention dispatch, C8 adaptation and native calls + hisparse/ + backend.py # Core backend interface, registration and DMA + binding.py # Shared-slab specs and key/scale views + cache_ops.py # Residency, refill, invalidation and byte copies + c8.py # C8 packing, views and physical-index helpers + ops/triton/hisparse.py # Cache-operation kernels; no Attention arithmetic +``` + +The common package contains no Attention invocation. The HiSparse Attention +implementation owns query preparation, current length tensors and the call to +the existing native SFA operation. The Triton and native Attention kernels are +unchanged by this organization. + Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces are used directly. The Ascend backend preserves Core's event placement and does not replace its Torch bindings. This PR tests its own NPU operations and changing diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 8cd4ab3a25c1..6c763dc66058 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -116,8 +116,9 @@ def test_c8_attention_uses_existing_native_operator(strided, capture): if not is_950(): pytest.skip("This case uses the native A5 E4M3FN C8 payload; A2/A3 use an INT8 format") + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope - from vllm_ascend.ops.hisparse_c8 import c8_physical_attention, pack_c8_components + from vllm_ascend.hisparse.c8 import pack_c8_components torch.manual_seed(71) rows, heads, blocks, block, topk = 3, 8, 2, 128, 128 @@ -142,9 +143,14 @@ def test_c8_attention_uses_existing_native_operator(strided, capture): indices[-1] = -1 selected = indices.to(device) query_end, key_length = (torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl.enable_sparse_sfa_c8 = True + impl.kv_lora_rank = 512 + impl.scale = 1 / math.sqrt(576) + impl._c8_length_workspace = lambda: (query_end, key_length) def forward(): - return c8_physical_attention(q, qr, cache, selected, 1 / math.sqrt(576), query_end, key_length) + return impl._attention_physical(q, qr, cache, selected) graph = None try: diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 5c35f9c8e622..72c6fa5c7717 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -8,7 +8,7 @@ import torch import torch_npu # noqa: F401 -from vllm_ascend.ops.hisparse import ( +from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, MappedHostTensor, ResolverWorkspace, @@ -651,7 +651,7 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec - from vllm_ascend.distributed.kv_transfer.hisparse.binding import _quantized_indexer_views + from vllm_ascend.hisparse.binding import _quantized_indexer_views if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("The native E4M3FN HiSparse format currently targets A5") diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 16268196e9ac..8975b88a2130 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -9,8 +9,8 @@ import torch_npu # noqa: F401 from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView, build_hisparse_prefill_staging_plan -from vllm_ascend.distributed.kv_transfer.hisparse.backend import AscendHiSparseBackend -from vllm_ascend.ops.hisparse import HiSparseNPUOps, allocate_mapped_host, scatter_packed_kv +from vllm_ascend.hisparse.backend import AscendHiSparseBackend +from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, allocate_mapped_host, scatter_packed_kv @pytest.mark.parametrize("resident_block_size", [128, 64, 32]) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index ce1383f4c1e1..f100adc46eec 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -15,14 +15,14 @@ from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.device.device_op import DeviceOperator -from vllm_ascend.distributed.kv_transfer.hisparse.backend import AscendHiSparseBackend -from vllm_ascend.distributed.kv_transfer.hisparse.binding import ( +from vllm_ascend.hisparse import cache_ops as hisparse_ops +from vllm_ascend.hisparse.backend import AscendHiSparseBackend +from vllm_ascend.hisparse.binding import ( _bind_indexer_views, align_hisparse_indexer_specs, get_hisparse_indexer_specs, ) -from vllm_ascend.ops import hisparse as hisparse_ops -from vllm_ascend.ops.hisparse_c8 import ( +from vllm_ascend.hisparse.c8 import ( c8_storage_view, compact_c8_indices, map_c8_physical_indices, diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 9285c92af657..33493544b70e 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -538,7 +538,7 @@ def write_cache(self, k_li, k_li_scale, slot_mapping) -> None: return super().write_cache(k_li, k_li_scale, slot_mapping) if k_li_scale is None: raise ValueError("HiSparse Indexer C8 requires FP32 scales") - from vllm_ascend.distributed.kv_transfer.hisparse.binding import write_hisparse_indexer_c8 + from vllm_ascend.hisparse.binding import write_hisparse_indexer_c8 write_hisparse_indexer_c8( self.k_cache.kv_cache[INDEXER_K_CACHE_SLOT], diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index cf7844d0268a..368bd35ba71f 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Ascend SFA execution over upstream HiSparse's resident and hot cache pool.""" +import math from dataclasses import dataclass import torch @@ -109,7 +110,7 @@ def __init__(self, *args, **kwargs): raise ValueError("Ascend HiSparse requires a model sparse-index group") builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) if self.enable_sparse_sfa_c8: - from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + from vllm_ascend.hisparse.c8 import C8_ROW_BYTES row_width = C8_ROW_BYTES else: @@ -200,12 +201,12 @@ def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=Non return super().forward(layer_name, hidden_states, kv_cache, attn_metadata, output) def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): - from vllm_ascend.ops.hisparse import scatter_packed_kv + from vllm_ascend.hisparse.cache_ops import scatter_packed_kv assert self.kv_a_layernorm is not None if self.enable_sparse_sfa_c8: from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope - from vllm_ascend.ops.hisparse_c8 import pack_c8_components + from vllm_ascend.hisparse.c8 import pack_c8_components values = kv_no_split.reshape(-1, self.num_kv_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim) rope, latent, scales = custom_kv_rmsnorm_rope( @@ -273,13 +274,65 @@ def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_s # which would duplicate the write and can overwrite unrelated slab rows. return k_pe, k_nope + def _c8_physical_attention(self, query, query_rope, cache, indices): + """Adapt native packed rows to the existing quantized SFA operation. + + Keep compaction stable, including duplicates. The two current lengths are + device writes captured on the consumer stream; no Host read or stream/event + policy is introduced here. + """ + from vllm_ascend.hisparse.c8 import c8_storage_view, compact_c8_indices, map_c8_physical_indices + + query_ends, key_length = self._c8_length_workspace() + scale = self.scale + view, block, stride_rows, capacity = c8_storage_view(cache) + rows = query.shape[0] + if query.ndim != 3 or query.dtype != torch.bfloat16 or query.shape[-1] != self.kv_lora_rank or rows <= 0: + raise ValueError("HiSparse C8 query must be nonempty BF16 [queries, heads, 512]") + if query_rope.shape != (*query.shape[:-1], 64) or query_rope.dtype != query.dtype: + raise ValueError("HiSparse C8 requires matching BF16 query RoPE") + if indices.ndim != 2 or indices.shape[0] != rows or indices.shape[1] <= 0 or indices.shape[1] % 128: + raise ValueError("HiSparse C8 top-k width must be positive and divisible by 128") + if query.device.type != "npu" or any( + tensor.device != query.device for tensor in (query_rope, cache, indices, query_ends, key_length) + ): + raise ValueError("HiSparse C8 inputs must share one NPU") + if any(tensor.shape != (1,) or tensor.dtype != torch.int32 for tensor in (query_ends, key_length)): + raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") + if not math.isfinite(scale) or scale <= 0: + raise ValueError("HiSparse C8 attention scale must be finite and positive") + mapped = compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) + packed_query = torch.cat((query, query_rope), dim=-1) + query_ends.fill_(rows) + key_length.fill_(view.shape[0]) + result = torch.ops._C_ascend.npu_kv_quant_sparse_flash_attention_vllm( + query=packed_query, + key=view, + value=view, + sparse_indices=mapped.unsqueeze(1), + scale_value=scale, + sparse_block_size=1, + actual_seq_lengths_query=query_ends, + actual_seq_lengths_kv=key_length, + layout_query="TND", + layout_kv="TND", + sparse_mode=0, + attention_mode=2, + quant_scale_repo_mode=1, + tile_size=128, + key_quant_mode=2, + value_quant_mode=2, + rope_head_dim=64, + return_softmax_lse=True, + ) + if isinstance(result, tuple): + result = result[0] + return torch.where((mapped >= 0).any(dim=-1)[:, None, None], result, 0) + def _attention_physical(self, query, query_rope, packed_cache, physical_topk): if self.enable_sparse_sfa_c8: - from vllm_ascend.ops.hisparse_c8 import c8_physical_attention - - lengths = self._c8_length_workspace() - return c8_physical_attention(query, query_rope, packed_cache, physical_topk, self.scale, *lengths) - from vllm_ascend.ops.hisparse import pack_attention_rows + return self._c8_physical_attention(query, query_rope, packed_cache, physical_topk) + from vllm_ascend.hisparse.cache_ops import pack_attention_rows latent, rope, block_table, indices, query_ends, seq_lens = self._workspace(query.dtype) positions = indices.view(1, -1).to(torch.float32) diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/__init__.py b/vllm_ascend/hisparse/__init__.py similarity index 59% rename from vllm_ascend/distributed/kv_transfer/hisparse/__init__.py rename to vllm_ascend/hisparse/__init__.py index 4b8ccbf15b4c..c1132d34e1e2 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/__init__.py +++ b/vllm_ascend/hisparse/__init__.py @@ -1,3 +1,3 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Ascend device operations for the upstream HiSparse connector.""" +"""Shared HiSparse cache operations, layouts and Core backend integration.""" diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py b/vllm_ascend/hisparse/backend.py similarity index 93% rename from vllm_ascend/distributed/kv_transfer/hisparse/backend.py rename to vllm_ascend/hisparse/backend.py index e8fe79772567..4a03a305b517 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""NPU streams, mapped host memory and DMA for the common HiSparse runtime.""" +"""Core HiSparse backend: NPU streams, mapped Host memory and DMA.""" from __future__ import annotations @@ -23,7 +23,7 @@ def check_available(self) -> None: from vllm_ascend.utils import bootstrap_custom_op_env # Import the device kernels before the connector allocates Host storage. - importlib.import_module("vllm_ascend.ops.hisparse") + importlib.import_module("vllm_ascend.hisparse.cache_ops") if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): # A5 disables the global custom-op selection policy. HiSparse still # needs its explicit DMA dependency; load the extension lazily here @@ -37,7 +37,7 @@ def check_available(self) -> None: raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") def pin_tensor(self, tensor: torch.Tensor) -> None: - from vllm_ascend.ops.hisparse import MappedHostTensor + from vllm_ascend.hisparse.cache_ops import MappedHostTensor tensor._hisparse_registration = MappedHostTensor(tensor) @@ -47,7 +47,7 @@ def unpin_tensor(self, tensor: torch.Tensor) -> None: del tensor._hisparse_registration def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: - from vllm_ascend.ops.hisparse import HiSparseNPUOps, ResolverWorkspace + from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, ResolverWorkspace group = runtime.index_group if runtime is group.leader: diff --git a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py b/vllm_ascend/hisparse/binding.py similarity index 99% rename from vllm_ascend/distributed/kv_transfer/hisparse/binding.py rename to vllm_ascend/hisparse/binding.py index 175c0201a813..1f98b844e81d 100644 --- a/vllm_ascend/distributed/kv_transfer/hisparse/binding.py +++ b/vllm_ascend/hisparse/binding.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Bind Ascend indexer views while retaining upstream HiSparse allocation.""" +"""Bind shared-slab views while retaining Core HiSparse allocation.""" import math from dataclasses import replace diff --git a/vllm_ascend/ops/hisparse_c8.py b/vllm_ascend/hisparse/c8.py similarity index 62% rename from vllm_ascend/ops/hisparse_c8.py rename to vllm_ascend/hisparse/c8.py index ce29c398edec..d2fe993fe3aa 100644 --- a/vllm_ascend/ops/hisparse_c8.py +++ b/vllm_ascend/hisparse/c8.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Native C8 payload/layout bridge; quantization and Attention stay in existing ops.""" - -import math +"""C8 payload packing, shared-slab views and physical index conversion.""" import torch @@ -85,55 +83,3 @@ def compact_c8_indices(indices: torch.Tensor) -> torch.Tensor: keys = positions + (indices < 0).to(torch.float32) * width order = torch.argsort(keys, dim=-1) return torch.gather(indices, -1, order) - - -def c8_physical_attention(query, query_rope, cache, indices, scale, query_ends, key_length): - """Adapt native packed rows to the existing quantized SFA operation. - - Keep compaction stable, including duplicates. The two current lengths are - device writes captured on the consumer stream; no Host read or stream/event - policy is introduced here. - """ - view, block, stride_rows, capacity = c8_storage_view(cache) - rows = query.shape[0] - if query.ndim != 3 or query.dtype != torch.bfloat16 or query.shape[-1] != _LATENT or rows <= 0: - raise ValueError("HiSparse C8 query must be nonempty BF16 [queries, heads, 512]") - if query_rope.shape != (*query.shape[:-1], 64) or query_rope.dtype != query.dtype: - raise ValueError("HiSparse C8 requires matching BF16 query RoPE") - if indices.ndim != 2 or indices.shape[0] != rows or indices.shape[1] <= 0 or indices.shape[1] % 128: - raise ValueError("HiSparse C8 top-k width must be positive and divisible by 128") - if query.device.type != "npu" or any( - tensor.device != query.device for tensor in (query_rope, cache, indices, query_ends, key_length) - ): - raise ValueError("HiSparse C8 inputs must share one NPU") - if any(tensor.shape != (1,) or tensor.dtype != torch.int32 for tensor in (query_ends, key_length)): - raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") - if not math.isfinite(scale) or scale <= 0: - raise ValueError("HiSparse C8 attention scale must be finite and positive") - mapped = compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) - packed_query = torch.cat((query, query_rope), dim=-1) - query_ends.fill_(rows) - key_length.fill_(view.shape[0]) - result = torch.ops._C_ascend.npu_kv_quant_sparse_flash_attention_vllm( - query=packed_query, - key=view, - value=view, - sparse_indices=mapped.unsqueeze(1), - scale_value=scale, - sparse_block_size=1, - actual_seq_lengths_query=query_ends, - actual_seq_lengths_kv=key_length, - layout_query="TND", - layout_kv="TND", - sparse_mode=0, - attention_mode=2, - quant_scale_repo_mode=1, - tile_size=128, - key_quant_mode=2, - value_quant_mode=2, - rope_head_dim=64, - return_softmax_lse=True, - ) - if isinstance(result, tuple): - result = result[0] - return torch.where((mapped >= 0).any(dim=-1)[:, None, None], result, 0) diff --git a/vllm_ascend/ops/hisparse.py b/vllm_ascend/hisparse/cache_ops.py similarity index 99% rename from vllm_ascend/ops/hisparse.py rename to vllm_ascend/hisparse/cache_ops.py index b2bbd7d387ea..87d335b6e6e8 100644 --- a/vllm_ascend/ops/hisparse.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Ascend HiSparse operators and explicit mapped-host memory ownership. +"""Shared HiSparse cache operations and explicit mapped-host memory ownership. Mapping and scratch allocation happen at cache initialization. Decode planning and payload copies stay on the NPU, including during ACL graph replay. diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 98058a9f6518..6eb4916031d9 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -167,7 +167,7 @@ def is_pin_memory_available(cls): @classmethod def get_hisparse_backend_cls(cls) -> str: - return "vllm_ascend.distributed.kv_transfer.hisparse.backend.AscendHiSparseBackend" + return "vllm_ascend.hisparse.backend.AscendHiSparseBackend" @classmethod def opaque_attention_op(cls) -> bool: diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 7364c427eca6..7d93f8594581 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -277,7 +277,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: # Keep Core's row-packed source spec and actual index-group ownership. source_layout = {} if getattr(attn_module.impl, "enable_sparse_sfa_c8", False): - from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + from vllm_ascend.hisparse.c8 import C8_ROW_BYTES source_layout.update(head_size=C8_ROW_BYTES, dtype=attn_module.impl.c8_cache_dtype) kv_cache_spec[layer_name] = replace( @@ -432,11 +432,11 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if vllm_config.attention_config.hisparse_config is not None: from vllm.v1.hisparse.binding import resolve_hisparse_specs - from vllm_ascend.distributed.kv_transfer.hisparse.binding import align_hisparse_indexer_specs + from vllm_ascend.hisparse.binding import align_hisparse_indexer_specs packed_row_bytes = None if get_ascend_config().enable_sparse_sfa_c8: - from vllm_ascend.ops.hisparse_c8 import C8_ROW_BYTES + from vllm_ascend.hisparse.c8 import C8_ROW_BYTES packed_row_bytes = C8_ROW_BYTES return align_hisparse_indexer_specs( From 5bdb93a2b5616164e2f1bbebb1b747cf2c986885 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 4 Nov 2026 23:37:36 +0800 Subject: [PATCH 012/120] test(hisparse): record device qualification and fix scatter reference Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 18 +++++++++++++----- .../pull_request/one_card/test_hisparse_ops.py | 3 +-- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 92ecea2fa2c9..d4f6d54be9cf 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -193,11 +193,19 @@ pytest --confcutdir=tests/e2e/pull_request/one_card \ tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py ``` -The full modules have not been rerun on this new split. Isolated stream-boundary -checks on Ascend 950 with Torch 2.10.0 / torch-npu 2.10.0.post7 passed for generic -and native interfaces, including the merged Core methods' leader/follower and -lazy-event paths. Those checks substitute small CANN work for resolver/refill; -they do not qualify the complete HiSparse graph or model. +On Ascend 950 with Torch 2.10.0, torch-npu 2.10.0.post7 and CANN 9.2.1, +the standard CPU module passed all 20 cases. Running the four complete NPU +modules passed 40 of 41 cases, including BF16/C8 Attention and changing graph +replays. The run reused the compiled Ascend extension and vendor package; +it was not an independent native rebuild. + +The remaining failure is `test_resident_prefill_plan_shared_layers[2-128]`. +With the paired Core pin unchanged, `res_blocks[:, offsets]` in +`HiSparsePrefillStagingPlan.ensure_gpu_sources` returned incorrect NPU rows +for the tested singleton-column input. An isolated CPU/NPU comparison found +the first difference at that indexing operation; equivalent `index_select` +matched the CPU result. The failing test remains enabled. This device stack +therefore has an unresolved mixed Resident/Host prefill correctness issue. Fresh full-model correctness, dataset accuracy, eager/Graph lifecycle and same-workload performance remain prerequisites for marking the foundation diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 72c6fa5c7717..a41ad581e1ef 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -55,8 +55,7 @@ def forward(): else: graph.replay() torch.npu.synchronize() - expected = torch.full_like(raw).cpu() - expected.fill_(165) + expected = torch.full_like(raw, 165).cpu() payload = values.cpu() for row, slot in enumerate(slots_cpu.tolist()): if row < count and slot >= 0: From 2b66c7bfb211b54566dbbf93fab27a52353905e2 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 00:01:26 +0800 Subject: [PATCH 013/120] fix(hisparse): preserve prefill row mapping on NPU Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 24 +++++--- .../one_card/test_hisparse_prefill_staging.py | 33 ++++++++++- vllm_ascend/attention/sfa_hisparse.py | 4 +- vllm_ascend/hisparse/prefill.py | 57 +++++++++++++++++++ 4 files changed, 107 insertions(+), 11 deletions(-) create mode 100644 vllm_ascend/hisparse/prefill.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index d4f6d54be9cf..5511c6a541de 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -110,6 +110,7 @@ vllm_ascend/ binding.py # Shared-slab specs and key/scale views cache_ops.py # Residency, refill, invalidation and byte copies c8.py # C8 packing, views and physical-index helpers + prefill.py # Core prefill plan with scoped NPU index selection ops/triton/hisparse.py # Cache-operation kernels; no Attention arithmetic ``` @@ -195,17 +196,24 @@ pytest --confcutdir=tests/e2e/pull_request/one_card \ On Ascend 950 with Torch 2.10.0, torch-npu 2.10.0.post7 and CANN 9.2.1, the standard CPU module passed all 20 cases. Running the four complete NPU -modules passed 40 of 41 cases, including BF16/C8 Attention and changing graph +modules passed all 42 cases, including BF16/C8 Attention and changing graph replays. The run reused the compiled Ascend extension and vendor package; it was not an independent native rebuild. -The remaining failure is `test_resident_prefill_plan_shared_layers[2-128]`. -With the paired Core pin unchanged, `res_blocks[:, offsets]` in -`HiSparsePrefillStagingPlan.ensure_gpu_sources` returned incorrect NPU rows -for the tested singleton-column input. An isolated CPU/NPU comparison found -the first difference at that indexing operation; equivalent `index_select` -matched the CPU result. The failing test remains enabled. This device stack -therefore has an unresolved mixed Resident/Host prefill correctness issue. +The original `test_resident_prefill_plan_shared_layers[2-128]` failure was +traced to NPU advanced column indexing in Core's `ensure_gpu_sources` for a +singleton-column input. The Ascend plan subclasses Core's plan and runs its +unchanged source-resolution method inside a local `TorchFunctionMode` that +uses equivalent `index_select` for NPU column selection. Core still constructs +the plan and owns representative pages, source-key caching and Resident/Host +masks. There is no global Torch replacement, copied planner or CPU readback. +Remove this adapter when the supported device stack fixes advanced indexing. + +The original case remains enabled and passes. A small regression also checks +changing request-state rows. Prefill metadata planning executes outside the +model graph: this stack's Core `scatter_reduce` path cannot be captured. The +existing eager/changing-graph payload and Attention cases continue to pass; +the tests do not claim graph capture of the complete Prefill planner. Fresh full-model correctness, dataset accuracy, eager/Graph lifecycle and same-workload performance remain prerequisites for marking the foundation diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 8975b88a2130..18c533f62df6 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -7,10 +7,11 @@ import pytest import torch import torch_npu # noqa: F401 -from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView, build_hisparse_prefill_staging_plan +from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, allocate_mapped_host, scatter_packed_kv +from vllm_ascend.hisparse.prefill import build_hisparse_prefill_staging_plan @pytest.mark.parametrize("resident_block_size", [128, 64, 32]) @@ -112,3 +113,33 @@ def test_backend_resident_payload_graph_replay(dtype, width): safe = torch.tensor(values).clamp_min(0) expected = source.view(torch.uint8).cpu()[safe // 2, safe % 2] assert torch.equal(destination.view(torch.uint8).cpu(), expected) + + +def test_prefill_singleton_column_mapping_tracks_current_residency(): + if not torch.npu.is_available(): + pytest.skip("HiSparse prefill mapping requires an NPU") + torch.npu.set_device(0) + block = 128 + host_table = torch.tensor([[1, 2]], dtype=torch.int32, device="npu:0") + lengths = torch.tensor([2 * block - 7], dtype=torch.int32, device="npu:0") + resident = torch.tensor([[2, 0], [0, 5]], dtype=torch.int32, device="npu:0") + state = torch.tensor([0], dtype=torch.int32, device="npu:0") + + def resolve(): + plan = build_hisparse_prefill_staging_plan(host_table, lengths, block, 2) + plan.ensure_gpu_sources(resident, state, block) + return plan.gpu_row_ids, plan.miss_mask + + # FULL_DECODE_ONLY builds Prefill metadata outside graph capture; Core's + # representative-page scatter_reduce is not a capturable NPU operation. + for state_row, pages in ((0, [2, 0]), (1, [0, 5]), (0, [2, 0])): + state.fill_(state_row) + gpu_rows, mask = resolve() + expected = torch.full((3 * block,), -1, dtype=torch.int32) + for logical, physical in enumerate(pages, start=1): + if physical: + expected[logical * block : (logical + 1) * block] = torch.arange( + physical * block, (physical + 1) * block, dtype=torch.int32 + ) + torch.testing.assert_close(gpu_rows.cpu().flatten(), expected, rtol=0, atol=0) + torch.testing.assert_close(mask.cpu().flatten(), (expected < 0).int(), rtol=0, atol=0) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 368bd35ba71f..0ffce7ab66aa 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -8,10 +8,10 @@ import torch import torch_npu from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder -from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan from vllm.v1.worker.workspace import current_workspace_manager from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder +from vllm_ascend.hisparse.prefill import AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan # Bound the temporary D2D packing workspace independently of prefill length. _ATTENTION_CHUNK_ROWS = 32 @@ -30,7 +30,7 @@ def finish_hisparse_kv_update(cache, metadata) -> None: @dataclass class AscendHiSparsePrefillMetadata: - host_staging_plan: HiSparsePrefillStagingPlan + host_staging_plan: AscendHiSparsePrefillStagingPlan @dataclass diff --git a/vllm_ascend/hisparse/prefill.py b/vllm_ascend/hisparse/prefill.py new file mode 100644 index 000000000000..0bf43a4b4808 --- /dev/null +++ b/vllm_ascend/hisparse/prefill.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reuse Core's prefill plan with NPU-safe column selection.""" + +import torch +from torch.overrides import TorchFunctionMode +from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan +from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan + + +class _PrefillColumnSelection(TorchFunctionMode): + """Keep a local workaround for the NPU singleton-column indexing defect. + + torch-npu 2.10.0.post7 / CANN 9.2.1 can interleave rows for ``x[:, ids]`` + when x has one column. index_select has the same semantics and returns + correct rows. Remove this scope when the supported device stack is fixed. + """ + + def __torch_function__(self, func, types, args=(), kwargs=None): + if func is torch.Tensor.__getitem__: + tensor, index = args + if ( + tensor.device.type == "npu" + and tensor.ndim == 2 + and isinstance(index, tuple) + and len(index) == 2 + and isinstance(index[0], slice) + and index[0] == slice(None) + and isinstance(index[1], torch.Tensor) + and index[1].ndim == 1 + ): + return torch.index_select(tensor, 1, index[1]) + return func(*args, **(kwargs or {})) + + +class AscendHiSparsePrefillStagingPlan(HiSparsePrefillStagingPlan): + def ensure_gpu_sources(self, resident_state_rows, state_indices, resident_block_size) -> None: + # Core retains source-key caching, representative-page selection and + # current Resident/Host masks. The override affects only this call. + with _PrefillColumnSelection(): + super().ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size) + + +def build_hisparse_prefill_staging_plan( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + staging_block_capacity: int, +) -> AscendHiSparsePrefillStagingPlan: + plan = build_core_prefill_plan(block_table, seq_lens, block_size, staging_block_capacity) + return AscendHiSparsePrefillStagingPlan( + block_table=plan.block_table, + row_ids=plan.row_ids, + dst_rows=plan.dst_rows, + miss_mask=plan.miss_mask, + block_size=plan.block_size, + ) From e282c9cfb48b235e8badc077814f94eccac591ac Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 00:19:26 +0800 Subject: [PATCH 014/120] fix(hisparse): retain indexer backend beyond model construction Signed-off-by: zhenwenqi2024 --- .../test_hisparse_foundation_cpu.py | 24 +++++++++++++++++++ vllm_ascend/patch/worker/patch_attn_utils.py | 18 ++++++++++---- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index f100adc46eec..0363927194f2 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -30,6 +30,30 @@ ) +@pytest.mark.parametrize("hisparse", [False, True]) +def test_indexer_backend_choice_survives_model_construction_context(hisparse): + from vllm.config import CacheConfig, VllmConfig, set_current_vllm_config + from vllm.config.attention import HiSparseConfig + from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache + + from vllm_ascend.attention.indexer import AscendSFAIndexerBackend + from vllm_ascend.patch.worker import patch_attn_utils # noqa: F401 + + config = VllmConfig() + config.attention_config.hisparse_config = HiSparseConfig() if hisparse else None + expected = AscendHiSparseIndexerBackend if hisparse else AscendSFAIndexerBackend + with set_current_vllm_config(config): + cache = DeepseekV32IndexerCache(132, torch.uint8, "regression.indexer.k_cache", CacheConfig()) + assert cache.get_attn_backend() is expected + # Core's post-load block-size discovery runs after the construction context + # exits. A later model's context must not change an existing cache either. + assert cache.get_attn_backend() is expected + other = VllmConfig() + other.attention_config.hisparse_config = None if hisparse else HiSparseConfig() + with set_current_vllm_config(other): + assert cache.get_attn_backend() is expected + + def test_failed_unregister_retains_host_owner_for_safe_retry(monkeypatch): tensor = torch.empty(64, dtype=torch.uint8) events = [] diff --git a/vllm_ascend/patch/worker/patch_attn_utils.py b/vllm_ascend/patch/worker/patch_attn_utils.py index 77c021560c54..b3a431b29813 100644 --- a/vllm_ascend/patch/worker/patch_attn_utils.py +++ b/vllm_ascend/patch/worker/patch_attn_utils.py @@ -23,10 +23,20 @@ ) -def _get_ascend_sfa_indexer_backend(_self): - if get_current_vllm_config().attention_config.hisparse_config is not None: - return AscendHiSparseIndexerBackend - return AscendSFAIndexerBackend +def _get_ascend_sfa_indexer_backend(self): + backend = getattr(self, "_ascend_indexer_backend", None) + if backend is None: + # IndexerWrapper first queries this while constructing the layer. + # Post-load platform discovery also queries it, outside that config + # context. Keep the choice on this cache rather than consulting a + # temporary global context again. + backend = ( + AscendHiSparseIndexerBackend + if get_current_vllm_config().attention_config.hisparse_config is not None + else AscendSFAIndexerBackend + ) + self._ascend_indexer_backend = backend + return backend # vLLM main (#63384) moved OLMoE to the Transformers backend, whose From 83e4655fceb78851182ad6cc2c9c9556c09b22de Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 00:31:13 +0800 Subject: [PATCH 015/120] fix(hisparse): preserve core indexer cache roles and layout Signed-off-by: zhenwenqi2024 --- .../kv_transfer/test_hisparse_foundation_cpu.py | 10 +++++++++- tests/ut/worker/test_attn_utils_v2.py | 6 +++++- vllm_ascend/core/kv_cache_interface.py | 16 +++++----------- vllm_ascend/worker/v2/attn_utils.py | 1 + 4 files changed, 20 insertions(+), 13 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 0363927194f2..2afed4894d5e 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -8,7 +8,7 @@ import pytest import torch -from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec, SparseCacheRole from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder, finish_hisparse_kv_update @@ -139,6 +139,14 @@ def indexer_spec(): ) +@pytest.mark.parametrize("role", [SparseCacheRole.SPARSE, SparseCacheRole.INDEXER]) +def test_indexer_spec_merge_retains_core_cache_role(indexer_spec, role): + spec = replace(indexer_spec, cache_role=role) + merged = AscendSFAIndexerCacheSpec.merge([spec, spec]) + assert merged.cache_role is role + assert merged.block_stride_alignment == spec.block_stride_alignment + + @pytest.mark.parametrize("blocks,stride", [(1, 1), (3, 2), (5, 3)]) def test_expanded_indexer_alias_preserves_shared_slab_offsets(indexer_spec, blocks, stride): spec = indexer_spec diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index e8ef130214bf..c502fa2bd865 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -18,6 +18,7 @@ MambaSpec, MLAAttentionSpec, SlidingWindowSpec, + SparseCacheRole, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu import attn_utils as upstream_attn_utils @@ -1243,7 +1244,9 @@ def test_sfa_indexer_cache_spec_runtime_ownership_and_dcp_replication( monkeypatch.setattr( indexer_module, "get_kv_cache_spec", - lambda _config: object(), + lambda _config: MLAAttentionSpec( + block_size=128, num_kv_heads=1, head_size=132, dtype=torch.uint8, cache_role=SparseCacheRole.INDEXER + ), ) layers = {layer_name: indexer_module} if owner != "unpaired": @@ -1296,6 +1299,7 @@ def test_sfa_indexer_cache_spec_runtime_ownership_and_dcp_replication( spec = specs[layer_name] assert isinstance(spec, AscendSFAIndexerCacheSpec) + assert spec.cache_role is SparseCacheRole.INDEXER assert spec.sfa_dcp_replicated_indexer_size == expected_size if li_c4: assert spec.head_size == 128 // 2 diff --git a/vllm_ascend/core/kv_cache_interface.py b/vllm_ascend/core/kv_cache_interface.py index 071d629e67a0..53823c22cbd2 100644 --- a/vllm_ascend/core/kv_cache_interface.py +++ b/vllm_ascend/core/kv_cache_interface.py @@ -325,17 +325,13 @@ def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, AscendSFAIndexerCacheSpec) for spec in specs), ( "All attention layers in the same KV cache group must be AscendSFAIndexerCacheSpec." ) - cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) - dtype_set = set(spec.dtype for spec in specs) scale_dim_set = set(spec.scale_dim for spec in specs) scale_dtype_set = set(spec.scale_dtype for spec in specs) cache_sparse_li_c8_set = set(spec.cache_sparse_li_c8 for spec in specs) cache_sparse_li_c4_set = set(spec.cache_sparse_li_c4 for spec in specs) sfa_dcp_replicated_indexer_size_set = set(spec.sfa_dcp_replicated_indexer_size for spec in specs) assert ( - len(cache_dtype_str_set) == 1 - and len(dtype_set) == 1 - and len(scale_dim_set) == 1 + len(scale_dim_set) == 1 and len(scale_dtype_set) == 1 and len(cache_sparse_li_c8_set) == 1 and len(cache_sparse_li_c4_set) == 1 @@ -345,12 +341,10 @@ def merge(cls, specs: list[Self]) -> Self: "the same dtype, scale layout, quantization method, LI quant flags " "and DCP replication size." ) - return cls( - block_size=specs[0].block_size, - num_kv_heads=specs[0].num_kv_heads, - head_size=specs[0].head_size, - dtype=dtype_set.pop(), - cache_dtype_str=cache_dtype_str_set.pop(), + # Preserve Core's role, padding and stride alignment, and reuse its + # standard-field validation before restoring the Ascend-only fields. + return replace( + super().merge(specs), scale_dim=scale_dim_set.pop(), scale_dtype=scale_dtype_set.pop(), cache_sparse_li_c8=cache_sparse_li_c8_set.pop(), diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 7d93f8594581..282dd646b435 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -352,6 +352,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: kv_cache_spec[layer_name] = AscendSFAIndexerCacheSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, + cache_role=spec.cache_role, head_size=head_dim // 2 if cache_sparse_li_c4 else head_dim, dtype=torch.uint8 if cache_sparse_li_c4 From 81100e0f3eba7a4a9bdb5708b981b51ed57be31f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 01:02:07 +0800 Subject: [PATCH 016/120] fix(hisparse): allocate pinned private pools before cache views Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 17 ++++++++- release_notes/fragments/20926.yaml | 6 +-- .../one_card/test_hisparse_ops.py | 37 +++++++++++++++++++ .../test_hisparse_foundation_cpu.py | 26 +++++++++++++ vllm_ascend/hisparse/backend.py | 11 ++++++ vllm_ascend/hisparse/cache_ops.py | 13 ++++--- vllm_ascend/patch/__init__.py | 8 ++++ .../patch/worker/patch_hisparse_host_pool.py | 32 ++++++++++++++++ 8 files changed, 141 insertions(+), 9 deletions(-) create mode 100644 vllm_ascend/patch/worker/patch_hisparse_host_pool.py diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 5511c6a541de..0be47999c09d 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -206,7 +206,7 @@ singleton-column input. The Ascend plan subclasses Core's plan and runs its unchanged source-resolution method inside a local `TorchFunctionMode` that uses equivalent `index_select` for NPU column selection. Core still constructs the plan and owns representative pages, source-key caching and Resident/Host -masks. There is no global Torch replacement, copied planner or CPU readback. +masks. This adapter adds no global Torch replacement, copied planner or CPU readback. Remove this adapter when the supported device stack fixes advanced indexing. The original case remains enabled and passes. A small regression also checks @@ -219,3 +219,18 @@ Fresh full-model correctness, dataset accuracy, eager/Graph lifecycle and same-workload performance remain prerequisites for marking the foundation ready. Existing C8 benchmark results are intentionally not presented as this implementation's performance. + +Full-model integration additionally preserves the Core Indexer cache role and +merge semantics, and retains each Indexer cache's backend choice after the model +construction context exits. With these fixes, the complete TP8/MTP3 model loaded, +initialized its KV caches, captured all eight ranks and reached serving. Its +first real request exposed a Host writeback failure; serving alone is not a +functional pass. + +The DMA reproduction distinguishes page locking from mapping: pageable and +PTA-pinned destinations passed, while a pageable allocation registered only for +mapping failed both batch and single-copy APIs on this stack. The private-pool +adapter reuses PR #19109's pinned-allocation approach before Core creates any +cache views. Core still owns completion, cache views and release. Other backends +keep their original allocator. The added writeback regression and a new full-model +run remain pending hardware validation for this allocation change. diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index 17d0715b4b32..f007c28f72bf 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -2,7 +2,7 @@ pr: 20926 entries: - category: feature text: >- - Added an experimental Ascend HiSparse backend and unquantized sparse - MLA/SFA data path, reusing Core allocation and lifecycle interfaces. - Device and full-model qualification of the foundation remains pending. + Added experimental Ascend HiSparse BF16 and native C8 sparse MLA/SFA + support using existing Attention operators and Core lifecycle interfaces. + Full-model qualification of the foundation remains pending. hardware: [A5] diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index a41ad581e1ef..53c7268f8d8a 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -8,6 +8,7 @@ import torch import torch_npu # noqa: F401 +from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, MappedHostTensor, @@ -25,6 +26,42 @@ def _require_mapped_host_device(): torch.npu.set_device(0) +def test_private_host_pool_retains_pinned_storage_for_dma_writeback(): + from vllm.v1.hisparse import runtime + + backend = AscendHiSparseBackend() + backend.check_available() + payload, registered = runtime.allocate_pinned_host_pool(32768) + try: + assert payload.numel() == 32768 + assert payload.is_pinned() and registered.is_pinned() + assert registered.data_ptr() % 4096 == 0 + mapping = registered._hisparse_registration + assert mapping.address_of(payload[656:1312]) == mapping.device_ptr + 656 + assert payload.untyped_storage().data_ptr() == registered.untyped_storage().data_ptr() + source = torch.arange(32768, dtype=torch.int32, device="npu").to(torch.uint8) + stream = torch.npu.Stream() + for change in (1, 3, 7): + source.bitwise_xor_(change) + payload.fill_(-1) + src = torch.tensor([source.data_ptr(), source.data_ptr() + 8192], dtype=torch.int64) + dst = torch.tensor([payload.data_ptr() + 656, payload.data_ptr() + 16384], dtype=torch.int64) + sizes = torch.tensor([656, 1312], dtype=torch.int64) + stream.wait_stream(torch.npu.current_stream()) + with stream: + backend.swap_blocks_batch(src, dst, sizes) + complete = torch.npu.Event() + complete.record() + complete.synchronize() + expected = torch.full((32768,), 255, dtype=torch.uint8) + expected[656:1312] = source[:656].cpu() + expected[16384:17696] = source[8192:9504].cpu() + torch.testing.assert_close(payload.view(torch.uint8), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + backend.unpin_tensor(registered) + + @pytest.mark.parametrize("capture", [False, True]) def test_c8_bytes_reuse_existing_scatter_with_padding_and_page_gaps(capture): rows, pages, block, width = 5, 4, 128, 656 diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 2afed4894d5e..e00f880473f5 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -80,6 +80,32 @@ def close(self): assert events == ["close", "close"] +def test_private_host_allocator_keeps_other_backends_on_core_path(monkeypatch): + from vllm.v1.hisparse import backend as core_backend + from vllm.v1.hisparse import runtime + + from vllm_ascend.patch.worker.patch_hisparse_host_pool import install + + class OtherBackend(core_backend.HiSparseBackend): + def pin_tensor(self, tensor): + tensor._test_registration = self + + backend = OtherBackend() + monkeypatch.setattr(runtime, "get_hisparse_backend", lambda: backend) + monkeypatch.setattr(core_backend, "get_hisparse_backend", lambda: backend) + # Restore the process-wide binding when this test exits. + monkeypatch.setattr(runtime, "allocate_pinned_host_pool", runtime.allocate_pinned_host_pool) + install() + installed = runtime.allocate_pinned_host_pool + install() + assert runtime.allocate_pinned_host_pool is installed + payload, registered = installed(1001) + assert registered._test_registration is backend + assert payload.numel() == 1001 and registered.data_ptr() % 4096 == 0 + payload.fill_(17) + torch.testing.assert_close(registered[:1001], payload) + + def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): backing = torch.arange(128, dtype=torch.uint8) mapping = hisparse_ops.MappedHostTensor.__new__(hisparse_ops.MappedHostTensor) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 4a03a305b517..2ba23c2690a5 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -20,8 +20,10 @@ class AscendHiSparseBackend(HiSparseBackend): stats_dtype = torch.int64 def check_available(self) -> None: + from vllm_ascend.patch.worker.patch_hisparse_host_pool import install from vllm_ascend.utils import bootstrap_custom_op_env + install() # Import the device kernels before the connector allocates Host storage. importlib.import_module("vllm_ascend.hisparse.cache_ops") if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): @@ -36,6 +38,15 @@ def check_available(self) -> None: if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") + def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Tensor]: + from vllm_ascend.hisparse.cache_ops import allocate_mapped_host + + payload = allocate_mapped_host((size,), torch.int8) + registration = payload._hisparse_registration + registered = registration.tensor.view(torch.int8) + registered._hisparse_registration = registration + return registered[:size], registered + def pin_tensor(self, tensor: torch.Tensor) -> None: from vllm_ascend.hisparse.cache_ops import MappedHostTensor diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 87d335b6e6e8..085dfd3f7087 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -99,18 +99,21 @@ def close(self) -> None: def allocate_mapped_host(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: - """Allocate page-aligned CPU storage and attach its explicit mapping owner.""" + """Allocate page-locked, aligned storage for both DMA and AIV access.""" count = math.prod(shape) nbytes = count * torch.empty((), dtype=dtype).element_size() if nbytes <= 0: raise ValueError("HiSparse host allocation must be nonempty") allocated_bytes = triton.cdiv(nbytes, mmap.PAGESIZE) * mmap.PAGESIZE - owner = mmap.mmap(-1, allocated_bytes) - raw = torch.frombuffer(owner, dtype=torch.uint8) + # Mapping permits AIV access; DMA also requires a page-locked allocation. + # Align a view of PTA's pinned allocation without creating a second pool. + owner = torch.empty(allocated_bytes + mmap.PAGESIZE - 1, dtype=torch.uint8, pin_memory=True) + offset = (-owner.data_ptr()) % mmap.PAGESIZE + raw = owner[offset : offset + allocated_bytes] mapping = MappedHostTensor(raw) tensor = raw[:nbytes].view(dtype).view(shape) - # The mapping owns raw; frombuffer keeps the mmap alive. Runtime owns the - # registration through this base tensor until its explicit shutdown. + # Both aliases retain the pinned allocation. Runtime owns the registration + # through this base tensor until its explicit, drained shutdown. tensor._hisparse_registration = mapping return tensor diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 12113a662d21..e57106438ef7 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -28,6 +28,14 @@ # What's Patched and how it works: # -------------------------------- +# * worker/patch_hisparse_host_pool.py: +# Target: vllm.v1.hisparse.runtime.allocate_pinned_host_pool. +# Why: AIV mapping alone does not enable DMA to pageable Host storage. +# How: Installed only by the HiSparse backend; allocate private pinned storage +# before Core builds cache views, retaining Core's release lifecycle. +# Related PR: No backend Host-allocation hook exists in the supported Core. +# Future plan: Replace this adapter with that upstream hook when available. +# -------------------------------- # * worker/patch_kda_recoverssm.py: # Target: Kimi KDA's KDARecoverSSMCommitContext module binding. # Why: The upstream metadata builder imports its CUDA context directly. diff --git a/vllm_ascend/patch/worker/patch_hisparse_host_pool.py b/vllm_ascend/patch/worker/patch_hisparse_host_pool.py new file mode 100644 index 000000000000..b8432769d2a6 --- /dev/null +++ b/vllm_ascend/patch/worker/patch_hisparse_host_pool.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Allocate pinned private HiSparse storage before Core creates cache views. + +Core registers an aligned pageable allocation through ``pin_tensor``. AIV +mapping alone does not make that storage a valid asynchronous DMA destination +on Ascend. Reuse PTA's pinned allocator and Core's existing release ownership. +Remove this adapter when Core exposes a backend Host-allocation hook. +""" + +from functools import wraps + + +def install() -> None: + from vllm.v1.hisparse import runtime + from vllm.v1.hisparse.backend import get_hisparse_backend + + from vllm_ascend.hisparse.backend import AscendHiSparseBackend + + original = runtime.allocate_pinned_host_pool + if getattr(original, "_ascend_pinned_private_pool", False): + return + + @wraps(original) + def allocate(size): + backend = get_hisparse_backend() + if isinstance(backend, AscendHiSparseBackend): + return backend.allocate_private_host_pool(size) + return original(size) + + allocate._ascend_pinned_private_pool = True + runtime.allocate_pinned_host_pool = allocate From 8042477929169d26d9d470acc8312e70b56bf19a Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 02:01:21 +0800 Subject: [PATCH 017/120] fix(hisparse): propagate captured side-stream waits before workspace allocation Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 69 ++++++++++++++++++- vllm_ascend/attention/sfa_hisparse.py | 3 + vllm_ascend/hisparse/backend.py | 33 +++++++++ 3 files changed, 104 insertions(+), 1 deletion(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 53c7268f8d8a..aafe5d81f2ea 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -8,7 +8,7 @@ import torch import torch_npu # noqa: F401 -from vllm_ascend.hisparse.backend import AscendHiSparseBackend +from vllm_ascend.hisparse.backend import AscendHiSparseBackend, capture_aware_stream from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, MappedHostTensor, @@ -62,6 +62,73 @@ def test_private_host_pool_retains_pinned_storage_for_dma_writeback(): backend.unpin_tensor(registered) +def test_resolver_sort_workspace_survives_cross_stream_graph_owners(): + # Production-sized Hot/TopK buffers exercise CANN's workspace sort path; + # the small resolver oracle cases do not need this external workspace. + width, topk, max_rows = 10240, 2048, 128 + groups = [] + for _ in range(22): + stream = torch.Stream(device="npu") + wrapped = capture_aware_stream(stream) + assert wrapped.stream_id == stream.stream_id and wrapped.device == stream.device + assert capture_aware_stream(wrapped) is wrapped + cache = torch.full((max_rows, width), 2**63 - 1, dtype=torch.int64, device="npu") + rows = torch.full((max_rows, topk), 2**63 - 1, dtype=torch.int64, device="npu") + groups.append( + ( + wrapped, + torch.Event(), + cache, + torch.empty_like(cache), + torch.empty_like(cache), + rows, + torch.empty_like(rows), + torch.empty_like(rows), + ) + ) + + def forward(count): + current = torch.accelerator.current_stream() + for stream, ready, cache, values, indices, rows, row_values, row_indices in groups: + stream.wait_stream(current) + with stream: + torch.sort(rows[:count], dim=-1, out=(row_values[:count], row_indices[:count])) + torch.sort(cache[:count], dim=-1, out=(values[:count], indices[:count])) + ready.record(stream) + current.wait_event(ready) + + graphs = {} + pool = torch.npu.graph_pool_handle() + try: + for count in (128, 64, 32, 16, 8, 4, 1): + forward(count) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph, pool=pool): + forward(count) + graphs[count] = graph + for step, count in enumerate((4, 1, 8, 4, 16, 4)): + hot_ids = torch.arange(width, dtype=torch.int64) + query_ids = torch.arange(topk, dtype=torch.int64) + hot = ((((hot_ids + step) % 127) << 32) | hot_ids).flip(0) + query = ((((query_ids + step) % 67) << 32) | query_ids).flip(0) + hot[::17] = query[::17] = 2**63 - 1 + expected_hot = torch.sort(hot).values.expand(count, -1) + expected_query = torch.sort(query).values.expand(count, -1) + for _, _, cache, _, _, rows, _, _ in groups: + cache.copy_(hot.npu()) + rows.copy_(query.npu()) + graphs[count].replay() + torch.npu.synchronize() + for _, _, _, values, _, _, row_values, _ in groups: + torch.testing.assert_close(values[:count].cpu(), expected_hot, rtol=0, atol=0) + torch.testing.assert_close(row_values[:count].cpu(), expected_query, rtol=0, atol=0) + finally: + torch.npu.synchronize() + for graph in graphs.values(): + graph.reset() + + @pytest.mark.parametrize("capture", [False, True]) def test_c8_bytes_reuse_existing_scatter_with_padding_and_page_gaps(capture): rows, pages, block, width = 5, 4, 128, 656 diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 0ffce7ab66aa..cd63d2213937 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -123,6 +123,9 @@ def __init__(self, *args, **kwargs): ) if not isinstance(self.index_group, HiSparseMLAIndexGroup): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") + from vllm_ascend.hisparse.backend import capture_aware_stream + + self.index_group.side_stream = capture_aware_stream(self.index_group.side_stream) self.hisparse_cache = self.index_group.cache(self.index_group_index) self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 2ba23c2690a5..680cdd88c3fc 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -15,6 +15,39 @@ from vllm.v1.kv_cache_interface import KVCacheConfig +class _CaptureAwareStream(torch.Stream): + """Propagate a captured dependency before PTA allocates CANN workspace. + + PTA's native NPUEvent wait drains Host submission queues during capture; + its generic event binding does not. Keep generic event points and stream + identity, adding the missing submission boundary only during capture. + Remove when the supported PTA generic binding has the same behavior. + """ + + def _finish_capture_propagation(self) -> None: + if torch.npu.is_current_stream_capturing(): + native = torch.npu.Stream( + stream_id=self.stream_id, device_index=self.device_index, device_type=self.device_type + ) + # Obtaining the native handle drains its Host task queue. This is + # not a device synchronization and does not run during graph replay. + _ = native.npu_stream + + def wait_event(self, event) -> None: + super().wait_event(event) + self._finish_capture_propagation() + + def wait_stream(self, stream) -> None: + super().wait_stream(stream) + self._finish_capture_propagation() + + +def capture_aware_stream(stream: torch.Stream) -> torch.Stream: + if stream.device.type != "npu" or isinstance(stream, _CaptureAwareStream): + return stream + return _CaptureAwareStream(stream.stream_id, stream.device_index, stream.device_type) + + class AscendHiSparseBackend(HiSparseBackend): # Inherit Core's private Host-pool policy; NPU event IPC is not implemented. stats_dtype = torch.int64 From 0267e49e1ab318a6030224f1ff6e53a3db4b5292 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 02:09:46 +0800 Subject: [PATCH 018/120] fix(hisparse): adapt the bound runtime stream before graph capture Signed-off-by: zhenwenqi2024 --- .../ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py | 2 +- vllm_ascend/attention/sfa_hisparse.py | 3 --- vllm_ascend/hisparse/backend.py | 3 +++ 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index e00f880473f5..50a05b9e8f5d 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -129,7 +129,7 @@ def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups backend = AscendHiSparseBackend() def group(): - owner = NS() + owner = NS(copy_stream=NS(device=torch.device("cpu"))) leader = NS( index_group=owner, max_swap_rows=max_rows, diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index cd63d2213937..0ffce7ab66aa 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -123,9 +123,6 @@ def __init__(self, *args, **kwargs): ) if not isinstance(self.index_group, HiSparseMLAIndexGroup): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") - from vllm_ascend.hisparse.backend import capture_aware_stream - - self.index_group.side_stream = capture_aware_stream(self.index_group.side_stream) self.hisparse_cache = self.index_group.cache(self.index_group_index) self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 680cdd88c3fc..4a35af878677 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -94,6 +94,9 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, ResolverWorkspace group = runtime.index_group + # Runtime copied the model group's stream reference at layer creation. + # Adapt the live runtime binding before any graph captures use it. + group.copy_stream = capture_aware_stream(group.copy_stream) if runtime is group.leader: workspace = ResolverWorkspace( max_rows=runtime.max_swap_rows, From 908bc2c30a0cfddd48eb44704d6f5afca88bcd41 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 02:22:54 +0800 Subject: [PATCH 019/120] docs(hisparse): record full-model foundation qualification Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 52 ++++++++++++------- 1 file changed, 34 insertions(+), 18 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 0be47999c09d..bed9e6fb44c7 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -4,8 +4,9 @@ This is the foundational split of [#19109](https://github.com/vllm-project/vllm-ascend/pull/19109). It supplies an Ascend backend for the merged vLLM HiSparse API and the minimal -SFA data path. It is a contribution starting point, with new-head NPU and full-model -qualification pending. Historical optimized C8 performance does not qualify this split. +SFA data path. It is a contribution starting point with NPU component checks +and a complete A5 model smoke test. Dataset accuracy and performance qualification +remain pending. Historical optimized C8 performance does not qualify this split. Core owns HMA allocation, Coordinator admission/eviction, prefix-cache state, request lifetimes, prefetch control flow, mirror completion and release/rollback. @@ -25,8 +26,7 @@ The data path covers unquantized KV and the native Ascend 950 C8 format with the RoPE SFA layout and Model Runner V2. Model-weight quantization is independent. Private Host pools are used per rank because NPU IPC events are not implemented. The mapped-Host capability is checked at registration; hardware evidence for -the original primitives is from Ascend 950. A3 and the new source pair still -need their own qualification. +the current primitives is from Ascend 950. A3 still needs its own qualification. The SFA kernel requires TopK divisible by 128. The resolver's int16 LRU format limits aligned Hot capacity to 32768 rows, and Hot must cover the complete @@ -119,10 +119,15 @@ implementation owns query preparation, current length tensors and the call to the existing native SFA operation. The Triton and native Attention kernels are unchanged by this organization. -Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces are -used directly. The Ascend backend preserves Core's event placement and does not -replace its Torch bindings. This PR tests its own NPU operations and changing -graph replays rather than duplicating generic Torch stream tests. +Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces retain +their event placement and stream identities. On the supported PTA stack, generic +waits do not drain Host submission queues to propagate a graph-capture fork before +CANN allocates workspace. The backend wraps only the bound runtime copy stream: +after a wait during capture, obtaining the native handle drains that Host queue. +It adds no device synchronization or replay-time Python work and replaces no +global Torch binding. A regression covers production-size sort workspace across +multiple graph owners and changing replays. Remove this adapter when PTA's generic +binding propagates the captured dependency itself. The model-owned index-group builder is carried in Core's existing `MLAModules.index_group_builder` field, preserving the common MLA wrapper ABI. @@ -195,10 +200,11 @@ pytest --confcutdir=tests/e2e/pull_request/one_card \ ``` On Ascend 950 with Torch 2.10.0, torch-npu 2.10.0.post7 and CANN 9.2.1, -the standard CPU module passed all 20 cases. Running the four complete NPU -modules passed all 42 cases, including BF16/C8 Attention and changing graph -replays. The run reused the compiled Ascend extension and vendor package; -it was not an independent native rebuild. +25 foundation CPU cases, the existing MLA-wrapper case and 32 Indexer-spec cases +passed: 58 in total. Running the four complete NPU modules passed all 44 cases, +including BF16/C8 Attention, pinned Host writeback, production-size cross-stream +sort workspace and changing graph replays. The run reused the compiled Ascend +extension and vendor package; it was not an independent native rebuild. The original `test_resident_prefill_plan_shared_layers[2-128]` failure was traced to NPU advanced column indexing in Core's `ensure_gpu_sources` for a @@ -215,10 +221,9 @@ model graph: this stack's Core `scatter_reduce` path cannot be captured. The existing eager/changing-graph payload and Attention cases continue to pass; the tests do not claim graph capture of the complete Prefill planner. -Fresh full-model correctness, dataset accuracy, eager/Graph lifecycle and -same-workload performance remain prerequisites for marking the foundation -ready. Existing C8 benchmark results are intentionally not presented as this -implementation's performance. +Dataset accuracy, additional lifecycle coverage and same-workload performance +remain prerequisites for production qualification. Existing C8 benchmark results +are intentionally not presented as this implementation's performance. Full-model integration additionally preserves the Core Indexer cache role and merge semantics, and retains each Indexer cache's backend choice after the model @@ -232,5 +237,16 @@ PTA-pinned destinations passed, while a pageable allocation registered only for mapping failed both batch and single-copy APIs on this stack. The private-pool adapter reuses PR #19109's pinned-allocation approach before Core creates any cache views. Core still owns completion, cache views and release. Other backends -keep their original allocator. The added writeback regression and a new full-model -run remain pending hardware validation for this allocation change. +keep their original allocator. The writeback regression passed on the NPU with +changing payloads, completion fences and untouched-neighbor checks. + +The next full-model attempt exposed invalid CANN sort workspace during captured +cross-stream execution. A production-size reproducer retained generic streams, +events and the sort operation: the original path failed, while capture-propagated +forks passed. With the narrow backend adapter, the complete GLM-5.2 W4A8C8 model +loaded with TP8, MTP3, V2 and native C8 KV; all eight ranks captured. Three natural +controls returned the expected answers. Two concurrent requests each consumed +exactly 16,384 input tokens and returned sensible summaries (24 and 23 output +tokens); all five responses were successful with actual token IDs. Host mirroring +remained enabled. The model exited cleanly, its workers were reaped and the lease +was released. This smoke test is not a dataset accuracy, P/D or performance result. From 6ea9e7875591631a4a82d6c06d0b1fa57083532b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 16:20:14 +0800 Subject: [PATCH 020/120] refactor(hisparse): use the fused native residency resolver Signed-off-by: zhenwenqi2024 --- CMakeLists.txt | 29 + csrc/hisparse_runtime/build.py | 95 +++ csrc/hisparse_runtime/resolver.asc | 603 ++++++++++++++++++ csrc/hisparse_runtime/torch_binding.cpp | 205 ++++++ .../user_guide/feature_guide/hisparse.md | 308 +++------ release_notes/fragments/20926.yaml | 5 +- .../one_card/test_hisparse_ops.py | 86 +-- .../test_hisparse_foundation_cpu.py | 49 +- vllm_ascend/hisparse/backend.py | 2 + vllm_ascend/hisparse/cache_ops.py | 188 +----- vllm_ascend/hisparse/resolver.py | 88 +++ vllm_ascend/ops/triton/hisparse.py | 410 +----------- 12 files changed, 1172 insertions(+), 896 deletions(-) create mode 100644 csrc/hisparse_runtime/build.py create mode 100644 csrc/hisparse_runtime/resolver.asc create mode 100644 csrc/hisparse_runtime/torch_binding.cpp create mode 100644 vllm_ascend/hisparse/resolver.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 1393308c5537..5250aa299627 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,6 +26,8 @@ if(NOT ${TORCH_VERSION} VERSION_EQUAL "2.10.0") message(FATAL_ERROR "Expected PyTorch version 2.10.0, but found ${TORCH_VERSION}") endif() + + set(RUN_MODE "npu" CACHE STRING "cpu/sim/npu") set(SOC_VERSION ${SOC_VERSION}) message(STATUS "Detected SOC version: ${SOC_VERSION}") @@ -232,3 +234,30 @@ if(SOC_VERSION MATCHES "ascend310p.*|ascend950") else() install(TARGETS vllm_ascend_C vllm_ascend_kernels DESTINATION ${VLLM_ASCEND_INSTALL_PATH}) endif() + +# The optional HiSparse resolver uses SIMT, independently of Attention kernels. +# Older toolkits and other SOCs keep their ordinary build unchanged; selecting +# HiSparse without this extension raises an explicit availability error. +if(SOC_VERSION MATCHES "ascend950") + find_path(HISPARSE_SIMT_INCLUDE_DIR simt_api/asc_simt.h + PATHS + ${ASCEND_HOME_PATH}/asc/include + ${ASCEND_HOME_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include + NO_DEFAULT_PATH) + if(HISPARSE_SIMT_INCLUDE_DIR) + set(HISPARSE_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/csrc/hisparse_runtime) + set(HISPARSE_OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/hisparse_runtime) + add_custom_command( + OUTPUT ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so + COMMAND ${PYTHON_EXECUTABLE} ${HISPARSE_SOURCE}/build.py + --cann-root ${ASCEND_HOME_PATH} --output ${HISPARSE_OUTPUT} + --npu-arch dav-3510 --resolver-ub-bytes 175616 + DEPENDS ${HISPARSE_SOURCE}/build.py ${HISPARSE_SOURCE}/resolver.asc ${HISPARSE_SOURCE}/torch_binding.cpp + VERBATIM) + add_custom_target(hisparse_resolver + DEPENDS ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so) + add_dependencies(vllm_ascend_C hisparse_resolver) + install(FILES ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so + DESTINATION ${VLLM_ASCEND_INSTALL_PATH}/hisparse) + endif() +endif() diff --git a/csrc/hisparse_runtime/build.py b/csrc/hisparse_runtime/build.py new file mode 100644 index 000000000000..ac4c368d7c21 --- /dev/null +++ b/csrc/hisparse_runtime/build.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Build the HiSparse fused residency resolver for an explicit CANN target. + +Run in the matching A5 CANN 9.2 / torch 2.10 / torch-npu 2.10 environment. +Outputs stay in --output; installing into a serving environment is explicit. +""" + +import argparse +import importlib.util +import subprocess +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cann-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--npu-arch", required=True, help="SIMT-capable bisheng target architecture") + parser.add_argument( + "--resolver-ub-bytes", type=int, required=True, help="Usable per-AIV UB bytes, including descriptors" + ) + parser.add_argument("--install", action="store_true", help="Install built libraries into this checkout") + args = parser.parse_args() + if args.resolver_ub_bytes <= 512: + parser.error("--resolver-ub-bytes must leave space beyond the 512-byte descriptor") + import torch + + spec = importlib.util.find_spec("torch_npu") + if spec is None or spec.origin is None: + raise RuntimeError("torch_npu must be installed in this build environment") + npu = Path(spec.origin).parent + torch_root = Path(torch.__file__).parent + source = Path(__file__).resolve().parent + output = args.output.resolve() + output.mkdir(parents=True, exist_ok=True) + subprocess.run( + [ + str(args.cann_root / "bin/bisheng"), + "-x", + "asc", + f"--npu-arch={args.npu_arch}", + f"-DHISPARSE_RESOLVER_UB_BYTES={args.resolver_ub_bytes}", + "-O2", + "-shared", + "-fPIC", + "-std=c++17", + str(source / "resolver.asc"), + "-o", + str(output / "hisparse_resolver60103_args.so"), + ], + check=True, + ) + common = [ + "g++", + "-shared", + "-fPIC", + "-std=c++17", + "-O2", + "-D_GLIBCXX_USE_CXX11_ABI=" + str(int(torch._C._GLIBCXX_USE_CXX11_ABI)), + "-I" + str(npu / "include"), + "-I" + str(torch_root / "include"), + "-I" + str(torch_root / "include/torch/csrc/api/include"), + "-I" + str(args.cann_root / "include"), + "-L" + str(torch_root / "lib"), + "-L" + str(npu / "lib"), + "-L" + str(args.cann_root / "lib64"), + ] + subprocess.run( + common + + [ + str(source / "torch_binding.cpp"), + "-o", + str(output / "torch_binding.so"), + "-L" + str(output), + "-l:hisparse_resolver60103_args.so", + "-Wl,-rpath,$ORIGIN", + "-ltorch", + "-ltorch_cpu", + "-lc10", + "-ltorch_npu", + "-lascendcl", + ], + check=True, + ) + + if args.install: + import shutil + + package = source.parents[1] / "vllm_ascend" + for name in ("hisparse_resolver60103_args.so", "torch_binding.so"): + shutil.copy2(output / name, package / "hisparse" / name) + + +if __name__ == "__main__": + main() diff --git a/csrc/hisparse_runtime/resolver.asc b/csrc/hisparse_runtime/resolver.asc new file mode 100644 index 000000000000..f67902500fe1 --- /dev/null +++ b/csrc/hisparse_runtime/resolver.asc @@ -0,0 +1,603 @@ +// SPDX-License-Identifier: Apache-2.0 +// Request-union algorithm adapted from vLLM PR #59235 (Apache-2.0). +// SIMT implementation: one AIV per contiguous request group; one launch. +// Descriptor and scratch are invocation-owned. All mutable inputs are re-read. +#include "kernel_operator.h" +#include "simt_api/asc_simt.h" +#include + +namespace { +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; +constexpr int kTotal = kWarps; +constexpr int kError = kWarps + 2; +static_assert(kError < 64, "Prefix/error scratch exceeds reservation"); +// The build specifies the target's usable UB budget, including descriptors. +#ifndef HISPARSE_RESOLVER_UB_BYTES +#error "Build with the target's HISPARSE_RESOLVER_UB_BYTES budget" +#endif +constexpr uint64_t kScratchBudgetBytes = HISPARSE_RESOLVER_UB_BYTES; +static_assert(kScratchBudgetBytes > 512, "Resolver UB budget must cover descriptors"); + +inline uint64_t scratch_ints(uint64_t hot, uint64_t hash_size) { + const uint64_t tiles = (hot + 31) / 32; + return (2 * hash_size + hot + 2 * tiles + 64 + 15) / 16 * 16; +} +constexpr int kHit = int(0x80000000U); +constexpr int kMiss = 0x40000000; +constexpr int kSlotMask = 0xFFFF; + +__simt_callee__ inline int minimum(int a, int b) { return a < b ? a : b; } +__simt_callee__ inline int maximum(int a, int b) { return a > b ? a : b; } + +// Per-warp sums suffice for a deterministic, stable block prefix scan. +// All threads participate, including masked lanes. No cross-AIV barriers. +__simt_callee__ inline int prefix(int value, __ubuf__ int* sums, int tid) { + const int lane = tid & 31; + const int warp = tid >> 5; + for (int offset = 1; offset < 32; offset *= 2) { + int other = asc_shfl_up(value, offset, 32); + if (lane >= offset) value += other; + } + if (lane == 31) sums[warp] = value; + asc_syncthreads(); + // One warp computes all warp prefixes instead of every thread repeatedly + // loading preceding sums. Both original block barriers remain in place. + if (warp == 0) { + int total = lane < kWarps ? sums[lane] : 0; + for (int offset = 1; offset < 32; offset *= 2) { + int other = asc_shfl_up(total, offset, 32); + if (lane >= offset) total += other; + } + if (lane < kWarps) sums[lane] = total; + if (lane == kWarps - 1) sums[kTotal] = total; + } + asc_syncthreads(); + return value + (warp == 0 ? 0 : sums[warp - 1]); +} + +__simt_callee__ inline int hash_slot(int key, int size) { + const uint32_t mixed = uint32_t(key) * 2654435761U; + return (size & (size - 1)) == 0 ? int(mixed & (size - 1)) : + int((uint64_t(mixed) * uint32_t(size)) >> 32); +} + +__simt_callee__ inline void store_miss(__ubuf__ uint64_t* d, int64_t index, int value) { + if (!d[15]) return; + if (d[56]) ((__gm__ uint8_t*)d[15])[index] = uint8_t(value); + else ((__gm__ int32_t*)d[15])[index] = value; +} + +__simt_callee__ inline int lookup(__ubuf__ int* keys, int key, int size) { + int h = hash_slot(key, size); + for (int probes = 0; probes < size; ++probes) { + int old = keys[h]; + if (old == key) return h; + if (old == -1) return -1; + h = h + 1 == size ? 0 : h + 1; + } + return -1; +} + +template +__simt_callee__ inline int attention_row(__ubuf__ uint64_t* d, int id) { + // Widen before multiplying: a later int64 addition cannot fix overflow. + const int block = kQualified ? 128 : int(d[48]); + const int64_t mapped = id >= 0 ? int64_t(id / block) * int(d[51]) + id % block : -1; + return mapped >= 0 && mapped < int64_t(d[55]) ? int(mapped) : -1; +} + +// Only the optional Attention output is compacted. Refill still consumes the +// original virtual-page row IDs and the unchanged request-union swap plan. +template +__simt_callee__ inline void compact_attention( + __ubuf__ uint64_t* d, __ubuf__ int* scratch, int row, bool already_mapped = false) { + const int tid = threadIdx.x; + const int topk = int(d[33]); + const int tiles = (topk + 31) / 32; + auto physical = (__gm__ int32_t*)d[8]; + auto attention = (__gm__ int32_t*)d[9]; + auto counts = (__gm__ int32_t*)d[11]; + const int64_t source_row = int64_t(row) * topk; + const int64_t target_row = int64_t(row) * d[50]; + __ubuf__ int* values = scratch; + auto masks = (__ubuf__ uint32_t*)(values + topk); + __ubuf__ int* offsets = (__ubuf__ int*)(masks + tiles); + __ubuf__ int* sums = offsets + tiles; + // The old hash/LRU allocation covers topk + two warp tables + prefix sums. + // Stage before scattering: neither aliasing nor thread order may lose a row. + for (int base = 0; base < topk; base += kThreads) { + const int col = base + tid; + const int value = col < topk ? (already_mapped ? attention[target_row + col] : + attention_row(d, physical[source_row + col])) : -1; + if (col < topk) values[col] = value; + const uint32_t mask = asc_ballot(int(value >= 0)); + if ((tid & 31) == 0 && col < topk) masks[col >> 5] = mask; + } + asc_syncthreads(); + int total = 0; + for (int base = 0; base < tiles; base += kThreads) { + const int tile = base + tid; + const int count = tile < tiles ? __popc(masks[tile]) : 0; + const int preceding = prefix(count, sums, tid); + const int tile_total = sums[kTotal]; + if (tile < tiles) offsets[tile] = total + preceding - count; + asc_syncthreads(); + total += tile_total; + } + for (int col = tid; col < topk; col += kThreads) { + const int value = values[col]; + if (value >= 0) { + const uint32_t bit = uint32_t(1) << (col & 31); + const int destination = offsets[col >> 5] + __popc(masks[col >> 5] & (bit - 1)); + attention[target_row + destination] = value; + } + } + for (int col = total + tid; col < topk; col += kThreads) attention[target_row + col] = -1; + if (tid == 0 && counts) counts[int64_t(row) * d[52]] = total; + asc_syncthreads(); +} + +template +__simt_vf__ __launch_bounds__(kThreads) inline void resolve_vf( + __ubuf__ uint64_t* d, __ubuf__ int* scratch, int first) { + const int tid = threadIdx.x; + auto logical = (__gm__ int32_t*)d[0]; + auto requests = (__gm__ int32_t*)d[1]; + auto states = (__gm__ int32_t*)d[2]; + auto source = (__gm__ int32_t*)d[3]; + auto resident = (__gm__ int32_t*)d[4]; + auto table = (__gm__ int32_t*)d[5]; + auto dgi = (__gm__ int32_t*)d[6]; + auto lru = (__gm__ int16_t*)d[7]; + auto output = (__gm__ int32_t*)d[8]; + auto attention = (__gm__ int32_t*)d[9]; + auto resolved = (__gm__ int32_t*)d[10]; + auto counts = (__gm__ int32_t*)d[11]; + auto swap_host = (__gm__ int32_t*)d[12]; + auto swap_device = (__gm__ int32_t*)d[13]; + auto swap_counts = (__gm__ int32_t*)d[14]; + auto stats = (__gm__ int64_t*)d[16]; + auto errors = (__gm__ int32_t*)d[17]; + auto num_valid_rows = (__gm__ int32_t*)d[18]; + const int rows = d[32]; + const int topk = kQualified ? 2048 : int(d[33]); + const int hot = kQualified ? 10240 : int(d[34]); + const int hash_size = kQualified ? 16384 : int(d[54]); + const bool compact = d[55] != 0; + const int valid_rows = num_valid_rows ? minimum(rows, maximum(0, num_valid_rows[0])) : rows; + if (first >= valid_rows) { + if (tid < 8 && errors) errors[first * 8 + tid] = 0; + for (int i = tid; i < topk; i += kThreads) { + const int64_t index = int64_t(first) * topk + i; + output[index] = -1; + if (attention) attention[int64_t(first) * d[50] + i] = -1; + if (resolved) resolved[index] = -1; + store_miss(d, index, 0); + } + if (tid == 0) { + if (counts) counts[int64_t(first) * d[52]] = 0; + if (swap_counts) swap_counts[first] = 0; + } + return; + } + int request = requests ? requests[first] : first; + int state = states ? ((request >= 0 && request < int(d[37])) ? states[request] : -1) : request; + bool state_valid = state >= 0 && state < int(d[38]); + // Source/Resident tables are immutable for this invocation. A complete + // device-table proof lets every AIV translate its own row independently; + // a mixed request still uses the original leader-owned union/LRU path. + if (kQualified && source && resident && requests && attention && num_valid_rows && + int(d[50]) == topk && state_valid && state < int(d[43]) && + request >= 0 && request < int(d[39]) && int(d[40]) > 0 && int(d[44]) > 0) { + const int host0 = source[int64_t(request) * d[41]]; + const int page0 = resident[int64_t(state) * d[45]]; + // The ordinary cold case returns to the old path without a page scan. + if (host0 <= 0 || (page0 >= 0 && page0 != int(d[47]))) { + int missing = 0; + for (int col = tid; col < int(d[40]); col += kThreads) { + const int host_page = source[int64_t(request) * d[41] + col]; + const int page = col < int(d[44]) ? resident[int64_t(state) * d[45] + col] : int(d[47]); + missing += host_page > 0 && (page < 0 || page > 0xFFFFFF || page == int(d[47])); + } + prefix(missing, scratch, tid); + const bool all_resident = scratch[kTotal] == 0; + if (compact && tid == 0) scratch[kError] = 0; + // Every lane must finish reading the proof before fallback code + // can repurpose this AIV's scratch for hash keys. + asc_syncthreads(); + if (all_resident) { + int valid_count = 0; + bool attention_invalid = false; + for (int i = tid; i < topk; i += kThreads) { + const int token = logical[int64_t(first) * d[36] + i]; + const int col = token >= 0 ? token / 128 : -1; + const int host_page = token >= 0 && col < int(d[40]) ? source[int64_t(request) * d[41] + col] : 0; + int host = token >= 0 && host_page > 0 ? host_page * 128 + token % 128 : -1; + const int page = token >= 0 && col < int(d[44]) ? resident[int64_t(state) * d[45] + col] : int(d[47]); + const int device = token >= 0 && page >= 0 && page != int(d[47]) ? page * 128 + token % 128 : -1; + if (host < 0 || int64_t(host) >= int64_t(d[35])) host = -1; + const int64_t index = int64_t(first) * topk + i; + output[index] = device; + const int mapped = compact ? attention_row(d, device) : + (device >= 0 ? device / 128 * int(d[51]) + device % 128 : -1); + attention[int64_t(first) * d[50] + i] = mapped; + attention_invalid |= mapped < 0; + if (resolved) resolved[index] = host; + store_miss(d, index, 0); + valid_count += device >= 0 || host >= 0; + } + if (compact) { + if (attention_invalid) asc_atomic_or(scratch + kError, 1); + asc_syncthreads(); + if (scratch[kError]) { + asc_syncthreads(); + compact_attention(d, scratch, first, true); + } else if (tid == 0 && counts) { + counts[int64_t(first) * d[52]] = topk; + } + } else if (counts) { + prefix(valid_count, scratch, tid); + if (tid == 0) counts[int64_t(first) * d[52]] = scratch[kTotal]; + } + if (tid == 0 && swap_counts) swap_counts[first] = 0; + // Nonleader error rows were untouched by the original path. + if (errors && (first == 0 || requests[first - 1] != request) && tid < 8) errors[first * 8 + tid] = 0; + return; + } + } + } + if (requests && first > 0 && requests[first - 1] == request) return; + int end = first + 1; + while (requests && end < valid_rows && requests[end] == request) ++end; + __ubuf__ int* keys = scratch; + __ubuf__ int* values = keys + hash_size; + __ubuf__ int* temp = values + hash_size; + __ubuf__ uint32_t* masks = (__ubuf__ uint32_t*)(temp + hot); + __ubuf__ int* offsets = (__ubuf__ int*)(masks + (hot + 31) / 32); + __ubuf__ int* sums = offsets + (hot + 31) / 32; + if (tid < kWarps + 8) sums[tid] = 0; + if (tid < 8 && errors) errors[first * 8 + tid] = 0; + asc_syncthreads(); + int cold_references = 0; + for (int row = first; row < end; ++row) { + int valid_count = 0; + for (int i = tid; i < topk; i += kThreads) { + int token = logical[int64_t(row) * d[36] + i]; + int host = token, device = -1; + if (source) { + int col = token >= 0 ? token / int(d[42]) : -1; + bool valid = state_valid && token >= 0 && request >= 0 && request < int(d[39]) && col < int(d[40]); + int page = valid ? source[int64_t(request) * d[41] + col] : 0; + host = valid && page > 0 ? page * int(d[42]) + token % int(d[42]) : -1; + if (resident) { + col = token >= 0 ? token / int(d[46]) : -1; + valid = state_valid && token >= 0 && state < int(d[43]) && col < int(d[44]); + page = valid ? resident[int64_t(state) * d[45] + col] : int(d[47]); + device = valid && page >= 0 && page != int(d[47]) ? page * int(d[46]) + token % int(d[46]) : -1; + } + } + if (!state_valid || host < 0 || int64_t(host) >= int64_t(d[35])) host = -1; + if (!state_valid) device = -1; + int64_t index = int64_t(row) * topk + i; + output[index] = device; + if (resolved) resolved[index] = host; + store_miss(d, index, 0); + valid_count += device >= 0 || host >= 0; + if (device < 0 && host >= 0) { + ++cold_references; + + } + } + if (counts && !compact) { + prefix(valid_count, sums, tid); + if (tid == 0) counts[int64_t(row) * d[52]] = sums[kTotal]; + asc_syncthreads(); + } + } + prefix(cold_references, sums, tid); + int total_references = sums[kTotal]; + asc_syncthreads(); + if (!state_valid || total_references == 0 || sums[kError]) { + const bool failed = sums[kError] != 0; + if (compact) asc_syncthreads(); + for (int row = first; row < end; ++row) { + if (tid == 0 && swap_counts) swap_counts[row] = 0; + if (compact) { + compact_attention(d, scratch, row); + } else for (int i = tid; i < topk; i += kThreads) { + int value = output[int64_t(row) * topk + i]; + if (attention) attention[int64_t(row) * d[50] + i] = + value >= 0 ? value / int(d[48]) * int(d[51]) + value % int(d[48]) : -1; + } + } + if (tid == 0 && failed && errors) errors[first * 8] = -1; + return; + } + // Match upstream's Resident-first path: initialize the union only when + // the request actually needs Host/Hot resolution. Re-read translated rows; + // resolved is invocation-owned and always bound by the wrapper. + for (int i = tid; i < hash_size; i += kThreads) { + keys[i] = -1; + values[i] = 0x7FFFFFFF; + } + asc_syncthreads(); + for (int row = first; row < end; ++row) { + for (int i = tid; i < topk; i += kThreads) { + const int64_t index = int64_t(row) * topk + i; + const int host = resolved[index]; + if (output[index] < 0 && host >= 0) { + int h = hash_slot(host, hash_size); + bool inserted = false; + for (int probe = 0; probe < hash_size; ++probe) { + // Keys only transition from empty to an immutable host + // ID during insertion. Occupied reads need no CAS. Values + // only decrease, so a settled earlier position needs no min. + int old = keys[h]; + if (old == -1) old = asc_atomic_cas(keys + h, -1, host); + if (old == -1 || old == host) { + int position = (row - first) * topk + i; + if (values[h] > position) asc_atomic_min(values + h, position); + inserted = true; + break; + } + h = h + 1 == hash_size ? 0 : h + 1; + } + if (!inserted) asc_atomic_or(sums + kError, 1); + } + } + } + asc_syncthreads(); + if (sums[kError]) { + if (tid == 0 && errors) errors[first * 8] = -1; + if (compact) asc_syncthreads(); + for (int row = first; row < end; ++row) { + if (tid == 0 && swap_counts) swap_counts[row] = 0; + if (compact) { + compact_attention(d, scratch, row); + } else for (int i = tid; i < topk; i += kThreads) { + const int value = output[int64_t(row) * topk + i]; + if (attention) attention[int64_t(row) * d[50] + i] = + value >= 0 ? value / int(d[48]) * int(d[51]) + value % int(d[48]) : -1; + } + } + return; + } + // Publish all hits before choosing ANY eviction. Duplicate cached host IDs + // protect every copy; the minimum slot is selected, matching sorted keys. + for (int slot = tid; slot < hot; slot += kThreads) { + int host = dgi[int64_t(state) * hot + slot]; + int h = host >= 0 ? lookup(keys, host, hash_size) : -1; + // Save membership while every Hot slot is visited. Reuse the existing + // partition scratch only until all LRU membership masks are complete. + temp[slot] = int(h >= 0); + if (h >= 0) asc_atomic_min(values + h, kHit | slot); + } + asc_syncthreads(); + // Compute all warp hit masks first, then scan the tile counts once. + // Preserve the exact stable LRU order while avoiding one block scan per + // 512-slot tile. No cross-AIV state or graph-lifetime changes. + for (int base = 0; base < hot; base += kThreads) { + int p = base + tid; + int slot = p < hot ? int(lru[int64_t(state) * hot + p]) : -1; + bool member = slot >= 0 && slot < hot && temp[slot] != 0; + uint32_t mask = asc_ballot(int(p < hot && member)); + if ((tid & 31) == 0 && p < hot) { + masks[p >> 5] = mask; + offsets[p >> 5] = __popc(mask); + } + } + asc_syncthreads(); + int tiles = (hot + 31) / 32; + int hit_count = 0; + for (int base = 0; base < tiles; base += kThreads) { + int tile = base + tid; + int count = tile < tiles ? offsets[tile] : 0; + int preceding = prefix(count, sums, tid); + int tile_total = sums[kTotal]; + if (tile < tiles) offsets[tile] = hit_count + preceding - count; + asc_syncthreads(); + hit_count += tile_total; + } + for (int p = tid; p < hot; p += kThreads) { + int slot = int(lru[int64_t(state) * hot + p]); + uint32_t mask = masks[p >> 5]; + uint32_t bit = uint32_t(1) << (p & 31); + bool hit = (mask & bit) != 0; + int rank = offsets[p >> 5] + __popc(mask & (bit - 1)) + int(hit); + temp[hit ? rank - 1 : hot - 1 - (p - rank)] = slot; + } + asc_syncthreads(); + // Count distinct misses first: bounds failure must not partially mutate LRU. + int own_misses = 0; + for (int h = tid; h < hash_size; h += kThreads) own_misses += keys[h] >= 0 && values[h] >= 0; + prefix(own_misses, sums, tid); + int total_misses = sums[kTotal]; + asc_syncthreads(); + if (hit_count + total_misses > int(d[53]) || total_misses > hot - hit_count) { + if (tid == 0 && errors) errors[first * 8] = hit_count + total_misses; + for (int row = first; row < end; ++row) { + if (tid == 0 && swap_counts) swap_counts[row] = 0; + if (compact && tid == 0 && counts) counts[int64_t(row) * d[52]] = 0; + for (int i = tid; i < topk; i += kThreads) { + output[int64_t(row) * topk + i] = -1; + if (attention) attention[int64_t(row) * d[50] + i] = -1; + } + } + return; + } + // Scan all first-reference miss flags before changing hash values. + // Stable global-position offsets preserve duplicate ordering and victims. + int positions = (end - first) * topk; + if (positions <= minimum(tiles, kThreads) * 32) { + for (int base = 0; base < positions; base += kThreads) { + int position = base + tid; + bool active = position < positions; + int row = first + position / topk, col = position % topk; + int host = active && output[int64_t(row) * topk + col] < 0 ? resolved[int64_t(row) * topk + col] : -1; + int h = host >= 0 ? lookup(keys, host, hash_size) : -1; + uint32_t mask = asc_ballot(int(h >= 0 && values[h] == position)); + if ((tid & 31) == 0 && active) { + masks[position >> 5] = mask; + offsets[position >> 5] = __popc(mask); + } + } + asc_syncthreads(); + int tiles = (positions + 31) / 32; + int preceding = prefix(tid < tiles ? offsets[tid] : 0, sums, tid); + if (tid < tiles) offsets[tid] = preceding - __popc(masks[tid]); + asc_syncthreads(); + for (int position = tid; position < positions; position += kThreads) { + uint32_t mask = masks[position >> 5]; + uint32_t bit = uint32_t(1) << (position & 31); + if (mask & bit) { + int rank = offsets[position >> 5] + __popc(mask & (bit - 1)); + int row = first + position / topk, col = position % topk; + int host = resolved[int64_t(row) * topk + col]; + int h = lookup(keys, host, hash_size); + int slot = temp[hot - 1 - rank]; + int block = table[int64_t(request) * d[49] + slot / int(d[48])]; + int physical = block * int(d[48]) + slot % int(d[48]); + values[h] = kMiss | slot; + dgi[int64_t(state) * hot + slot] = host; + if (swap_host) swap_host[int64_t(first) * topk + rank] = host; + if (swap_device) swap_device[int64_t(first) * topk + rank] = physical; + store_miss(d, int64_t(row) * topk + col, 1); + } + } + asc_syncthreads(); + } else { + int consumed = 0; + for (int base = 0; base < (end - first) * topk; base += kThreads) { + int position = base + tid; + bool active = position < (end - first) * topk; + int row = first + position / topk, col = position % topk; + int host = active && output[int64_t(row) * topk + col] < 0 ? resolved[int64_t(row) * topk + col] : -1; + int h = host >= 0 ? lookup(keys, host, hash_size) : -1; + bool miss = h >= 0 && values[h] == position; + int rank = prefix(miss, sums, tid) - 1 + consumed; + int tile_misses = sums[kTotal]; + if (miss) { + int slot = temp[hot - 1 - rank]; + int block = table[int64_t(request) * d[49] + slot / int(d[48])]; + int physical = block * int(d[48]) + slot % int(d[48]); + values[h] = kMiss | slot; + dgi[int64_t(state) * hot + slot] = host; + if (swap_host) swap_host[int64_t(first) * topk + rank] = host; + if (swap_device) swap_device[int64_t(first) * topk + rank] = physical; + store_miss(d, int64_t(row) * topk + col, 1); + } + asc_syncthreads(); + consumed += tile_misses; + } + } + int remaining = hot - hit_count - total_misses; + for (int p = tid; p < hot; p += kThreads) { + int source_pos = p < remaining ? hot - 1 - total_misses - p : + (p < remaining + total_misses ? hot - 1 - (p - remaining) : p - remaining - total_misses); + lru[int64_t(state) * hot + p] = int16_t(temp[source_pos]); + } + // The successful union path leaves kError clear. Its final barrier below + // also retires all LRU readers before any compaction reuses scratch. + bool attention_invalid = false; + for (int row = first; row < end; ++row) { + if (tid == 0 && swap_counts) swap_counts[row] = minimum(topk, maximum(0, total_misses - (row - first) * topk)); + for (int i = tid; i < topk; i += kThreads) { + int64_t index = int64_t(row) * topk + i; + int physical = output[index]; + int host = resolved[index]; + if (physical < 0 && host >= 0) { + int h = lookup(keys, host, hash_size); + int slot = values[h] & kSlotMask; + int block = table[int64_t(request) * d[49] + slot / int(d[48])]; + physical = block * int(d[48]) + slot % int(d[48]); + } + output[index] = physical; + if (attention) { + const int mapped = compact ? attention_row(d, physical) : + (physical >= 0 ? physical / int(d[48]) * int(d[51]) + physical % int(d[48]) : -1); + attention[int64_t(row) * d[50] + i] = mapped; + attention_invalid |= mapped < 0; + } + } + } + if (compact) { + if (attention_invalid) asc_atomic_or(sums + kError, 1); + asc_syncthreads(); + if (sums[kError]) { + asc_syncthreads(); + for (int row = first; row < end; ++row) compact_attention(d, scratch, row, true); + } else if (tid == 0 && counts) { + for (int row = first; row < end; ++row) counts[int64_t(row) * d[52]] = topk; + } + } + if (tid == 0 && stats) { + asc_atomic_add(stats, int64_t(total_references - total_misses)); + asc_atomic_add(stats + 1, int64_t(total_misses)); + } +} +} // namespace + +template +__global__ __vector__ void hisparse_union_resolve60103(uint64_t a0, uint64_t a1, uint64_t a2, uint64_t a3, uint64_t a4, uint64_t a5, uint64_t a6, uint64_t a7, uint64_t a8, uint64_t a9, uint64_t a10, uint64_t a11, uint64_t a12, uint64_t a13, uint64_t a14, uint64_t a15, uint64_t a16, uint64_t a17, uint64_t a18, uint64_t a32, uint64_t a33, uint64_t a34, uint64_t a35, uint64_t a36, uint64_t a37, uint64_t a38, uint64_t a39, uint64_t a40, uint64_t a41, uint64_t a42, uint64_t a43, uint64_t a44, uint64_t a45, uint64_t a46, uint64_t a47, uint64_t a48, uint64_t a49, uint64_t a50, uint64_t a51, uint64_t a52, uint64_t a53, uint64_t a54, uint64_t a55, uint64_t a56) { + extern __ubuf__ int32_t local_scratch[]; + const uint64_t words = (2 * a54 + a34 + 2 * ((a34 + 31) / 32) + 64 + 15) / 16 * 16; + auto descriptor = (__ubuf__ uint64_t*)(local_scratch + words); + descriptor[0] = a0; + descriptor[1] = a1; + descriptor[2] = a2; + descriptor[3] = a3; + descriptor[4] = a4; + descriptor[5] = a5; + descriptor[6] = a6; + descriptor[7] = a7; + descriptor[8] = a8; + descriptor[9] = a9; + descriptor[10] = a10; + descriptor[11] = a11; + descriptor[12] = a12; + descriptor[13] = a13; + descriptor[14] = a14; + descriptor[15] = a15; + descriptor[16] = a16; + descriptor[17] = a17; + descriptor[18] = a18; + descriptor[32] = a32; + descriptor[33] = a33; + descriptor[34] = a34; + descriptor[35] = a35; + descriptor[36] = a36; + descriptor[37] = a37; + descriptor[38] = a38; + descriptor[39] = a39; + descriptor[40] = a40; + descriptor[41] = a41; + descriptor[42] = a42; + descriptor[43] = a43; + descriptor[44] = a44; + descriptor[45] = a45; + descriptor[46] = a46; + descriptor[47] = a47; + descriptor[48] = a48; + descriptor[49] = a49; + descriptor[50] = a50; + descriptor[51] = a51; + descriptor[52] = a52; + descriptor[53] = a53; + descriptor[54] = a54; + descriptor[55] = a55; + descriptor[56] = a56; + asc_vf_call>(dim3(kThreads), descriptor, local_scratch, int(AscendC::GetBlockIdx())); +} + +extern "C" uint64_t hisparse_resolver_scratch_capacity() { + return kScratchBudgetBytes; +} + +extern "C" void launch_hisparse_union_resolve60103(void** pointers, uint64_t* geometry, int rows, void* stream) { + if (geometry[1] == 2048 && geometry[2] == 10240 && geometry[10] == 128 && geometry[14] == 128 && geometry[16] == 128 && geometry[22] == 16384) { + hisparse_union_resolve60103<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); + } else { + hisparse_union_resolve60103<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); + } +} diff --git a/csrc/hisparse_runtime/torch_binding.cpp b/csrc/hisparse_runtime/torch_binding.cpp new file mode 100644 index 000000000000..6db1d8af9b55 --- /dev/null +++ b/csrc/hisparse_runtime/torch_binding.cpp @@ -0,0 +1,205 @@ +// SPDX-License-Identifier: Apache-2.0 +// Tensor metadata is decoded here; launches keep TorchNPU queue ordering. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" void launch_hisparse_union_resolve60103(void**, uint64_t*, int, void*); +extern "C" uint64_t hisparse_resolver_scratch_capacity(); + +namespace { +at::Tensor tensor_or_undefined(const std::optional& tensor) { + return tensor.value_or(at::Tensor()); +} + +void resolve_tensors( + const at::Tensor& logical, + const std::optional& requests, + const std::optional& states, + const std::optional& source, + const std::optional& resident, + const at::Tensor& hot_table, + const at::Tensor& dgi, + const at::Tensor& lru, + const at::Tensor& output, + const std::optional& attention, + const at::Tensor& resolved, + const std::optional& counts, + const std::optional& swap_host, + const std::optional& swap_device, + const std::optional& swap_counts, + const std::optional& miss_mask, + const std::optional& stats, + const at::Tensor& errors, + const std::optional& valid_rows, + int64_t host_rows, int64_t hot_block_size, int64_t source_block_size, + int64_t resident_block_size, int64_t resident_null_block, + int64_t attention_block_stride, int64_t union_rows, int64_t attention_capacity) { + TORCH_CHECK(logical.dim() == 2 && dgi.dim() == 2 && hot_table.dim() == 2, + "HiSparse requires two-dimensional TopK, hot tags and block table"); + const auto rows = logical.size(0), topk = logical.size(1), hot = dgi.size(1); + TORCH_CHECK(rows > 0 && topk > 0 && rows <= ((int64_t(1) << 30) - 1) / topk, + "HiSparse first-reference positions exceed the int32 tag representation"); + TORCH_CHECK(union_rows > 0 && union_rows <= hot && hot <= 32768 && topk <= hot, + "HiSparse union must fit the hot cache and int16 LRU slots"); + TORCH_CHECK(host_rows > 0 && host_rows < (int64_t(1) << 31), "Invalid HiSparse host row capacity"); + TORCH_CHECK(hot_block_size > 0 && attention_block_stride >= hot_block_size, + "Invalid hot/attention block strides"); + TORCH_CHECK(!source || source_block_size > 0, "Invalid source block size"); + TORCH_CHECK(!resident || resident_block_size > 0, "Invalid resident block size"); + TORCH_CHECK(attention_capacity >= 0 && attention_capacity < (int64_t(1) << 31), + "HiSparse compact Attention capacity must fit int32"); + TORCH_CHECK(attention_capacity == 0 || attention, + "HiSparse compact Attention requires an output tensor"); + TORCH_CHECK(attention_capacity == 0 || !attention->is_alias_of(output), + "HiSparse compact Attention must not overwrite virtual refill indices"); + TORCH_CHECK(!states || states->dim() == 1, "HiSparse state indices must be a vector"); + TORCH_CHECK(!source || source->dim() == 2, "HiSparse source table must be a matrix"); + TORCH_CHECK(!resident || resident->dim() == 2, "HiSparse resident table must be a matrix"); + TORCH_CHECK(lru.sizes() == dgi.sizes(), "HiSparse LRU/tag geometry mismatch"); + int64_t hash_size = 1; + while (hash_size <= union_rows) hash_size <<= 1; + const uint64_t tiles = (hot + 31) / 32; + const uint64_t budget = hisparse_resolver_scratch_capacity(); + const uint64_t available_words = (budget - 512) / 64 * 16; + const uint64_t fixed_words = hot + 2 * tiles + 64; + const uint64_t minimum_bytes = (2 * (union_rows + 1) + fixed_words + 15) / 16 * 64 + 512; + TORCH_CHECK(minimum_bytes <= budget, + "HiSparse hot_size=", hot, ", max_union_rows=", union_rows, + " needs at least ", minimum_bytes, " UB bytes; the compiled target allows ", budget, + ". Lower device_buffer_size or speculative tokens."); + hash_size = std::min(hash_size, int64_t((available_words - fixed_words) / 2)); + + const std::array tensors = { + logical, tensor_or_undefined(requests), tensor_or_undefined(states), + tensor_or_undefined(source), tensor_or_undefined(resident), hot_table, + dgi, lru, output, tensor_or_undefined(attention), resolved, + tensor_or_undefined(counts), tensor_or_undefined(swap_host), + tensor_or_undefined(swap_device), tensor_or_undefined(swap_counts), + tensor_or_undefined(miss_mask), tensor_or_undefined(stats), errors, + tensor_or_undefined(valid_rows)}; + c10_npu::NPUGuard guard(logical.device()); + auto npu_stream = c10_npu::getCurrentNPUStream(); + std::array pointers{}; + std::vector owners; + owners.reserve(tensors.size()); + for (size_t i = 0; i < tensors.size(); ++i) { + const auto& tensor = tensors[i]; + if (!tensor.defined()) continue; + TORCH_CHECK(tensor.device() == logical.device(), "HiSparse resolver device mismatch"); + const auto dtype = i == 7 ? at::kShort : (i == 16 ? at::kLong : at::kInt); + if (i == 15) { + TORCH_CHECK(tensor.scalar_type() == at::kInt || tensor.scalar_type() == at::kByte || + tensor.scalar_type() == at::kBool, "HiSparse miss mask must be int32, uint8 or bool"); + } else { + TORCH_CHECK(tensor.scalar_type() == dtype, "HiSparse resolver dtype mismatch at ", i); + } + if (tensor.dim() == 2) { + TORCH_CHECK(tensor.stride(1) == 1 && tensor.stride(0) >= tensor.size(1), + "HiSparse resolver matrix rows must be dense and non-overlapping"); + } + const bool dense_topk = i == 8 || i == 10 || i == 12 || i == 13; + if (i == 2 || i == 6 || i == 7) { + TORCH_CHECK(tensor.is_contiguous(), "HiSparse state/LRU/tag storage must be contiguous"); + } + if (dense_topk) { + TORCH_CHECK(tensor.dim() == 2 && tensor.size(0) == rows && tensor.size(1) == topk && tensor.is_contiguous(), + "HiSparse resolver requires contiguous TopK-shaped output at ", i); + } + if (i == 9) { + TORCH_CHECK(tensor.dim() == 2 && tensor.size(0) == rows && tensor.size(1) == topk, + "Invalid HiSparse attention index shape"); + } + if (i == 1 || i == 14) { + TORCH_CHECK(tensor.dim() == 1 && tensor.numel() == rows && tensor.is_contiguous(), + "Invalid HiSparse per-row vector at ", i); + } + if (i == 11) { + TORCH_CHECK((tensor.dim() == 1 || (tensor.dim() == 2 && tensor.size(1) == 1)) && tensor.size(0) == rows, + "Invalid HiSparse valid-count shape"); + TORCH_CHECK(tensor.stride(0) > 0, "HiSparse valid-count rows must not overlap"); + } + if (i == 15) { + TORCH_CHECK(tensor.dim() == 2 && tensor.size(0) == rows && tensor.size(1) == topk && tensor.is_contiguous(), + "Invalid HiSparse miss-mask shape"); + } + if (i == 16) TORCH_CHECK(tensor.numel() == 2 && tensor.is_contiguous(), "Invalid HiSparse stats shape"); + if (i == 17) { + TORCH_CHECK(tensor.dim() == 2 && tensor.size(0) == rows && tensor.size(1) == 8 && tensor.is_contiguous(), + "Invalid HiSparse error scratch shape"); + } + if (i == 18) TORCH_CHECK(tensor.numel() == 1, "Invalid HiSparse valid-row scalar"); + // Separate TensorImpls retain storage until queue submission without + // Python-object destruction on a queue/release thread. + owners.emplace_back(tensor.detach()); + TORCH_CHECK(owners.back().unsafeGetTensorImpl()->pyobj_slot()->load_pyobj() == nullptr && + owners.back().unsafeGetTensorImpl() != tensor.unsafeGetTensorImpl(), + "HiSparse queue owner must have an independent, Python-free TensorImpl"); + // Queue callback completion is not device completion. The allocator + // must keep every storage unavailable for reuse until this stream's + // work completes, including when graph capture is active. + c10_npu::NPUCachingAllocator::recordStream(tensor.storage().data_ptr(), npu_stream); + pointers[i] = tensor.numel() ? tensor.data_ptr() : nullptr; + } + TORCH_CHECK(bool(swap_host) == bool(swap_device) && bool(swap_host) == bool(swap_counts), + "All compact HiSparse swap outputs are required together"); + TORCH_CHECK(hot_table.size(1) >= (hot + hot_block_size - 1) / hot_block_size, + "HiSparse hot block table is too narrow"); + TORCH_CHECK(hot_table.size(0) >= (states ? states->numel() : dgi.size(0)), + "HiSparse hot block table does not cover request slots"); + std::array sizes = { + uint64_t(rows), uint64_t(topk), uint64_t(hot), uint64_t(host_rows), uint64_t(logical.stride(0)), + uint64_t(states ? states->numel() : 0), uint64_t(dgi.size(0)), + uint64_t(source ? source->size(0) : 0), uint64_t(source ? source->size(1) : 0), + uint64_t(source ? source->stride(0) : 0), uint64_t(source_block_size), + uint64_t(resident ? resident->size(0) : 0), uint64_t(resident ? resident->size(1) : 0), + uint64_t(resident ? resident->stride(0) : 0), uint64_t(resident_block_size), uint64_t(resident_null_block), + uint64_t(hot_block_size), uint64_t(hot_table.stride(0)), + uint64_t(attention ? attention->stride(0) : topk), uint64_t(attention_block_stride), + uint64_t(counts ? counts->stride(0) : 1), uint64_t(union_rows), uint64_t(hash_size), + uint64_t(attention_capacity), uint64_t(miss_mask && miss_mask->element_size() == 1)}; + auto stream = npu_stream.stream(); + at_npu::native::OpCommand command; + command.Name("HiSparseUnionResolve"); + command.SetCustomHandler([pointers, sizes, stream, owners = std::move(owners)]() mutable -> int { + // Release queued tensor owners after submission, not on queue-slot reuse. + auto submission_owners = std::move(owners); + launch_hisparse_union_resolve60103(pointers.data(), sizes.data(), int(sizes[0]), stream); + return 0; + }); + command.Run(); +} + +// The operator mutates caller-owned tensors and returns no new tensors. +void resolve_meta( + const at::Tensor&, const std::optional&, const std::optional&, + const std::optional&, const std::optional&, const at::Tensor&, + const at::Tensor&, const at::Tensor&, const at::Tensor&, const std::optional&, + const at::Tensor&, const std::optional&, const std::optional&, + const std::optional&, const std::optional&, const std::optional&, + const std::optional&, const at::Tensor&, const std::optional&, + int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t) {} +} // namespace + +TORCH_LIBRARY(hisparse_fused60103, m) { + m.def("resolve_tensors(Tensor logical, Tensor? requests, Tensor? states, Tensor? source, Tensor? resident, " + "Tensor hot_table, Tensor(a!) dgi, Tensor(a!) lru, Tensor(a!) output, Tensor(a!)? attention, " + "Tensor(a!) resolved, Tensor(a!)? counts, Tensor(a!)? swap_host, Tensor(a!)? swap_device, " + "Tensor(a!)? swap_counts, Tensor(a!)? miss_mask, Tensor(a!)? stats, Tensor(a!) errors, Tensor? valid_rows, " + "int host_rows, int hot_block_size, int source_block_size, int resident_block_size, int resident_null_block, " + "int attention_block_stride, int union_rows, int attention_capacity=0) -> ()"); + m.def("scratch_capacity() -> int", []() -> int64_t { return hisparse_resolver_scratch_capacity(); }); + m.def("supports_attention_compaction() -> bool", []() -> bool { return true; }); + m.def("supports_budgeted_hash() -> bool", []() -> bool { return true; }); +} +TORCH_LIBRARY_IMPL(hisparse_fused60103, PrivateUse1, m) { m.impl("resolve_tensors", TORCH_FN(resolve_tensors)); } +TORCH_LIBRARY_IMPL(hisparse_fused60103, Meta, m) { m.impl("resolve_tensors", TORCH_FN(resolve_meta)); } diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index bed9e6fb44c7..b92397ff7e47 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -1,252 +1,96 @@ -# HiSparse foundation on Ascend - -## Scope - -This is the foundational split of [#19109](https://github.com/vllm-project/vllm-ascend/pull/19109). -It supplies an Ascend backend for the merged vLLM HiSparse API and the minimal -SFA data path. It is a contribution starting point with NPU component checks -and a complete A5 model smoke test. Dataset accuracy and performance qualification -remain pending. Historical optimized C8 performance does not qualify this split. - -Core owns HMA allocation, Coordinator admission/eviction, prefix-cache state, -request lifetimes, prefetch control flow, mirror completion and release/rollback. -Ascend implements device operations and cache-layout adaptation. Attention uses -the existing CANN sparse-flash-attention and fused RMSNorm/RoPE operations; -this PR adds no Attention core or native C++/AscendC kernel. - -```text -Core sparse index group -> logical TopK -> NPU residency/refill plan - | - Resident / Hot hit / mapped-Host miss - | - packed KV -> existing native SFA operation -``` - -The data path covers unquantized KV and the native Ascend 950 C8 format with -the RoPE SFA layout and Model Runner V2. Model-weight quantization is independent. Private -Host pools are used per rank because NPU IPC events are not implemented. -The mapped-Host capability is checked at registration; hardware evidence for -the current primitives is from Ascend 950. A3 still needs its own qualification. +# HiSparse on Ascend -The SFA kernel requires TopK divisible by 128. The resolver's int16 LRU format -limits aligned Hot capacity to 32768 rows, and Hot must cover the complete -verification-row union. Context parallelism is not implemented in this split. -Request counts, lengths, layer/group counts and scratch capacities come from -configuration, rather than benchmark-specific constants. +HiSparse keeps active sparse-attention KV rows in NPU Resident/Hot storage and +refills selected missing rows from Host memory. Core owns HMA allocation, +Coordinator admission/eviction, request lifetimes, prefetch and mirror completion. +The Ascend backend supplies residency resolution, transfers and the existing SFA +Attention interface. Attention core arithmetic is unchanged. -## Configuration +## Supported configurations -Use the matching merged Core pin in `.github/vllm-main-verified.commit`. -No unmerged Core PR, copied scheduler, experiment profile or source-file hash -allowlist is required. As with other Ascend backends, paired-version ABI changes -still need review and device validation. +- Experimental sparse RoPE MLA/SFA with Model Runner V2. +- BF16 KV and native Ascend 950 C8: E4M3FN main/Indexer values with FP32 scales. + Main and Indexer cache dtypes are independently selectable. +- Private Host pools per tensor-parallel rank and Core eager Host mirroring. +- Ascend 950 with a SIMT-capable CANN 9.2 toolchain. A2/A3 are not newly qualified. +- TopK must satisfy the existing SFA requirement of divisibility by 128. -Set `VLLM_USE_V2_MODEL_RUNNER=1` and add these options to a compatible model's -normal launch command: +C4, TurboQuant, INT8-format C8, context parallelism and external P/D transport +are outside this backend's current support. -```bash ---block-size 128 \ ---kv-cache-dtype bfloat16 \ ---attention-config '{"indexer_kv_dtype":"bf16","hisparse_config":{"eager_host_mirror":true}}' \ ---kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":4}}' -``` +## Build the resolver -The example reserves 4 GiB Host storage per rank. Size it for the model and context -budget; it is not a model capacity claim. Keep Core's eager mirror policy for -this foundation. External P/D transfer is not provided by this PR. - -For native A5 C8, replace the two dtype options above with: +Use the paired merged vLLM version in `.github/vllm-main-verified.commit` and +install vllm-ascend normally. Its CMake build includes the fused resolver when +the Ascend 950 SIMT headers are available. A separate resolver rebuild in the +same Torch/torch-npu/CANN environment is also supported: ```bash ---kv-cache-dtype fp8 \ ---attention-config '{"indexer_kv_dtype":"fp8","hisparse_config":{"eager_host_mirror":true}}' +python csrc/hisparse_runtime/build.py \ + --cann-root "$ASCEND_HOME_PATH" \ + --output build/hisparse-resolver \ + --npu-arch dav-3510 \ + --resolver-ub-bytes 175616 \ + --install ``` -The existing SFA quantizer produces each main KV row: 512 E4M3FN NoPE bytes, -64 BF16 RoPE values and four FP32 block scales, totaling 656 bytes. The adapter -copies these components as opaque bytes and calls the existing quantized SFA -Attention operation. Indexer C8 uses 128 E4M3FN keys and one FP32 scale per token. -Main and Indexer dtypes remain independently selectable. C4, TurboQuant, -INT8-format C8 and other payload geometries are not implemented by this bridge. - -Core's planner receives the complete key/scale page size and packed-row stride -alignment before choosing the shared-slab capacity. Strided views preserve the -original allocation and layer offset. The existing scatter handles packed bytes; -prefill Resident gathers reinterpret FP8 as bytes before using Core's gather. -HiSparse writes through its own current destination once; the ordinary SFA C8 -source-slot store is skipped. No new fused scatter or Attention kernel is added. - -## Extension points for contributors - -Use existing interfaces, rather than adding unused placeholder hooks: - -| Boundary | Current consumer and responsibility | Follow-up work | -| --- | --- | --- | -| `NPUPlatform.get_hisparse_backend_cls` | Core selects the device backend only for HiSparse | Backend capability/version support | -| `HiSparseBackend.create_cache_ops` | Runtime invokes resolver, refill and invalidation operations | Faster resolver/refill, vectorization and balancing | -| `validate_kv_cache_config` / `prepare_kv_cache_views` | Core validates BF16/C8 before allocation and binds views inside rollback scope | Additional payload formats with layout tests | -| Backend `pin_tensor` / `unpin_tensor` / `swap_blocks_batch` | Core owns completion, explicit registration release and transfer lifetime | DMA implementation and supported memory capabilities | -| `AscendHiSparseMetadataBuilder` / `AscendHiSparseImpl` | SFA factories select metadata and the BF16/C8 layout bridge | A smaller layout-conversion cost | -| `AscendHiSparseIndexerBackend.write_cache` | The existing cache backend interface selects HiSparse once; its override handles C8 page strides | Other Indexer payload layouts | - -The model's Indexer wrapper reuses the cache layer's `get_attn_backend()` choice, -preserving model-specific backend factories. HiSparse owns its metadata builder -and C8 write override. The ordinary Indexer write path remains unchanged, and -BF16 HiSparse writes reuse it after metadata remaps the slots. Cache binding -creates the key/scale views; it does not attach payload-format routing strings. - -`hisparse/cache_ops.py` owns operation validation and persistent per-group workspace; -`ops/triton/hisparse.py` implements the reference NPU path. Followers share their -leader's scratch only on that group's ordered stream. Independent groups and -current request/state mappings remain separate. - -Code ownership follows these boundaries: - -```text -vllm_ascend/ - attention/sfa_hisparse.py # Attention dispatch, C8 adaptation and native calls - hisparse/ - backend.py # Core backend interface, registration and DMA - binding.py # Shared-slab specs and key/scale views - cache_ops.py # Residency, refill, invalidation and byte copies - c8.py # C8 packing, views and physical-index helpers - prefill.py # Core prefill plan with scoped NPU index selection - ops/triton/hisparse.py # Cache-operation kernels; no Attention arithmetic -``` - -The common package contains no Attention invocation. The HiSparse Attention -implementation owns query preparation, current length tensors and the call to -the existing native SFA operation. The Triton and native Attention kernels are -unchanged by this organization. - -Core's generic `torch.Stream`, `torch.Event` and current-stream interfaces retain -their event placement and stream identities. On the supported PTA stack, generic -waits do not drain Host submission queues to propagate a graph-capture fork before -CANN allocates workspace. The backend wraps only the bound runtime copy stream: -after a wait during capture, obtaining the native handle drains that Host queue. -It adds no device synchronization or replay-time Python work and replaces no -global Torch binding. A regression covers production-size sort workspace across -multiple graph owners and changing replays. Remove this adapter when PTA's generic -binding propagates the captured dependency itself. - -The model-owned index-group builder is carried in Core's existing -`MLAModules.index_group_builder` field, preserving the common MLA wrapper ABI. +Use a dedicated output directory. Select architecture and usable per-AIV UB +budget for the actual build target; these are hardware/compiler parameters. +The kernel sizes its hash table from the request union and available UB. An +unsupported capacity or missing extension raises an error before execution; +there is no Triton resolver fallback. Other vllm-ascend features do not load this +optional extension. -HiSparse groups use Core's coordinator and shared-slab memory budget. Ascend's -existing private-buffer budget correction and hybrid coordinator are bypassed -for those groups. V2 retains Core's source/group ownership and skips metadata -construction for Resident/Hot groups that have no Attention consumer. +## Enable HiSparse -Follow-up contributions should be separate and measurable: other packed formats, -native resolver/refill improvements, Host-writeback optimization, metadata -reuse, additional hardware and P/D transport. PR19109 retains these experiments. -AVSG, lazy-MTP policy, native optimized kernels, private Core adapters and -machine-specific benchmarking tools are excluded here. - -Preserve completion-before-unregister, graph allocation owners, invalidation on -overwrite and protection of every selected token before choosing an LRU victim. -A same-shaped batch is not permission to reuse old lengths or request mappings. - -## Validation - -CPU boundary checks in the normal installed development environment: +Set `VLLM_USE_V2_MODEL_RUNNER=1` and add these options to the compatible model's +normal serving command: ```bash -python -m pytest -q \ - tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +--block-size 128 \ +--kv-cache-dtype bfloat16 \ +--attention-config '{"indexer_kv_dtype":"bf16","hisparse_config":{"eager_host_mirror":true}}' \ +--kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":4}}' ``` -These check current Host offsets, failed-unregister ownership, group scratch -isolation, strided indexer aliases, planner alignment, unsupported formats, -and capture-only mirror suppression. C8 checks also -cover opaque component bytes, nonzero slab offsets, stable duplicate/padding -ordering, Indexer key/scale writes and planner alignment before allocation. -The tests directly import the production modules and Core/Ascend cache specs, -using the repository's normal CPU-UT conftest. Install the matching Core and -Ascend test dependencies before running them. CPU checks do not qualify device -execution. - -The weight-loading wrapper keeps Ascend main's explicit activation-dtype ABI, -already covered by `tests/ut/ops/test_mla_weight_processing.py`; this foundation -does not add another calling convention or a duplicate test module. - -Reuse the existing suites for behavior this PR does not implement: - -| Behavior | Existing coverage | -| --- | --- | -| Core planning and HMA budgets | Core `tests/v1/core/test_kv_cache_planning.py` and `test_kv_cache_utils.py` | -| Core Host-mirror ordering, zero-token steps and request residency | Core `tests/v1/kv_connector/unit/test_hisparse_connector.py` | -| Core prefill plan and sparse layout rules | Core `tests/v1/attention/test_sparse_mla_backends.py` and `test_sparse_mla_kv_cache_layout.py` | -| Native SFA preprocessing and operator dispatch | Ascend `tests/ut/attention/test_sfa_v1.py` | - -These suites retain their original platform/dependency requirements. In -particular, Core's CUDA-kernel cases do not qualify Ascend kernels. The NPU -modules below keep coverage for Ascend resolver/refill/invalidation, strided -Resident/Host staging and the BF16/C8 data bridge. Preprocessing uses the existing -native SFA operations as its reference; it contains no embedded activation dump -or independent-versus-fused rounding requirement. - -Real-NPU cases are included for resolver/refill/overwrite invalidation, independent -FP32 Attention references, selection order/duplicates, native BF16/C8 preprocessing, -C8 key/scale writes, packed-byte scatter and Resident/Host prefill staging, in -eager and changing-graph cases: +For native A5 C8, replace the two dtype options with: ```bash -pytest --confcutdir=tests/e2e/pull_request/one_card \ - tests/e2e/pull_request/one_card/test_hisparse_ops.py \ - tests/e2e/pull_request/one_card/test_hisparse_attention.py \ - tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py \ - tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +--kv-cache-dtype fp8 \ +--attention-config '{"indexer_kv_dtype":"fp8","hisparse_config":{"eager_host_mirror":true}}' ``` -On Ascend 950 with Torch 2.10.0, torch-npu 2.10.0.post7 and CANN 9.2.1, -25 foundation CPU cases, the existing MLA-wrapper case and 32 Indexer-spec cases -passed: 58 in total. Running the four complete NPU modules passed all 44 cases, -including BF16/C8 Attention, pinned Host writeback, production-size cross-stream -sort workspace and changing graph replays. The run reused the compiled Ascend -extension and vendor package; it was not an independent native rebuild. - -The original `test_resident_prefill_plan_shared_layers[2-128]` failure was -traced to NPU advanced column indexing in Core's `ensure_gpu_sources` for a -singleton-column input. The Ascend plan subclasses Core's plan and runs its -unchanged source-resolution method inside a local `TorchFunctionMode` that -uses equivalent `index_select` for NPU column selection. Core still constructs -the plan and owns representative pages, source-key caching and Resident/Host -masks. This adapter adds no global Torch replacement, copied planner or CPU readback. -Remove this adapter when the supported device stack fixes advanced indexing. - -The original case remains enabled and passes. A small regression also checks -changing request-state rows. Prefill metadata planning executes outside the -model graph: this stack's Core `scatter_reduce` path cannot be captured. The -existing eager/changing-graph payload and Attention cases continue to pass; -the tests do not claim graph capture of the complete Prefill planner. - -Dataset accuracy, additional lifecycle coverage and same-workload performance -remain prerequisites for production qualification. Existing C8 benchmark results -are intentionally not presented as this implementation's performance. - -Full-model integration additionally preserves the Core Indexer cache role and -merge semantics, and retains each Indexer cache's backend choice after the model -construction context exits. With these fixes, the complete TP8/MTP3 model loaded, -initialized its KV caches, captured all eight ranks and reached serving. Its -first real request exposed a Host writeback failure; serving alone is not a -functional pass. - -The DMA reproduction distinguishes page locking from mapping: pageable and -PTA-pinned destinations passed, while a pageable allocation registered only for -mapping failed both batch and single-copy APIs on this stack. The private-pool -adapter reuses PR #19109's pinned-allocation approach before Core creates any -cache views. Core still owns completion, cache views and release. Other backends -keep their original allocator. The writeback regression passed on the NPU with -changing payloads, completion fences and untouched-neighbor checks. - -The next full-model attempt exposed invalid CANN sort workspace during captured -cross-stream execution. A production-size reproducer retained generic streams, -events and the sort operation: the original path failed, while capture-propagated -forks passed. With the narrow backend adapter, the complete GLM-5.2 W4A8C8 model -loaded with TP8, MTP3, V2 and native C8 KV; all eight ranks captured. Three natural -controls returned the expected answers. Two concurrent requests each consumed -exactly 16,384 input tokens and returned sensible summaries (24 and 23 output -tokens); all five responses were successful with actual token IDs. Host mirroring -remained enabled. The model exited cleanly, its workers were reaped and the lease -was released. This smoke test is not a dataset accuracy, P/D or performance result. +| Parameter | Meaning | +| --- | --- | +| `hisparse_config` | Enables Core HiSparse for the supported sparse Attention backend. | +| `eager_host_mirror` | Keep `true` for this foundation; Core maintains the Host KV mirror. | +| `host_pool_gib` | Host pool size per rank. Size it for the workload; the example reserves 4 GiB per rank. | +| `device_buffer_size` | Per-request Hot capacity from Core's HiSparse configuration. It must cover the selected verification-row union. | +| `kv_cache_dtype` / `indexer_kv_dtype` | Select main and Indexer KV formats independently from model-weight quantization. | + +## Capacity and correctness notes + +The resolver preserves the complete MTP query-row union before choosing LRU +victims, and produces the refill plan on device. Hot capacity must fit both the +int16 LRU representation (at most 32768 rows) and the compiled target's UB budget. +These checks do not change the configured TopK or discard selected tokens. + +Native C8 main rows contain 512 E4M3FN bytes, 64 BF16 RoPE values and four FP32 +scales (656 bytes total). Indexer rows contain 128 E4M3FN keys and one FP32 scale. +Views preserve the shared-slab page strides and layer offsets. The existing +quantizer and Attention operations are reused; payloads are copied as bytes. + +Keep completion-before-unregister and graph allocation ownership intact. +Overwriting KV still invalidates any old Hot copy. Equal batch shapes do not +permit reusing stale request mappings or lengths. No upstream vLLM source change +or unmerged Core PR is required. + +## Validation and performance + +The previous Triton-based foundation passed its CPU/NPU component suites and a +TP8/MTP3 full-model smoke. That evidence does not qualify this native-resolver +integration. The fused kernel is reused from the optimized implementation; +this foundation's integration and full-model results are recorded separately. + +Dataset accuracy and same-workload ON/OFF performance are not yet qualified for +this foundation. No throughput, TTFT or TPOT improvement is claimed here. diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index f007c28f72bf..673f799094a5 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -3,6 +3,7 @@ entries: - category: feature text: >- Added experimental Ascend HiSparse BF16 and native C8 sparse MLA/SFA - support using existing Attention operators and Core lifecycle interfaces. - Full-model qualification of the foundation remains pending. + support using a fused AscendC residency resolver, existing Attention + operators and Core lifecycle interfaces. The resolver requires its + compiled extension and does not fall back to Triton. hardware: [A5] diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index aafe5d81f2ea..e6123f872fa7 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -62,73 +62,6 @@ def test_private_host_pool_retains_pinned_storage_for_dma_writeback(): backend.unpin_tensor(registered) -def test_resolver_sort_workspace_survives_cross_stream_graph_owners(): - # Production-sized Hot/TopK buffers exercise CANN's workspace sort path; - # the small resolver oracle cases do not need this external workspace. - width, topk, max_rows = 10240, 2048, 128 - groups = [] - for _ in range(22): - stream = torch.Stream(device="npu") - wrapped = capture_aware_stream(stream) - assert wrapped.stream_id == stream.stream_id and wrapped.device == stream.device - assert capture_aware_stream(wrapped) is wrapped - cache = torch.full((max_rows, width), 2**63 - 1, dtype=torch.int64, device="npu") - rows = torch.full((max_rows, topk), 2**63 - 1, dtype=torch.int64, device="npu") - groups.append( - ( - wrapped, - torch.Event(), - cache, - torch.empty_like(cache), - torch.empty_like(cache), - rows, - torch.empty_like(rows), - torch.empty_like(rows), - ) - ) - - def forward(count): - current = torch.accelerator.current_stream() - for stream, ready, cache, values, indices, rows, row_values, row_indices in groups: - stream.wait_stream(current) - with stream: - torch.sort(rows[:count], dim=-1, out=(row_values[:count], row_indices[:count])) - torch.sort(cache[:count], dim=-1, out=(values[:count], indices[:count])) - ready.record(stream) - current.wait_event(ready) - - graphs = {} - pool = torch.npu.graph_pool_handle() - try: - for count in (128, 64, 32, 16, 8, 4, 1): - forward(count) - torch.npu.synchronize() - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph, pool=pool): - forward(count) - graphs[count] = graph - for step, count in enumerate((4, 1, 8, 4, 16, 4)): - hot_ids = torch.arange(width, dtype=torch.int64) - query_ids = torch.arange(topk, dtype=torch.int64) - hot = ((((hot_ids + step) % 127) << 32) | hot_ids).flip(0) - query = ((((query_ids + step) % 67) << 32) | query_ids).flip(0) - hot[::17] = query[::17] = 2**63 - 1 - expected_hot = torch.sort(hot).values.expand(count, -1) - expected_query = torch.sort(query).values.expand(count, -1) - for _, _, cache, _, _, rows, _, _ in groups: - cache.copy_(hot.npu()) - rows.copy_(query.npu()) - graphs[count].replay() - torch.npu.synchronize() - for _, _, _, values, _, _, row_values, _ in groups: - torch.testing.assert_close(values[:count].cpu(), expected_hot, rtol=0, atol=0) - torch.testing.assert_close(row_values[:count].cpu(), expected_query, rtol=0, atol=0) - finally: - torch.npu.synchronize() - for graph in graphs.values(): - graph.reset() - - @pytest.mark.parametrize("capture", [False, True]) def test_c8_bytes_reuse_existing_scatter_with_padding_and_page_gaps(capture): rows, pages, block, width = 5, 4, 128, 656 @@ -396,7 +329,8 @@ def forward(): @pytest.mark.parametrize("query_rows,hot_size", [(2, 4096), (4, 10240)]) -def test_topk_2048_speculative_union_graph(query_rows, hot_size): +@pytest.mark.parametrize("cross_stream", [False, True]) +def test_topk_2048_speculative_union_graph(query_rows, hot_size, cross_stream): topk, width, block = 2048, 576, 128 union = query_rows * topk blocks = hot_size // block + 1 @@ -418,7 +352,10 @@ def test_topk_2048_speculative_union_graph(query_rows, hot_size): sh, sd, sc = torch.empty_like(out), torch.empty_like(out), torch.empty_like(counts) stats = torch.zeros(2, dtype=torch.int64, device="npu") - def forward(): + stream = capture_aware_stream(torch.Stream(device="npu")) if cross_stream else None + ready = torch.Event() if cross_stream else None + + def resolve_and_gather(): ops.hisparse_resolve_residency( host, hot, @@ -449,6 +386,17 @@ def forward(): ) ops.hisparse_gather_compact(host, hot, sh, sd, sc) + def forward(): + if stream is None: + resolve_and_gather() + return + current = torch.accelerator.current_stream() + stream.wait_stream(current) + with stream: + resolve_and_gather() + ready.record(stream) + current.wait_event(ready) + forward() torch.npu.synchronize() expected = torch.arange(block, block + union, dtype=torch.int32).view(query_rows, topk) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 50a05b9e8f5d..7537b0fcf6f6 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -125,7 +125,15 @@ def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): @pytest.mark.parametrize("max_rows,hot_rows,union_rows", [(1, 16, 8), (9, 48, 24), (37, 128, 72)]) -def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups(max_rows, hot_rows, union_rows): +def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( + monkeypatch, max_rows, hot_rows, union_rows +): + from vllm_ascend.hisparse import resolver + + # Loading a device extension is the only mocked boundary; the group-owned + # tensors below are real CPU storage, not a replacement resolver algorithm. + native = NS(scratch_capacity=lambda: 175616) + monkeypatch.setattr(resolver, "_load_resolver", lambda: native) backend = AscendHiSparseBackend() def group(): @@ -149,13 +157,40 @@ def group(): assert follower.host_mapping is first.registered_host_pool._hisparse_registration assert second.cache_ops.workspace is not first.cache_ops.workspace first_workspace, second_workspace = first.cache_ops.workspace, second.cache_ops.workspace - assert first_workspace.row_keys.shape == (max_rows * hot_rows,) - assert first_workspace.cache_keys.shape == (max_rows, hot_rows) + assert first_workspace.first_refs.shape == (max_rows * hot_rows,) + assert first_workspace.errors.shape == (max_rows, 8) assert first_workspace.max_union_rows == union_rows - first_workspace.row_keys.fill_(13) - second_workspace.row_keys.zero_() - assert torch.all(follower.workspace.row_keys == 13) - assert torch.count_nonzero(second_workspace.row_keys) == 0 + first_workspace.first_refs.fill_(13) + second_workspace.first_refs.zero_() + assert torch.all(follower.workspace.first_refs == 13) + assert torch.count_nonzero(second_workspace.first_refs) == 0 + + +@pytest.mark.parametrize("hot,union", [(16, 8), (8192, 8192), (12288, 64), (32768, 32768)]) +def test_native_resolver_scratch_follows_request_capacity(hot, union): + from vllm_ascend.hisparse.resolver import scratch_layout + + buckets, size = scratch_layout(hot, union) + assert buckets > union and buckets & (buckets - 1) == 0 + words = 2 * buckets + hot + 2 * ((hot + 31) // 32) + 64 + assert size == (words + 15) // 16 * 64 + 512 + + +def test_native_resolver_uses_available_ub_without_truncating_the_union(): + from vllm_ascend.hisparse.resolver import scratch_layout + + buckets, size = scratch_layout(12288, 8192, 175616) + assert buckets == 15328 and buckets > 8192 and size == 175616 + with pytest.raises(RuntimeError, match="needs at least 118592 UB bytes"): + scratch_layout(12288, 8192, 118591) + + +def test_missing_native_resolver_raises_instead_of_falling_back(monkeypatch): + from vllm_ascend.hisparse import resolver + + monkeypatch.setattr(resolver.Path, "is_file", lambda path: False) + with pytest.raises(RuntimeError, match="requires its compiled AscendC resolver"): + resolver._load_resolver() @pytest.fixture diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 4a35af878677..c1d00eabbcad 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -53,12 +53,14 @@ class AscendHiSparseBackend(HiSparseBackend): stats_dtype = torch.int64 def check_available(self) -> None: + from vllm_ascend.hisparse.resolver import _load_resolver from vllm_ascend.patch.worker.patch_hisparse_host_pool import install from vllm_ascend.utils import bootstrap_custom_op_env install() # Import the device kernels before the connector allocates Host storage. importlib.import_module("vllm_ascend.hisparse.cache_ops") + _load_resolver() if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): # A5 disables the global custom-op selection policy. HiSparse still # needs its explicit DMA dependency; load the extension lazily here diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 085dfd3f7087..5e169d05bcf1 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -16,18 +16,10 @@ from vllm.triton_utils import triton from vllm_ascend.ops.triton.hisparse import ( - assign_residency_kernel, - classify_residency_kernel, - finish_residency_kernel, gather_rows_kernel, invalidate_kernel, pack_attention_rows_kernel, - partition_residency_kernel, - prepare_residency_kernel, - protect_residency_kernel, - residency_stats_kernel, scatter_packed_kv_kernel, - update_lru_kernel, ) _ACL_HOST_REGISTER_MAPPED = 0 @@ -127,21 +119,11 @@ def __init__(self, max_rows: int, hot_size: int, max_union_rows: int, device: to self.max_rows = max_rows self.hot_size = hot_size self.max_union_rows = max_union_rows - # Packed int64 keys retain exact int32 host IDs and deterministic first - # reference positions. torch.sort's out= ABI reuses these NPU buffers. - self.row_keys = torch.empty(max_rows * hot_size, dtype=torch.int64, device=device) - self.sorted_rows = torch.empty_like(self.row_keys) - self.row_sort_indices = torch.empty_like(self.row_keys) - self.cache_keys = torch.empty((max_rows, hot_size), dtype=torch.int64, device=device) - self.sorted_cache = torch.empty_like(self.cache_keys) - self.cache_sort_indices = torch.empty_like(self.cache_keys) + from vllm_ascend.hisparse.resolver import _load_resolver, scratch_layout + + self.native = _load_resolver() + scratch_layout(hot_size, max_union_rows, self.native.scratch_capacity()) self.first_refs = torch.empty(max_rows * hot_size, dtype=torch.int32, device=device) - self.first_misses = torch.empty_like(self.first_refs) - self.metadata = torch.empty((max_rows, 8), dtype=torch.int32, device=device) - self.row_stats = torch.empty((max_rows, 4), dtype=torch.int64, device=device) - self.lru_temp = torch.empty((max_rows, hot_size), dtype=torch.int32, device=device) - self.protected = torch.empty_like(self.lru_temp) - self.hit_prefix = torch.empty_like(self.lru_temp) self.errors = torch.zeros((max_rows, 8), dtype=torch.int32, device=device) @@ -260,9 +242,9 @@ def hisparse_resolve_residency( raise ValueError("HiSparse resolver input exceeds its prepared workspace") if region_stride != device_global_indices.size(1) or region_stride != hot_size or hot_size < topk: raise ValueError("HiSparse region and LRU widths must match and cover top-k") - if topk <= 0 or topk % 8 or hot_size % 16: - raise ValueError("Ascend HiSparse requires top-k aligned to 8 and hot capacity aligned to 16") - if workspace.row_keys.device != device: + if topk <= 0: + raise ValueError("HiSparse top-k must be positive") + if workspace.first_refs.device != device: raise ValueError("HiSparse resolver workspace must be on the cache device") if host_cache.size(0) >= 2**31: raise ValueError("HiSparse host row IDs must fit int32") @@ -312,164 +294,14 @@ def hisparse_resolve_residency( _check_tensor(swap_counts, "swap_counts", device, (torch.int32,), shape=(rows,)) if stats is not None: _check_tensor(stats, "stats", device, (torch.int64,), shape=(2,)) - dummy = hot_indices - source = source_block_table - resident = resident_block_table - has_swaps = swap_counts is not None if ( len({tensor is not None for tensor in (swap_counts, swap_host_physical_rows, swap_device_physical_rows)}) != 1 ): raise ValueError("All three compact HiSparse swap outputs are required together") - row_keys = workspace.row_keys[: rows * topk].view(rows, topk) - sorted_rows = workspace.sorted_rows[: rows * topk].view(rows, topk) - row_sort_indices = workspace.row_sort_indices[: rows * topk].view(rows, topk) - cache_keys = workspace.cache_keys[:rows] - sorted_cache = workspace.sorted_cache[:rows] - prepare_residency_kernel[(rows,)]( - global_indices, - request_ids if request_ids is not None else dummy, - request_state_indices if request_state_indices is not None else dummy, - source if source is not None else dummy, - resident if resident is not None else dummy, - device_global_indices, - hot_indices, - resolved_global_indices if resolved_global_indices is not None else dummy, - valid_counts if valid_counts is not None else dummy, - row_keys, - cache_keys, - workspace.metadata, - num_valid_rows if num_valid_rows is not None else dummy, - ROWS=rows, - TOP_K=topk, - HOT_SIZE=hot_size, - INPUT_STRIDE=global_indices.stride(0), - COUNT_STRIDE=valid_counts.stride(0) if valid_counts is not None else 1, - STATE_ROWS=device_global_indices.size(0), - STATE_COUNT=request_state_indices.numel() if request_state_indices is not None else 0, - HOST_ROWS=host_cache.size(0), - SOURCE_STRIDE=source.stride(0) if source is not None else 0, - SOURCE_ROWS=source.size(0) if source is not None else 0, - SOURCE_COLS=source.size(1) if source is not None else 0, - SOURCE_BLOCK=source_block_size if source is not None else 1, - RESIDENT_STRIDE=resident.stride(0) if resident is not None else 0, - RESIDENT_ROWS=resident.size(0) if resident is not None else 0, - RESIDENT_COLS=resident.size(1) if resident is not None else 0, - RESIDENT_BLOCK=resident_block_size if resident is not None else 1, - RESIDENT_NULL=resident_null_block, - HAS_REQUESTS=request_ids is not None, - HAS_STATE=request_state_indices is not None, - HAS_SOURCE=source is not None, - HAS_RESIDENT=resident is not None, - HAS_RESOLVED=resolved_global_indices is not None, - HAS_COUNTS=valid_counts is not None, - HAS_VALID_ROWS=num_valid_rows is not None, - BK=triton.next_power_of_2(topk), - multibuffer=False, - ) - torch.sort(row_keys, dim=-1, out=(sorted_rows, row_sort_indices)) - torch.sort(cache_keys, dim=-1, out=(sorted_cache, workspace.cache_sort_indices[:rows])) - classify_residency_kernel[(rows,)]( - row_keys, - sorted_rows, - sorted_cache, - workspace.metadata, - hot_block_table, - hot_indices, - workspace.first_refs, - workspace.first_misses, - workspace.row_stats, - TOP_K=topk, - HOT_SIZE=hot_size, - HOT_BLOCK=hot_cache.size(1), - TABLE_STRIDE=hot_block_table.stride(0), - BK=triton.next_power_of_2(topk), - multibuffer=False, - ) - protect_residency_kernel[(rows, triton.cdiv(hot_size, 256))]( - sorted_rows, - workspace.metadata, - device_global_indices, - lru_slots, - workspace.protected, - TOP_K=topk, - HOT_SIZE=hot_size, - STATE_ROWS=device_global_indices.size(0), - BLOCK=256, - multibuffer=False, - ) - torch.cumsum(workspace.protected[:rows], dim=-1, dtype=torch.int32, out=workspace.hit_prefix[:rows]) - partition_residency_kernel[(rows, triton.cdiv(hot_size, 256))]( - workspace.metadata, - workspace.protected, - workspace.hit_prefix, - workspace.row_stats, - lru_slots, - workspace.lru_temp, - workspace.errors, - HOT_SIZE=hot_size, - MAX_UNION=max_union_rows, - STATE_ROWS=device_global_indices.size(0), - BLOCK=256, - multibuffer=False, - ) - update_lru_kernel[(rows, triton.cdiv(hot_size, 256))]( - workspace.metadata, - workspace.lru_temp, - lru_slots, - workspace.errors, - HOT_SIZE=hot_size, - BLOCK=256, - multibuffer=False, - ) - assign_residency_kernel[(rows,)]( - row_keys, - workspace.metadata, - workspace.first_misses, - hot_block_table, - hot_indices, - device_global_indices, - workspace.lru_temp, - workspace.errors, - swap_host_physical_rows if has_swaps else dummy, - swap_device_physical_rows if has_swaps else dummy, - swap_counts if has_swaps else dummy, - TOP_K=topk, - HOT_SIZE=hot_size, - HOT_BLOCK=hot_cache.size(1), - TABLE_STRIDE=hot_block_table.stride(0), - HAS_SWAPS=has_swaps, - BK=triton.next_power_of_2(topk), - multibuffer=False, - ) - finish_residency_kernel[(rows,)]( - workspace.metadata, - workspace.first_refs, - workspace.first_misses, - hot_indices, - attention_indices if attention_indices is not None else dummy, - miss_mask if miss_mask is not None else dummy, - valid_counts if valid_counts is not None else dummy, - workspace.errors, - TOP_K=topk, - HOT_BLOCK=hot_cache.size(1), - ATTENTION_BLOCK=attention_block_stride, - ATTENTION_STRIDE=attention_indices.stride(0) if attention_indices is not None else topk, - COUNT_STRIDE=valid_counts.stride(0) if valid_counts is not None else 1, - HAS_ATTENTION=attention_indices is not None, - HAS_MISS_MASK=miss_mask is not None, - HAS_COUNTS=valid_counts is not None, - BK=triton.next_power_of_2(topk), - multibuffer=False, - ) - if stats is not None: - residency_stats_kernel[(1,)]( - workspace.row_stats, - stats, - ROWS=rows, - BLOCK=triton.next_power_of_2(rows), - multibuffer=False, - ) + from vllm_ascend.hisparse.resolver import resolve + + resolve(self, locals()) def hisparse_gather_compact( self, host_cache, hot_cache, miss_global_indices, miss_hot_indices, miss_counts diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py new file mode 100644 index 000000000000..bee43b434f2b --- /dev/null +++ b/vllm_ascend/hisparse/resolver.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tensor-only native resolver dispatch; no address-keyed launch ownership.""" + +from pathlib import Path + +import torch + + +def scratch_layout(hot, max_union_rows, capacity=None): + """Size hash/LRU scratch from the request capacity, not a model shape.""" + if not 0 < max_union_rows <= hot <= 32768: + raise ValueError("HiSparse union must fit the hot cache and int16 LRU slot representation") + hash_size = 1 << max_union_rows.bit_length() + tiles = (hot + 31) // 32 + fixed_words = hot + 2 * tiles + 64 + if capacity is not None: + available_words = max(0, (capacity - 512) // 64) * 16 + max_hash_size = (available_words - fixed_words) // 2 + if max_hash_size <= max_union_rows: + minimum_bytes = ((2 * (max_union_rows + 1) + fixed_words + 15) // 16) * 64 + 512 + raise RuntimeError( + f"HiSparse hot_size={hot}, max_union_rows={max_union_rows} needs at least {minimum_bytes} UB bytes; " + f"the compiled target allows {capacity}. Lower device_buffer_size or speculative tokens." + ) + hash_size = min(hash_size, max_hash_size) + scratch_ints = ((2 * hash_size + hot + 2 * tiles + 64 + 15) // 16) * 16 + return hash_size, scratch_ints * 4 + 512 + + +def _load_resolver(): + library = Path(__file__).with_name("torch_binding.so") + if not library.is_file(): + raise RuntimeError("HiSparse requires its compiled AscendC resolver. Build csrc/hisparse_runtime first.") + if str(library) not in torch.ops.loaded_libraries: + torch.ops.load_library(str(library)) + namespace = torch.ops.hisparse_fused60103 + if not all(hasattr(namespace, name) for name in ("scratch_capacity", "resolve_tensors", "supports_budgeted_hash")): + raise RuntimeError("Rebuild the HiSparse native extension for the budget-aware Tensor resolver ABI") + if not namespace.supports_budgeted_hash(): + raise RuntimeError("HiSparse requires a native resolver with budget-aware hash support") + return namespace + + +def resolve_tensors(native, arguments, resolved_scratch, errors): + """Pass live tensors; C++ derives geometry and protects queued storage. + + During replay, external tensors remain owned by the runtime/workspace. + Callers must destroy graphs before releasing or rebinding those owners. + """ + a = arguments + native.resolve_tensors( + a["global_indices"], + a["request_ids"], + a["request_state_indices"], + a["source_block_table"], + a["resident_block_table"], + a["hot_block_table"], + a["device_global_indices"], + a["lru_slots"], + a["hot_indices"], + a["attention_indices"], + resolved_scratch, + a["valid_counts"], + a["swap_host_physical_rows"], + a["swap_device_physical_rows"], + a["swap_counts"], + a["miss_mask"], + a["stats"], + errors, + a["num_valid_rows"], + a["host_cache"].shape[0], + a["hot_cache"].shape[1], + a["source_block_size"], + a["resident_block_size"], + a["resident_null_block"], + a["attention_block_stride"], + a["max_union_rows"], + ) + + +def resolve(owner, scope): + """Resolve with the native request-union kernel; no Triton fallback.""" + rows, topk = scope["global_indices"].shape + workspace = owner.workspace + resolved = scope["resolved_global_indices"] + if resolved is None: + resolved = workspace.first_refs[: rows * topk].view(rows, topk) + resolve_tensors(workspace.native, scope, resolved, workspace.errors[:rows]) diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 9e383e003289..09d7f219e5cc 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -1,419 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Device-only HiSparse planning and mapped-host row copies for Ascend. +"""Mapped-Host refill, invalidation and packed KV layout operations. -Device sorts and binary searches form the union of a request's contiguous -verification rows. All selected slots are protected before any eviction. -Scratch belongs to the caller and is allocated before graph capture. +Residency planning is provided by the native AscendC resolver. """ from vllm.triton_utils import tl, triton -@triton.jit -def _translate( - token, - request, - state, - source, - resident, - HOST_ROWS: tl.constexpr, - SOURCE_STRIDE: tl.constexpr, - SOURCE_ROWS: tl.constexpr, - SOURCE_COLS: tl.constexpr, - SOURCE_BLOCK: tl.constexpr, - RESIDENT_STRIDE: tl.constexpr, - RESIDENT_ROWS: tl.constexpr, - RESIDENT_COLS: tl.constexpr, - RESIDENT_BLOCK: tl.constexpr, - RESIDENT_NULL: tl.constexpr, - HAS_SOURCE: tl.constexpr, - HAS_RESIDENT: tl.constexpr, -): - host = token - device = tl.full(token.shape, -1, tl.int32) - if HAS_SOURCE: - col = token // SOURCE_BLOCK - valid = (token >= 0) & (request >= 0) & (request < SOURCE_ROWS) & (col < SOURCE_COLS) - block = tl.load(source + request.to(tl.int64) * SOURCE_STRIDE + col, mask=valid, other=0) - host = tl.where(valid & (block > 0), block * SOURCE_BLOCK + token % SOURCE_BLOCK, -1) - if HAS_RESIDENT: - col = token // RESIDENT_BLOCK - valid = (token >= 0) & (state >= 0) & (state < RESIDENT_ROWS) & (col < RESIDENT_COLS) - block = tl.load(resident + state.to(tl.int64) * RESIDENT_STRIDE + col, mask=valid, other=RESIDENT_NULL) - device = tl.where( - valid & (block >= 0) & (block != RESIDENT_NULL), - block * RESIDENT_BLOCK + token % RESIDENT_BLOCK, - -1, - ) - host = tl.where((host >= 0) & (host < HOST_ROWS), host, -1) - return host, device - - -@triton.jit -def _lookup(sorted_keys, value, active, N: tl.constexpr): - # Keys pack (nonnegative host row, original position). Lower-bound also - # identifies the first reference when a query contains duplicate indices. - low = tl.full(value.shape, 0, tl.int32) - high = tl.full(value.shape, N, tl.int32) - target = value.to(tl.int64) << 32 - for _ in range(0, tl.constexpr(N.bit_length())): - middle = (low + high) // 2 - key = tl.load(sorted_keys + middle, mask=active & (middle < N), other=9223372036854775807) - left = key < target - low = tl.where(left, middle + 1, low) - high = tl.where(left, high, middle) - key = tl.load(sorted_keys + low, mask=active & (low < N), other=9223372036854775807) - found = active & (low < N) & ((key >> 32) == value) - return (key & 4294967295).to(tl.int32), found - - -@triton.jit -def prepare_residency_kernel( - indices, - requests, - states, - source, - resident, - dgi, - hot_indices, - resolved, - counts, - row_keys, - cache_keys, - metadata, - num_valid_rows, - ROWS: tl.constexpr, - TOP_K: tl.constexpr, - HOT_SIZE: tl.constexpr, - INPUT_STRIDE: tl.constexpr, - COUNT_STRIDE: tl.constexpr, - STATE_ROWS: tl.constexpr, - STATE_COUNT: tl.constexpr, - HOST_ROWS: tl.constexpr, - SOURCE_STRIDE: tl.constexpr, - SOURCE_ROWS: tl.constexpr, - SOURCE_COLS: tl.constexpr, - SOURCE_BLOCK: tl.constexpr, - RESIDENT_STRIDE: tl.constexpr, - RESIDENT_ROWS: tl.constexpr, - RESIDENT_COLS: tl.constexpr, - RESIDENT_BLOCK: tl.constexpr, - RESIDENT_NULL: tl.constexpr, - HAS_REQUESTS: tl.constexpr, - HAS_STATE: tl.constexpr, - HAS_SOURCE: tl.constexpr, - HAS_RESIDENT: tl.constexpr, - HAS_RESOLVED: tl.constexpr, - HAS_COUNTS: tl.constexpr, - HAS_VALID_ROWS: tl.constexpr, - BK: tl.constexpr, -): - row = tl.program_id(0) - active_rows = tl.minimum(tl.maximum(tl.load(num_valid_rows), 0), ROWS) if HAS_VALID_ROWS else ROWS - request = tl.load(requests + row) if HAS_REQUESTS else row - request = tl.where(row < active_rows, request, -1) - first = row - end = row + 1 - if HAS_REQUESTS: - searching = first > 0 - while searching: - same = tl.load(requests + first - 1, mask=first > 0, other=-2147483648) == request - same = same & (first > 0) - first -= same.to(tl.int32) - searching = same - searching = end < active_rows - while searching: - same = tl.load(requests + end, mask=end < active_rows, other=-2147483648) == request - same = same & (end < active_rows) - end += same.to(tl.int32) - searching = same - state = tl.load(states + request, mask=(request >= 0) & (request < STATE_COUNT), other=-1) if HAS_STATE else request - valid_state = (state >= 0) & (state < STATE_ROWS) - m = tl.arange(0, 8) - meta = tl.where(m == 0, request, tl.where(m == 1, state, tl.where(m == 2, first, end))) - tl.store(metadata + row * 8 + m, meta) - i = tl.arange(0, BK) - token = tl.load(indices + row * INPUT_STRIDE + i, mask=(i < TOP_K) & valid_state, other=-1) - host, device = _translate( - token, - request, - state, - source, - resident, - HOST_ROWS, - SOURCE_STRIDE, - SOURCE_ROWS, - SOURCE_COLS, - SOURCE_BLOCK, - RESIDENT_STRIDE, - RESIDENT_ROWS, - RESIDENT_COLS, - RESIDENT_BLOCK, - RESIDENT_NULL, - HAS_SOURCE, - HAS_RESIDENT, - ) - tl.store(hot_indices + row * TOP_K + i, device, mask=i < TOP_K) - if HAS_RESOLVED: - tl.store(resolved + row * TOP_K + i, host, mask=i < TOP_K) - if HAS_COUNTS: - tl.store(counts + row * COUNT_STRIDE, tl.sum(((i < TOP_K) & ((device >= 0) | (host >= 0))).to(tl.int32), 0)) - cold = (i < TOP_K) & (device < 0) & (host >= 0) - key = tl.where(cold, (host.to(tl.int64) << 32) | i.to(tl.int64), 9223372036854775807) - tl.store(row_keys + row * TOP_K + i, key, mask=i < TOP_K) - for start in range(0, HOT_SIZE, 256): - p = start + tl.arange(0, 256) - cached = tl.load( - dgi + state.to(tl.int64) * HOT_SIZE + p, mask=valid_state & (row == first) & (p < HOT_SIZE), other=-1 - ) - cache_key = tl.where(cached >= 0, (cached.to(tl.int64) << 32) | p.to(tl.int64), 9223372036854775807) - tl.store(cache_keys + row * HOT_SIZE + p, cache_key, mask=p < HOT_SIZE) - - -@triton.jit -def classify_residency_kernel( - row_keys, - sorted_rows, - sorted_cache, - metadata, - hot_table, - hot_indices, - first_refs, - first_misses, - row_stats, - TOP_K: tl.constexpr, - HOT_SIZE: tl.constexpr, - HOT_BLOCK: tl.constexpr, - TABLE_STRIDE: tl.constexpr, - BK: tl.constexpr, -): - row = tl.program_id(0) - request = tl.load(metadata + row * 8) - first = tl.load(metadata + row * 8 + 2) - i = tl.arange(0, BK) - key = tl.load(row_keys + row * TOP_K + i, mask=i < TOP_K, other=9223372036854775807) - host = (key >> 32).to(tl.int32) - active = (i < TOP_K) & (host != 2147483647) - slot, hit = _lookup(sorted_cache + first * HOT_SIZE, host, active, HOT_SIZE) - reference = tl.full((BK,), 2147483647, tl.int32) - for query in range(first, row + 1): - pos, found = _lookup(sorted_rows + query * TOP_K, host, active, TOP_K) - reference = tl.minimum(reference, tl.where(found, query * TOP_K + pos, 2147483647)) - miss = active & ~hit & (reference == row * TOP_K + i) - tl.store(first_refs + row * TOP_K + i, tl.where(active, reference, -1), mask=i < TOP_K) - tl.store(first_misses + row * TOP_K + i, miss, mask=i < TOP_K) - block = tl.load(hot_table + request.to(tl.int64) * TABLE_STRIDE + slot // HOT_BLOCK, mask=hit, other=0) - tl.store(hot_indices + row * TOP_K + i, block * HOT_BLOCK + slot % HOT_BLOCK, mask=hit) - # Independent padded records avoid adjacent programs sharing a transfer line. - s = tl.arange(0, 4) - misses = tl.sum(miss.to(tl.int32), 0) - references = tl.sum(active.to(tl.int32), 0) - tl.store(row_stats + row * 4 + s, tl.where(s == 0, references - misses, tl.where(s == 1, misses, 0))) - - -@triton.jit -def protect_residency_kernel( - sorted_rows, - metadata, - dgi, - lru, - protected, - TOP_K: tl.constexpr, - HOT_SIZE: tl.constexpr, - STATE_ROWS: tl.constexpr, - BLOCK: tl.constexpr, -): - first = tl.program_id(0) - tile = tl.program_id(1) - leader = tl.load(metadata + first * 8 + 2) - p = tile * BLOCK + tl.arange(0, BLOCK) - hit = tl.full((BLOCK,), False, tl.int1) - if first == leader: - state = tl.load(metadata + first * 8 + 1) - end = tl.load(metadata + first * 8 + 3) - valid = (state >= 0) & (state < STATE_ROWS) & (p < HOT_SIZE) - slot = tl.load(lru + state.to(tl.int64) * HOT_SIZE + p, mask=valid, other=-1).to(tl.int32) - valid = valid & (slot >= 0) & (slot < HOT_SIZE) - cached = tl.load(dgi + state.to(tl.int64) * HOT_SIZE + slot, mask=valid, other=-1) - for row in range(first, end): - _, found = _lookup(sorted_rows + row * TOP_K, cached, valid & (cached >= 0), TOP_K) - hit = hit | found - tl.store(protected + first * HOT_SIZE + p, hit, mask=p < HOT_SIZE) - - -@triton.jit -def partition_residency_kernel( - metadata, - protected, - hit_prefix, - row_stats, - lru, - temp_lru, - errors, - HOT_SIZE: tl.constexpr, - MAX_UNION: tl.constexpr, - STATE_ROWS: tl.constexpr, - BLOCK: tl.constexpr, -): - first = tl.program_id(0) - tile = tl.program_id(1) - leader = tl.load(metadata + first * 8 + 2) - if first == leader: - state = tl.load(metadata + first * 8 + 1) - end = tl.load(metadata + first * 8 + 3) - p = tile * BLOCK + tl.arange(0, BLOCK) - valid = (state >= 0) & (state < STATE_ROWS) & (p < HOT_SIZE) - hit = tl.load(protected + first * HOT_SIZE + p, mask=valid, other=0) != 0 - hits = tl.load(hit_prefix + first * HOT_SIZE + HOT_SIZE - 1) - misses = 0 - for row in range(first, end): - misses += tl.load(row_stats + row * 4 + 1).to(tl.int32) - union = hits + misses - failed = union > MAX_UNION - if tile == 0: - err = tl.arange(0, 8) - tl.store(errors + first * 8 + err, tl.where(failed, union, 0)) - tl.store(metadata + first * 8 + 4, hits) - tl.store(metadata + first * 8 + 5, misses) - tl.device_assert(~failed, "HiSparse request union exceeds hot capacity") - if (union > 0) & ~failed: - slot = tl.load(lru + state.to(tl.int64) * HOT_SIZE + p, mask=valid, other=-1).to(tl.int32) - prefix = tl.load(hit_prefix + first * HOT_SIZE + p, mask=valid, other=0) - hit_rank = prefix - 1 - evictable = valid & ~hit - evict_rank = p - prefix - temp = temp_lru + first * HOT_SIZE - tl.store(temp + hit_rank, slot, mask=hit) - tl.store(temp + HOT_SIZE - 1 - evict_rank, slot, mask=evictable) - - -@triton.jit -def update_lru_kernel(metadata, temp_lru, lru, errors, HOT_SIZE: tl.constexpr, BLOCK: tl.constexpr): - first = tl.program_id(0) - tile = tl.program_id(1) - leader = tl.load(metadata + first * 8 + 2) - if first == leader: - hits = tl.load(metadata + first * 8 + 4) - misses = tl.load(metadata + first * 8 + 5) - failed = tl.load(errors + first * 8) != 0 - if ((hits + misses) > 0) & ~failed: - state = tl.load(metadata + first * 8 + 1) - p = tile * BLOCK + tl.arange(0, BLOCK) - remaining = HOT_SIZE - hits - misses - source = tl.where( - p < remaining, - HOT_SIZE - 1 - misses - p, - tl.where(p < remaining + misses, HOT_SIZE - 1 - (p - remaining), p - remaining - misses), - ) - updated = tl.load(temp_lru + first * HOT_SIZE + source, mask=p < HOT_SIZE, other=0) - tl.store(lru + state.to(tl.int64) * HOT_SIZE + p, updated, mask=p < HOT_SIZE) - - -@triton.jit -def assign_residency_kernel( - row_keys, - metadata, - first_misses, - hot_table, - hot_indices, - dgi, - temp_lru, - errors, - swap_host, - swap_device, - swap_counts, - TOP_K: tl.constexpr, - HOT_SIZE: tl.constexpr, - HOT_BLOCK: tl.constexpr, - TABLE_STRIDE: tl.constexpr, - HAS_SWAPS: tl.constexpr, - BK: tl.constexpr, -): - first = tl.program_id(0) - leader = tl.load(metadata + first * 8 + 2) - if first == leader: - request = tl.load(metadata + first * 8) - state = tl.load(metadata + first * 8 + 1) - end = tl.load(metadata + first * 8 + 3) - misses = tl.load(metadata + first * 8 + 5) - failed = tl.load(errors + first * 8) != 0 - if HAS_SWAPS: - for row in range(first, end): - count = tl.minimum(TOP_K, tl.maximum(0, misses - (row - first) * TOP_K)) - tl.store(swap_counts + row, tl.where(failed, 0, count)) - if (misses > 0) & ~failed: - i = tl.arange(0, BK) - consumed = 0 - temp = temp_lru + first * HOT_SIZE - for row in range(first, end): - miss = tl.load(first_misses + row * TOP_K + i, mask=i < TOP_K, other=0) != 0 - rank = tl.cumsum(miss.to(tl.int32), 0) - 1 + consumed - evict = tl.load(temp + HOT_SIZE - 1 - rank, mask=miss, other=-1) - good = miss & (evict >= 0) & (evict < HOT_SIZE) - key = tl.load(row_keys + row * TOP_K + i, mask=miss, other=0) - host = (key >> 32).to(tl.int32) - tl.store(dgi + state.to(tl.int64) * HOT_SIZE + evict, host, mask=good) - block = tl.load( - hot_table + request.to(tl.int64) * TABLE_STRIDE + evict // HOT_BLOCK, mask=good, other=0 - ) - physical = tl.where(good, block * HOT_BLOCK + evict % HOT_BLOCK, -1) - tl.store(hot_indices + row * TOP_K + i, physical, mask=miss) - if HAS_SWAPS: - tl.store(swap_host + first * TOP_K + rank, host, mask=miss) - tl.store(swap_device + first * TOP_K + rank, physical, mask=miss) - consumed += tl.sum(miss.to(tl.int32), 0) - - -@triton.jit -def finish_residency_kernel( - metadata, - first_refs, - first_misses, - hot_indices, - attention, - miss_mask, - counts, - errors, - TOP_K: tl.constexpr, - HOT_BLOCK: tl.constexpr, - ATTENTION_BLOCK: tl.constexpr, - ATTENTION_STRIDE: tl.constexpr, - COUNT_STRIDE: tl.constexpr, - HAS_ATTENTION: tl.constexpr, - HAS_MISS_MASK: tl.constexpr, - HAS_COUNTS: tl.constexpr, - BK: tl.constexpr, -): - row = tl.program_id(0) - i = tl.arange(0, BK) - first = tl.load(metadata + row * 8 + 2) - failed = tl.load(errors + first * 8) != 0 - reference = tl.load(first_refs + row * TOP_K + i, mask=i < TOP_K, other=-1) - own = tl.load(hot_indices + row * TOP_K + i, mask=i < TOP_K, other=-1) - physical = tl.load(hot_indices + reference, mask=(i < TOP_K) & (reference >= 0), other=-1) - physical = tl.where(failed, -1, tl.where(reference >= 0, physical, own)) - tl.store(hot_indices + row * TOP_K + i, physical, mask=i < TOP_K) - if HAS_ATTENTION: - encoded = tl.where(physical >= 0, physical // HOT_BLOCK * ATTENTION_BLOCK + physical % HOT_BLOCK, -1) - tl.store(attention + row * ATTENTION_STRIDE + i, encoded, mask=i < TOP_K) - if HAS_MISS_MASK: - miss = tl.load(first_misses + row * TOP_K + i, mask=i < TOP_K, other=0) - tl.store(miss_mask + row * TOP_K + i, tl.where(failed, 0, miss), mask=i < TOP_K) - if HAS_COUNTS: - if failed: - tl.store(counts + row * COUNT_STRIDE, 0) - - -@triton.jit -def residency_stats_kernel(row_stats, stats, ROWS: tl.constexpr, BLOCK: tl.constexpr): - i = tl.arange(0, BLOCK) - hit = tl.load(row_stats + i * 4, mask=i < ROWS, other=0) - miss = tl.load(row_stats + i * 4 + 1, mask=i < ROWS, other=0) - tl.store(stats, tl.load(stats) + tl.sum(hit, 0)) - tl.store(stats + 1, tl.load(stats + 1) + tl.sum(miss, 0)) - - @triton.jit def gather_rows_kernel( host_address, From a1dc5e68db83ef2bf1af1e23333d44f30d40bfbe Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 17:21:16 +0800 Subject: [PATCH 021/120] refactor(hisparse): integrate resolver with AscendC operator build Signed-off-by: zhenwenqi2024 --- CMakeLists.txt | 22 +---- csrc/hisparse/CMakeLists.txt | 14 +++ .../hisparse_resolver.cpp} | 8 +- .../hisparse_resolver_torch_adpt.h} | 27 ++---- csrc/hisparse_runtime/build.py | 95 ------------------- csrc/torch_binding.cpp | 15 +++ .../user_guide/feature_guide/hisparse.md | 28 +++--- .../test_hisparse_foundation_cpu.py | 4 +- vllm_ascend/hisparse/backend.py | 2 +- vllm_ascend/hisparse/cache_ops.py | 2 +- vllm_ascend/hisparse/resolver.py | 19 ++-- 11 files changed, 66 insertions(+), 170 deletions(-) create mode 100644 csrc/hisparse/CMakeLists.txt rename csrc/{hisparse_runtime/resolver.asc => hisparse/hisparse_resolver.cpp} (92%) rename csrc/{hisparse_runtime/torch_binding.cpp => hisparse/hisparse_resolver_torch_adpt.h} (88%) delete mode 100644 csrc/hisparse_runtime/build.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 5250aa299627..dbb2fe8328c7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -235,9 +235,8 @@ else() install(TARGETS vllm_ascend_C vllm_ascend_kernels DESTINATION ${VLLM_ASCEND_INSTALL_PATH}) endif() -# The optional HiSparse resolver uses SIMT, independently of Attention kernels. -# Older toolkits and other SOCs keep their ordinary build unchanged; selecting -# HiSparse without this extension raises an explicit availability error. +# SIMT HiSparse kernels use CANN's ASC CMake language integration. Other SOCs +# retain their existing build and report HiSparse unavailable when selected. if(SOC_VERSION MATCHES "ascend950") find_path(HISPARSE_SIMT_INCLUDE_DIR simt_api/asc_simt.h PATHS @@ -245,19 +244,8 @@ if(SOC_VERSION MATCHES "ascend950") ${ASCEND_HOME_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include NO_DEFAULT_PATH) if(HISPARSE_SIMT_INCLUDE_DIR) - set(HISPARSE_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/csrc/hisparse_runtime) - set(HISPARSE_OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/hisparse_runtime) - add_custom_command( - OUTPUT ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so - COMMAND ${PYTHON_EXECUTABLE} ${HISPARSE_SOURCE}/build.py - --cann-root ${ASCEND_HOME_PATH} --output ${HISPARSE_OUTPUT} - --npu-arch dav-3510 --resolver-ub-bytes 175616 - DEPENDS ${HISPARSE_SOURCE}/build.py ${HISPARSE_SOURCE}/resolver.asc ${HISPARSE_SOURCE}/torch_binding.cpp - VERBATIM) - add_custom_target(hisparse_resolver - DEPENDS ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so) - add_dependencies(vllm_ascend_C hisparse_resolver) - install(FILES ${HISPARSE_OUTPUT}/hisparse_resolver60103_args.so ${HISPARSE_OUTPUT}/torch_binding.so - DESTINATION ${VLLM_ASCEND_INSTALL_PATH}/hisparse) + add_subdirectory(csrc/hisparse) + target_compile_definitions(vllm_ascend_C PRIVATE VLLM_ASCEND_ENABLE_HISPARSE_RESOLVER) + target_link_libraries(vllm_ascend_C PRIVATE vllm_ascend_hisparse) endif() endif() diff --git a/csrc/hisparse/CMakeLists.txt b/csrc/hisparse/CMakeLists.txt new file mode 100644 index 000000000000..ddee06d4e7c3 --- /dev/null +++ b/csrc/hisparse/CMakeLists.txt @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# CANN's ASC integration supplies the compiler, archive rules and runtime links. +find_package(ASC REQUIRED CONFIG PATHS + ${ASCEND_HOME_PATH}/lib64/cmake + ${ASCEND_HOME_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/lib64/cmake + NO_DEFAULT_PATH) +set(CMAKE_ASC_ARCHITECTURES dav-3510) +enable_language(ASC) + +ascendc_library(vllm_ascend_hisparse STATIC hisparse_resolver.cpp) +target_include_directories(vllm_ascend_hisparse PRIVATE ${HISPARSE_SIMT_INCLUDE_DIR}) +target_compile_options(vllm_ascend_hisparse PRIVATE --enable-simt -O2) +# Reserve space for the descriptor and runtime overhead within the target UB. +target_compile_definitions(vllm_ascend_hisparse PRIVATE HISPARSE_RESOLVER_UB_BYTES=175616) diff --git a/csrc/hisparse_runtime/resolver.asc b/csrc/hisparse/hisparse_resolver.cpp similarity index 92% rename from csrc/hisparse_runtime/resolver.asc rename to csrc/hisparse/hisparse_resolver.cpp index f67902500fe1..404b21cb259b 100644 --- a/csrc/hisparse_runtime/resolver.asc +++ b/csrc/hisparse/hisparse_resolver.cpp @@ -539,7 +539,7 @@ __simt_vf__ __launch_bounds__(kThreads) inline void resolve_vf( } // namespace template -__global__ __vector__ void hisparse_union_resolve60103(uint64_t a0, uint64_t a1, uint64_t a2, uint64_t a3, uint64_t a4, uint64_t a5, uint64_t a6, uint64_t a7, uint64_t a8, uint64_t a9, uint64_t a10, uint64_t a11, uint64_t a12, uint64_t a13, uint64_t a14, uint64_t a15, uint64_t a16, uint64_t a17, uint64_t a18, uint64_t a32, uint64_t a33, uint64_t a34, uint64_t a35, uint64_t a36, uint64_t a37, uint64_t a38, uint64_t a39, uint64_t a40, uint64_t a41, uint64_t a42, uint64_t a43, uint64_t a44, uint64_t a45, uint64_t a46, uint64_t a47, uint64_t a48, uint64_t a49, uint64_t a50, uint64_t a51, uint64_t a52, uint64_t a53, uint64_t a54, uint64_t a55, uint64_t a56) { +__global__ __vector__ void hisparse_resolve_residency(uint64_t a0, uint64_t a1, uint64_t a2, uint64_t a3, uint64_t a4, uint64_t a5, uint64_t a6, uint64_t a7, uint64_t a8, uint64_t a9, uint64_t a10, uint64_t a11, uint64_t a12, uint64_t a13, uint64_t a14, uint64_t a15, uint64_t a16, uint64_t a17, uint64_t a18, uint64_t a32, uint64_t a33, uint64_t a34, uint64_t a35, uint64_t a36, uint64_t a37, uint64_t a38, uint64_t a39, uint64_t a40, uint64_t a41, uint64_t a42, uint64_t a43, uint64_t a44, uint64_t a45, uint64_t a46, uint64_t a47, uint64_t a48, uint64_t a49, uint64_t a50, uint64_t a51, uint64_t a52, uint64_t a53, uint64_t a54, uint64_t a55, uint64_t a56) { extern __ubuf__ int32_t local_scratch[]; const uint64_t words = (2 * a54 + a34 + 2 * ((a34 + 31) / 32) + 64 + 15) / 16 * 16; auto descriptor = (__ubuf__ uint64_t*)(local_scratch + words); @@ -594,10 +594,10 @@ extern "C" uint64_t hisparse_resolver_scratch_capacity() { return kScratchBudgetBytes; } -extern "C" void launch_hisparse_union_resolve60103(void** pointers, uint64_t* geometry, int rows, void* stream) { +extern "C" void launch_hisparse_resolve_residency(void** pointers, uint64_t* geometry, int rows, void* stream) { if (geometry[1] == 2048 && geometry[2] == 10240 && geometry[10] == 128 && geometry[14] == 128 && geometry[16] == 128 && geometry[22] == 16384) { - hisparse_union_resolve60103<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); + hisparse_resolve_residency<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); } else { - hisparse_union_resolve60103<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); + hisparse_resolve_residency<<>>(uint64_t(reinterpret_cast(pointers[0])), uint64_t(reinterpret_cast(pointers[1])), uint64_t(reinterpret_cast(pointers[2])), uint64_t(reinterpret_cast(pointers[3])), uint64_t(reinterpret_cast(pointers[4])), uint64_t(reinterpret_cast(pointers[5])), uint64_t(reinterpret_cast(pointers[6])), uint64_t(reinterpret_cast(pointers[7])), uint64_t(reinterpret_cast(pointers[8])), uint64_t(reinterpret_cast(pointers[9])), uint64_t(reinterpret_cast(pointers[10])), uint64_t(reinterpret_cast(pointers[11])), uint64_t(reinterpret_cast(pointers[12])), uint64_t(reinterpret_cast(pointers[13])), uint64_t(reinterpret_cast(pointers[14])), uint64_t(reinterpret_cast(pointers[15])), uint64_t(reinterpret_cast(pointers[16])), uint64_t(reinterpret_cast(pointers[17])), uint64_t(reinterpret_cast(pointers[18])), geometry[0], geometry[1], geometry[2], geometry[3], geometry[4], geometry[5], geometry[6], geometry[7], geometry[8], geometry[9], geometry[10], geometry[11], geometry[12], geometry[13], geometry[14], geometry[15], geometry[16], geometry[17], geometry[18], geometry[19], geometry[20], geometry[21], geometry[22], geometry[23], geometry[24]); } } diff --git a/csrc/hisparse_runtime/torch_binding.cpp b/csrc/hisparse/hisparse_resolver_torch_adpt.h similarity index 88% rename from csrc/hisparse_runtime/torch_binding.cpp rename to csrc/hisparse/hisparse_resolver_torch_adpt.h index 6db1d8af9b55..0177fd280aed 100644 --- a/csrc/hisparse_runtime/torch_binding.cpp +++ b/csrc/hisparse/hisparse_resolver_torch_adpt.h @@ -1,4 +1,5 @@ // SPDX-License-Identifier: Apache-2.0 +#pragma once // Tensor metadata is decoded here; launches keep TorchNPU queue ordering. #include #include @@ -13,15 +14,15 @@ #include #include -extern "C" void launch_hisparse_union_resolve60103(void**, uint64_t*, int, void*); +extern "C" void launch_hisparse_resolve_residency(void**, uint64_t*, int, void*); extern "C" uint64_t hisparse_resolver_scratch_capacity(); -namespace { +namespace vllm_ascend::hisparse { at::Tensor tensor_or_undefined(const std::optional& tensor) { return tensor.value_or(at::Tensor()); } -void resolve_tensors( +void resolve_residency( const at::Tensor& logical, const std::optional& requests, const std::optional& states, @@ -173,14 +174,14 @@ void resolve_tensors( command.SetCustomHandler([pointers, sizes, stream, owners = std::move(owners)]() mutable -> int { // Release queued tensor owners after submission, not on queue-slot reuse. auto submission_owners = std::move(owners); - launch_hisparse_union_resolve60103(pointers.data(), sizes.data(), int(sizes[0]), stream); + launch_hisparse_resolve_residency(pointers.data(), sizes.data(), int(sizes[0]), stream); return 0; }); command.Run(); } // The operator mutates caller-owned tensors and returns no new tensors. -void resolve_meta( +void resolve_residency_meta( const at::Tensor&, const std::optional&, const std::optional&, const std::optional&, const std::optional&, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, const std::optional&, @@ -188,18 +189,4 @@ void resolve_meta( const std::optional&, const std::optional&, const std::optional&, const std::optional&, const at::Tensor&, const std::optional&, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t) {} -} // namespace - -TORCH_LIBRARY(hisparse_fused60103, m) { - m.def("resolve_tensors(Tensor logical, Tensor? requests, Tensor? states, Tensor? source, Tensor? resident, " - "Tensor hot_table, Tensor(a!) dgi, Tensor(a!) lru, Tensor(a!) output, Tensor(a!)? attention, " - "Tensor(a!) resolved, Tensor(a!)? counts, Tensor(a!)? swap_host, Tensor(a!)? swap_device, " - "Tensor(a!)? swap_counts, Tensor(a!)? miss_mask, Tensor(a!)? stats, Tensor(a!) errors, Tensor? valid_rows, " - "int host_rows, int hot_block_size, int source_block_size, int resident_block_size, int resident_null_block, " - "int attention_block_stride, int union_rows, int attention_capacity=0) -> ()"); - m.def("scratch_capacity() -> int", []() -> int64_t { return hisparse_resolver_scratch_capacity(); }); - m.def("supports_attention_compaction() -> bool", []() -> bool { return true; }); - m.def("supports_budgeted_hash() -> bool", []() -> bool { return true; }); -} -TORCH_LIBRARY_IMPL(hisparse_fused60103, PrivateUse1, m) { m.impl("resolve_tensors", TORCH_FN(resolve_tensors)); } -TORCH_LIBRARY_IMPL(hisparse_fused60103, Meta, m) { m.impl("resolve_tensors", TORCH_FN(resolve_meta)); } +} // namespace vllm_ascend::hisparse diff --git a/csrc/hisparse_runtime/build.py b/csrc/hisparse_runtime/build.py deleted file mode 100644 index ac4c368d7c21..000000000000 --- a/csrc/hisparse_runtime/build.py +++ /dev/null @@ -1,95 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Build the HiSparse fused residency resolver for an explicit CANN target. - -Run in the matching A5 CANN 9.2 / torch 2.10 / torch-npu 2.10 environment. -Outputs stay in --output; installing into a serving environment is explicit. -""" - -import argparse -import importlib.util -import subprocess -from pathlib import Path - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--cann-root", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--npu-arch", required=True, help="SIMT-capable bisheng target architecture") - parser.add_argument( - "--resolver-ub-bytes", type=int, required=True, help="Usable per-AIV UB bytes, including descriptors" - ) - parser.add_argument("--install", action="store_true", help="Install built libraries into this checkout") - args = parser.parse_args() - if args.resolver_ub_bytes <= 512: - parser.error("--resolver-ub-bytes must leave space beyond the 512-byte descriptor") - import torch - - spec = importlib.util.find_spec("torch_npu") - if spec is None or spec.origin is None: - raise RuntimeError("torch_npu must be installed in this build environment") - npu = Path(spec.origin).parent - torch_root = Path(torch.__file__).parent - source = Path(__file__).resolve().parent - output = args.output.resolve() - output.mkdir(parents=True, exist_ok=True) - subprocess.run( - [ - str(args.cann_root / "bin/bisheng"), - "-x", - "asc", - f"--npu-arch={args.npu_arch}", - f"-DHISPARSE_RESOLVER_UB_BYTES={args.resolver_ub_bytes}", - "-O2", - "-shared", - "-fPIC", - "-std=c++17", - str(source / "resolver.asc"), - "-o", - str(output / "hisparse_resolver60103_args.so"), - ], - check=True, - ) - common = [ - "g++", - "-shared", - "-fPIC", - "-std=c++17", - "-O2", - "-D_GLIBCXX_USE_CXX11_ABI=" + str(int(torch._C._GLIBCXX_USE_CXX11_ABI)), - "-I" + str(npu / "include"), - "-I" + str(torch_root / "include"), - "-I" + str(torch_root / "include/torch/csrc/api/include"), - "-I" + str(args.cann_root / "include"), - "-L" + str(torch_root / "lib"), - "-L" + str(npu / "lib"), - "-L" + str(args.cann_root / "lib64"), - ] - subprocess.run( - common - + [ - str(source / "torch_binding.cpp"), - "-o", - str(output / "torch_binding.so"), - "-L" + str(output), - "-l:hisparse_resolver60103_args.so", - "-Wl,-rpath,$ORIGIN", - "-ltorch", - "-ltorch_cpu", - "-lc10", - "-ltorch_npu", - "-lascendcl", - ], - check=True, - ) - - if args.install: - import shutil - - package = source.parents[1] / "vllm_ascend" - for name in ("hisparse_resolver60103_args.so", "torch_binding.so"): - shutil.copy2(output / name, package / "hisparse" / name) - - -if __name__ == "__main__": - main() diff --git a/csrc/torch_binding.cpp b/csrc/torch_binding.cpp index bf2d24d0b760..045f417d3873 100644 --- a/csrc/torch_binding.cpp +++ b/csrc/torch_binding.cpp @@ -55,6 +55,9 @@ #include "attention/kda_layout_swap12/kda_layout_swap12_torch_adpt.h" #include "attention/recurrent_gated_delta_rule_v310/recurrent_gated_delta_rule_310_torch_adpt.h" #include "attention/k2q_csr/k2q_csr_torch_adpt.h" +#ifdef VLLM_ASCEND_ENABLE_HISPARSE_RESOLVER +#include "hisparse/hisparse_resolver_torch_adpt.h" +#endif #include "attention/sparse_attention_score/sparse_attention_score_torch_adpt.h" #include "moe/dequant_situ_quant/dequant_situ_quant_torch_adpt.h" #include "moe/situ_mx_quant/situ_mx_quant_torch_adpt.h" @@ -2804,6 +2807,18 @@ TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) // Pybind on other platform TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) { +#ifdef VLLM_ASCEND_ENABLE_HISPARSE_RESOLVER + ops.def("hisparse_resolve_residency(Tensor logical, Tensor? requests, Tensor? states, Tensor? source, Tensor? resident, " + "Tensor hot_table, Tensor(a!) dgi, Tensor(a!) lru, Tensor(a!) output, Tensor(a!)? attention, " + "Tensor(a!) resolved, Tensor(a!)? counts, Tensor(a!)? swap_host, Tensor(a!)? swap_device, " + "Tensor(a!)? swap_counts, Tensor(a!)? miss_mask, Tensor(a!)? stats, Tensor(a!) errors, Tensor? valid_rows, " + "int host_rows, int hot_block_size, int source_block_size, int resident_block_size, int resident_null_block, " + "int attention_block_stride, int union_rows, int attention_capacity=0) -> ()"); + ops.impl("hisparse_resolve_residency", torch::kPrivateUse1, &vllm_ascend::hisparse::resolve_residency); + ops.impl("hisparse_resolve_residency", torch::kMeta, &vllm_ascend::hisparse::resolve_residency_meta); + ops.def("hisparse_resolver_scratch_capacity() -> int", + []() -> int64_t { return hisparse_resolver_scratch_capacity(); }); +#endif ops.def("get_npu_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor?"); ops.impl("get_npu_view_from_cpu_tensor", c10::DispatchKey::CompositeExplicitAutograd, &vllm_ascend::get_npu_view_from_cpu_tensor); diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index b92397ff7e47..68127dc9baa5 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -21,25 +21,19 @@ are outside this backend's current support. ## Build the resolver Use the paired merged vLLM version in `.github/vllm-main-verified.commit` and -install vllm-ascend normally. Its CMake build includes the fused resolver when -the Ascend 950 SIMT headers are available. A separate resolver rebuild in the -same Torch/torch-npu/CANN environment is also supported: - -```bash -python csrc/hisparse_runtime/build.py \ - --cann-root "$ASCEND_HOME_PATH" \ - --output build/hisparse-resolver \ - --npu-arch dav-3510 \ - --resolver-ub-bytes 175616 \ - --install -``` - -Use a dedicated output directory. Select architecture and usable per-AIV UB -budget for the actual build target; these are hardware/compiler parameters. +install vllm-ascend normally in the matching Torch/torch-npu/CANN environment. +The standard CMake build compiles `csrc/hisparse/hisparse_resolver.cpp` using +CANN's ASC integration when the Ascend 950 SIMT headers are available. It links +the kernel into `vllm_ascend_C` and registers +`torch.ops._C_ascend.hisparse_resolve_residency` with the other Ascend operators; +there is no separate resolver build or library-loading step. + +The architecture and usable per-AIV UB budget are selected for the build target +in `csrc/hisparse/CMakeLists.txt`; these are hardware/compiler parameters. The kernel sizes its hash table from the request union and available UB. An unsupported capacity or missing extension raises an error before execution; -there is no Triton resolver fallback. Other vllm-ascend features do not load this -optional extension. +there is no Triton resolver fallback. Builds for other targets retain their +existing operators without registering this resolver. ## Enable HiSparse diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 7537b0fcf6f6..2a05b66bdc3b 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -132,7 +132,7 @@ def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups # Loading a device extension is the only mocked boundary; the group-owned # tensors below are real CPU storage, not a replacement resolver algorithm. - native = NS(scratch_capacity=lambda: 175616) + native = NS(hisparse_resolver_scratch_capacity=lambda: 175616) monkeypatch.setattr(resolver, "_load_resolver", lambda: native) backend = AscendHiSparseBackend() @@ -188,7 +188,7 @@ def test_native_resolver_uses_available_ub_without_truncating_the_union(): def test_missing_native_resolver_raises_instead_of_falling_back(monkeypatch): from vllm_ascend.hisparse import resolver - monkeypatch.setattr(resolver.Path, "is_file", lambda path: False) + monkeypatch.setattr(resolver.torch.ops, "_C_ascend", NS()) with pytest.raises(RuntimeError, match="requires its compiled AscendC resolver"): resolver._load_resolver() diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index c1d00eabbcad..2c29daa2010a 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -60,7 +60,6 @@ def check_available(self) -> None: install() # Import the device kernels before the connector allocates Host storage. importlib.import_module("vllm_ascend.hisparse.cache_ops") - _load_resolver() if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): # A5 disables the global custom-op selection policy. HiSparse still # needs its explicit DMA dependency; load the extension lazily here @@ -72,6 +71,7 @@ def check_available(self) -> None: raise RuntimeError("Ascend HiSparse requires the compiled Ascend DMA extension.") from error if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") + _load_resolver() def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Tensor]: from vllm_ascend.hisparse.cache_ops import allocate_mapped_host diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 5e169d05bcf1..054dcfc64df7 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -122,7 +122,7 @@ def __init__(self, max_rows: int, hot_size: int, max_union_rows: int, device: to from vllm_ascend.hisparse.resolver import _load_resolver, scratch_layout self.native = _load_resolver() - scratch_layout(hot_size, max_union_rows, self.native.scratch_capacity()) + scratch_layout(hot_size, max_union_rows, self.native.hisparse_resolver_scratch_capacity()) self.first_refs = torch.empty(max_rows * hot_size, dtype=torch.int32, device=device) self.errors = torch.zeros((max_rows, 8), dtype=torch.int32, device=device) diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py index bee43b434f2b..4ef0ad260492 100644 --- a/vllm_ascend/hisparse/resolver.py +++ b/vllm_ascend/hisparse/resolver.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 """Tensor-only native resolver dispatch; no address-keyed launch ownership.""" -from pathlib import Path - import torch @@ -28,16 +26,11 @@ def scratch_layout(hot, max_union_rows, capacity=None): def _load_resolver(): - library = Path(__file__).with_name("torch_binding.so") - if not library.is_file(): - raise RuntimeError("HiSparse requires its compiled AscendC resolver. Build csrc/hisparse_runtime first.") - if str(library) not in torch.ops.loaded_libraries: - torch.ops.load_library(str(library)) - namespace = torch.ops.hisparse_fused60103 - if not all(hasattr(namespace, name) for name in ("scratch_capacity", "resolve_tensors", "supports_budgeted_hash")): - raise RuntimeError("Rebuild the HiSparse native extension for the budget-aware Tensor resolver ABI") - if not namespace.supports_budgeted_hash(): - raise RuntimeError("HiSparse requires a native resolver with budget-aware hash support") + namespace = torch.ops._C_ascend + if not all( + hasattr(namespace, name) for name in ("hisparse_resolver_scratch_capacity", "hisparse_resolve_residency") + ): + raise RuntimeError("HiSparse requires its compiled AscendC resolver. Rebuild vllm-ascend for this target.") return namespace @@ -48,7 +41,7 @@ def resolve_tensors(native, arguments, resolved_scratch, errors): Callers must destroy graphs before releasing or rebinding those owners. """ a = arguments - native.resolve_tensors( + native.hisparse_resolve_residency( a["global_indices"], a["request_ids"], a["request_state_indices"], From 5934d4ac3b8ae5a9fb83fa9d5fca707e247621a5 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 17:24:51 +0800 Subject: [PATCH 022/120] fix(hisparse): retain mixed AscendC SIMT compilation mode Signed-off-by: zhenwenqi2024 --- csrc/hisparse/CMakeLists.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/csrc/hisparse/CMakeLists.txt b/csrc/hisparse/CMakeLists.txt index ddee06d4e7c3..a08c06ca6b43 100644 --- a/csrc/hisparse/CMakeLists.txt +++ b/csrc/hisparse/CMakeLists.txt @@ -9,6 +9,8 @@ enable_language(ASC) ascendc_library(vllm_ascend_hisparse STATIC hisparse_resolver.cpp) target_include_directories(vllm_ascend_hisparse PRIVATE ${HISPARSE_SIMT_INCLUDE_DIR}) -target_compile_options(vllm_ascend_hisparse PRIVATE --enable-simt -O2) +# The AscendC entry launches a SIMT vector function through asc_vf_call. Keep +# the mixed programming mode; --enable-simt selects a pure SIMT translation unit. +target_compile_options(vllm_ascend_hisparse PRIVATE -O2) # Reserve space for the descriptor and runtime overhead within the target UB. target_compile_definitions(vllm_ascend_hisparse PRIVATE HISPARSE_RESOLVER_UB_BYTES=175616) From 056aa918ad0612bbea4e17f7ce39edfd508c1ee3 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 17:31:56 +0800 Subject: [PATCH 023/120] test(hisparse): preserve copy stream during graph capture Signed-off-by: zhenwenqi2024 --- tests/e2e/pull_request/one_card/test_hisparse_ops.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index e6123f872fa7..e8d5cdb169f7 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -410,13 +410,13 @@ def forward(): torch.npu.synchronize() torch.testing.assert_close(sc.cpu(), torch.zeros(query_rows, dtype=torch.int32)) torch.testing.assert_close(stats.cpu(), torch.tensor([union, union], dtype=torch.int64)) - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): + capture_stream = torch.npu.Stream() + capture_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(capture_stream): forward() torch.npu.synchronize() graph = torch.npu.NPUGraph() - with torch.npu.graph(graph, stream=stream): + with torch.npu.graph(graph, stream=capture_stream): forward() dgi.fill_(-1) lru.copy_(torch.arange(hot_size, dtype=torch.int16)[None, :]) From 37a42a2caa77de17b14df89ec0ddb8df4b5f6b9f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 20:19:02 +0800 Subject: [PATCH 024/120] perf(hisparse): reuse native C8 Prolog with owned staging Keep Core lifecycle and Resident row mapping. Qualify this private candidate on NPU before publishing any performance claim. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_kv_preprocess.py | 101 ++++++++++++++++++ vllm_ascend/attention/sfa_hisparse.py | 53 ++++++++- 2 files changed, 151 insertions(+), 3 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index d97b05bb7cad..2c2e41d510ec 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -113,3 +113,104 @@ def forward(): torch.npu.synchronize() if graph is not None: graph.reset() + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("rows", [4, 16]) +def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): + _setup() + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.attention.sfa_v1 import AscendSFAImpl + from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod + + torch.manual_seed(71) + device, block, width = "npu", 128, 656 + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl._quant_type = AscendW8A8DynamicLinearMethod + impl.enable_sparse_sfa_c8 = True + impl.c8_cache_dtype = torch.float8_e4m3fn + impl.q_lora_rank, impl.kv_lora_rank, impl.qk_rope_head_dim = 2048, 512, 64 + impl.local_num_heads, impl.has_indexer = 8, True + impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False + impl.sfa_qsfa_tile_size = block + impl._kv_staging_capacity = block + impl.fused_qkv_a_proj = SimpleNamespace( + weight=torch.randint(-4, 5, (6144, 2624), dtype=torch.int8, device=device), + weight_scale=torch.full((2624,), 0.02, dtype=torch.float32, device=device), + ) + impl.q_proj = SimpleNamespace( + weight=torch.randint(-4, 5, (2048, 2048), dtype=torch.int8, device=device), + weight_scale=torch.full((2048,), 0.02, dtype=torch.float32, device=device), + ) + impl.W_UK_T = torch.randn(8, 192, 512, dtype=torch.bfloat16, device=device) * 0.02 + impl.q_a_layernorm = SimpleNamespace( + weight=torch.ones(2048, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 + ) + impl.kv_a_layernorm = SimpleNamespace( + weight=torch.ones(512, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 + ) + # Use the parent's production quantized-weight preparation and Prolog as + # the oracle. This case tests ownership/addressing, not Prolog arithmetic. + impl._process_weights_for_fused_prolog_v3() + impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16, device=device) + staging = torch.empty((1, block, 1, width), dtype=impl.c8_cache_dtype, device=device) + local_slots = torch.arange(block, dtype=torch.int64, device=device) + impl._c8_prolog_workspace = lambda: (staging, local_slots) + raw = torch.full((3, 2, block, width), 90, dtype=torch.uint8, device=device) + destination = raw[:, 1].view(impl.c8_cache_dtype) + logical = [-1, 129, 0, 257, 128][:rows] + list(range(16, 16 + max(0, rows - 5))) + slots = torch.tensor(logical, dtype=torch.int64, device=device) + impl.hisparse_cache = SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)) + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, rows], dtype=torch.int32, device=device), hisparse_for_capture=True + ) + x = torch.randn(rows, 6144, dtype=torch.bfloat16, device=device) + cos = torch.ones(rows, 1, 1, 64, dtype=x.dtype, device=device) + sin = torch.zeros_like(cos) + reference = torch.empty_like(staging) + + def forward(): + return impl._sfa_preprocess_prolog_v3(x, (), cos, sin, slots, attn_metadata=metadata) + + def compare(actual, expected): + if isinstance(actual, tuple): + assert len(actual) == len(expected) + for a, e in zip(actual, expected): + compare(a, e) + elif actual is not None: + torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0) + else: + assert expected is None + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + captured = forward() + for valid in (rows - 1, 2): + x.neg_() + metadata.query_start_loc[-1:].fill_(valid) + raw.fill_(90) + if graph is None: + actual = forward() + else: + graph.replay() + actual = captured + expected = AscendSFAImpl._sfa_preprocess_prolog_v3( + impl, x, (reference,), cos, sin, local_slots[:rows], attn_metadata=metadata + ) + torch.npu.synchronize() + compare(actual, expected) + payload = reference.view(torch.uint8).reshape(block, width).cpu() + wanted = torch.full(raw.shape, 90, dtype=torch.uint8) + for row, slot in enumerate(logical): + if row < valid and slot >= 0: + wanted[slot // block, 1, slot % block] = payload[row] + torch.testing.assert_close(raw.cpu(), wanted, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 0ffce7ab66aa..3ecafb286ee3 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -11,6 +11,7 @@ from vllm.v1.worker.workspace import current_workspace_manager from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder +from vllm_ascend.attention.utils import PreprocessType from vllm_ascend.hisparse.prefill import AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan # Bound the temporary D2D packing workspace independently of prefill length. @@ -134,9 +135,11 @@ def __init__(self, *args, **kwargs): self._kv_staging_capacity = (max_rows + 127) // 128 * 128 def _fused_preprocess_type(self): - # Fused prologs write the ordinary split cache directly. HiSparse owns - # both the resident destination and optional host-mirror staging. This - # also keeps quantized linear weight layout under its normal owner. + # Reuse SFA's supported quantization/weight preparation. Its C8 Prolog + # writes owned staging below; the HiSparse handle retains the final + # Resident destination, row mapping and Host-mirror completion policy. + if self.enable_sparse_sfa_c8: + return super()._fused_preprocess_type() return None def _compose_sfa_kv_cache(self, kv_cache): @@ -188,9 +191,53 @@ def allocate(): return current_workspace_manager().get_persistent_resource(("ascend_hisparse_c8_lengths",), allocate) + def _c8_prolog_workspace(self): + from vllm_ascend.hisparse.c8 import C8_ROW_BYTES + + def allocate(): + device = self.index_group.logical_topk_indices.device + return ( + torch.empty( + (self._kv_staging_capacity // 128, 128, 1, C8_ROW_BYTES), + dtype=self.c8_cache_dtype, + device=device, + ), + torch.arange(self._kv_staging_capacity, dtype=torch.int64, device=device), + ) + + # Like the existing preprocessing workspace, these rows are consumed + # on the compute stream before the next layer/graph can overwrite them. + # Host mirroring reads the final Resident cache, never this scratch. + return current_workspace_manager().get_persistent_resource( + ("ascend_hisparse_c8_prolog", self._kv_staging_capacity, self.c8_cache_dtype), allocate + ) + + def _sfa_preprocess_prolog_v3(self, hidden_states, kv_cache, cos, sin, slot_mapping, *, attn_metadata=None): + from vllm_ascend.hisparse.c8 import C8_ROW_BYTES + from vllm_ascend.hisparse.cache_ops import scatter_packed_kv + + if attn_metadata is None: + raise ValueError("HiSparse Prolog requires current row-validity metadata") + rows = hidden_states.shape[0] + if rows > self._kv_staging_capacity: + raise ValueError("HiSparse Prolog exceeds its profiled token capacity") + staging, local_slots = self._c8_prolog_workspace() + result = super()._sfa_preprocess_prolog_v3( + hidden_states, (staging,), cos, sin, local_slots[:rows], attn_metadata=attn_metadata + ) + destination, mapping, count = self.hisparse_cache.write_target(rows, slot_mapping.numel()) + packed = staging.view(-1, C8_ROW_BYTES)[:count].view(torch.uint8) + scatter_packed_kv( + packed, packed[:, :0], destination.view(torch.uint8), mapping, attn_metadata.query_start_loc[-1:] + ) + finish_hisparse_kv_update(self.hisparse_cache, attn_metadata) + return result + def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=None): if self.enable_sparse_sfa_c8: self._c8_length_workspace() + if self.preprocess_type == PreprocessType.PROLOG_V3: + self._c8_prolog_workspace() else: self._workspace(hidden_states.dtype) if self.qk_rope_head_dim and not self.enable_sparse_sfa_c8: From 5b786fcfeea52b6e2cb8a45440e88b834bc1c94f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 20:23:42 +0800 Subject: [PATCH 025/120] test(hisparse): match native C8 Prolog MXFP8 contract Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_kv_preprocess.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index 2c2e41d510ec..0c9362336b34 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -121,12 +121,12 @@ def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): _setup() from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl from vllm_ascend.attention.sfa_v1 import AscendSFAImpl - from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod + from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod torch.manual_seed(71) device, block, width = "npu", 128, 656 impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) - impl._quant_type = AscendW8A8DynamicLinearMethod + impl._quant_type = AscendW8A8MXFP8DynamicLinearMethod impl.enable_sparse_sfa_c8 = True impl.c8_cache_dtype = torch.float8_e4m3fn impl.q_lora_rank, impl.kv_lora_rank, impl.qk_rope_head_dim = 2048, 512, 64 @@ -134,14 +134,17 @@ def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False impl.sfa_qsfa_tile_size = block impl._kv_staging_capacity = block - impl.fused_qkv_a_proj = SimpleNamespace( - weight=torch.randint(-4, 5, (6144, 2624), dtype=torch.int8, device=device), - weight_scale=torch.full((2624,), 0.02, dtype=torch.float32, device=device), - ) - impl.q_proj = SimpleNamespace( - weight=torch.randint(-4, 5, (2048, 2048), dtype=torch.int8, device=device), - weight_scale=torch.full((2048,), 0.02, dtype=torch.float32, device=device), - ) + + def mx_layer(inputs, outputs): + weight, scale = torch_npu.npu_dynamic_mx_quant( + torch.randn(outputs, inputs, dtype=torch.bfloat16, device=device) * 0.02, + dst_type=torch.float8_e4m3fn, + ) + return SimpleNamespace(weight=weight.T.contiguous(), weight_scale=scale.transpose(0, 1).contiguous()) + + impl.fused_qkv_a_proj = mx_layer(6144, 2624) + impl.q_proj = mx_layer(2048, 2048) + impl.sfa_qsfa_k_nope_clip_alpha = None impl.W_UK_T = torch.randn(8, 192, 512, dtype=torch.bfloat16, device=device) * 0.02 impl.q_a_layernorm = SimpleNamespace( weight=torch.ones(2048, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 @@ -178,6 +181,8 @@ def compare(actual, expected): for a, e in zip(actual, expected): compare(a, e) elif actual is not None: + if actual.dtype == torch.float8_e4m3fn: + actual, expected = actual.view(torch.uint8), expected.view(torch.uint8) torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0) else: assert expected is None From e954dcffcd0a14f32aff27938ea59254f2fc54a5 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 20:48:40 +0800 Subject: [PATCH 026/120] perf(hisparse): consume native compact C8 indices in Attention Keep virtual refill rows separate and qualify bound HMA geometry. Leave Prefill mapping unchanged. Private candidate pending actual NPU and model qualification. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_attention.py | 19 ++++++++++++---- .../one_card/test_hisparse_ops.py | 21 +++++++++++++++--- vllm_ascend/attention/sfa_hisparse.py | 22 ++++++++++++++----- vllm_ascend/hisparse/backend.py | 22 ++++++++++++++++++- vllm_ascend/hisparse/cache_ops.py | 11 +++++++++- vllm_ascend/hisparse/resolver.py | 14 +++++++++--- 6 files changed, 91 insertions(+), 18 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 6c763dc66058..d1112e49c03c 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -111,7 +111,8 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) -def test_c8_attention_uses_existing_native_operator(strided, capture): +@pytest.mark.parametrize("indices_are_mapped", [False, True]) +def test_c8_attention_uses_existing_native_operator(strided, capture, indices_are_mapped): from vllm_ascend.device.device_config import is_950 if not is_950(): @@ -141,7 +142,17 @@ def test_c8_attention_uses_existing_native_operator(strided, capture): indices[0] = -1 indices[0, :3] = torch.tensor([3, 3, 7], dtype=torch.int32) indices[-1] = -1 - selected = indices.to(device) + + def selection(): + if not indices_are_mapped: + return indices + output = torch.full_like(indices, -1) + for original, target in zip(indices.tolist(), output): + values = [value // block * (cache.stride(0) // 656) + value % block for value in original if value >= 0] + target[: len(values)] = torch.tensor(values, dtype=torch.int32) + return output + + selected = selection().to(device) query_end, key_length = (torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) impl.enable_sparse_sfa_c8 = True @@ -150,7 +161,7 @@ def test_c8_attention_uses_existing_native_operator(strided, capture): impl._c8_length_workspace = lambda: (query_end, key_length) def forward(): - return impl._attention_physical(q, qr, cache, selected) + return impl._attention_physical(q, qr, cache, selected, indices_are_mapped=indices_are_mapped) graph = None try: @@ -162,7 +173,7 @@ def forward(): with torch.npu.graph(graph): output = forward() indices[0, :4] = torch.tensor([9, 2, 9, 0], dtype=torch.int32) - selected.copy_(indices) + selected.copy_(selection()) q.neg_() graph.replay() else: diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index e8d5cdb169f7..3d4cc1020f60 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -161,7 +161,8 @@ def reference(topk, requests, states, source, resident, hot_table, dgi, lru, blo @pytest.mark.parametrize("device_row_count", [False, True]) -def test_resolver_strides_padding_graph(device_row_count): +@pytest.mark.parametrize("compact_attention", [False, True]) +def test_resolver_strides_padding_graph(device_row_count, compact_attention): torch.npu.set_device(0) raw = torch.empty(8192, dtype=torch.uint8) offset = (-raw.data_ptr()) % 4096 @@ -192,7 +193,11 @@ def test_resolver_strides_padding_graph(device_row_count): lru = torch.cat((lru, torch.arange(8, 16, dtype=torch.int16)[None, :].expand(2, -1)), dim=1) expected = reference(topk, requests, states, source, resident, hot_table, dgi, lru, 4) workspace = ResolverWorkspace(4, 16, 16, torch.device("npu:0")) - ops = HiSparseNPUOps(mapping, workspace) + stride_rows = hot.stride(0) // hot.size(2) + capacity = (hot.size(0) - 1) * stride_rows + hot.size(1) + ops = HiSparseNPUOps( + mapping, workspace, attention_capacity=capacity if compact_attention else 0, attention_block_stride=stride_rows + ) topk_n, requests_n, states_n, source_n, resident_n, hot_table_n, dgi_n, lru_n = [ t.npu() for t in (topk, requests, states, source, resident, hot_table, dgi, lru) ] @@ -246,7 +251,16 @@ def forward(): ("indices", "counts", "swap_host", "swap_device", "swap_counts", "dgi", "lru", "stats"), actual, expected ): torch.testing.assert_close(a.cpu(), e, msg=name) - expected_attn = torch.where(expected[0] >= 0, expected[0] // 4 * 8 + expected[0] % 4, -1) + + def attention_reference(physical): + mapped = torch.where(physical >= 0, physical // 4 * 8 + physical % 4, -1) + if compact_attention: + for row in mapped: + valid = [value for value in row.tolist() if value >= 0] + row.copy_(torch.tensor(valid + [-1] * (row.numel() - len(valid)), dtype=row.dtype)) + return mapped + + expected_attn = attention_reference(expected[0]) torch.testing.assert_close(attn.cpu(), expected_attn) for row in range(4): for col in range(int(expected[4][row])): @@ -324,6 +338,7 @@ def forward(): torch.npu.synchronize() for actual_value, expected_value in zip(actual, oracle): torch.testing.assert_close(actual_value.cpu(), expected_value) + torch.testing.assert_close(attn.cpu(), attention_reference(oracle[0])) graph.reset() mapping.close() diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 3ecafb286ee3..baf832a10723 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -321,7 +321,7 @@ def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_s # which would duplicate the write and can overwrite unrelated slab rows. return k_pe, k_nope - def _c8_physical_attention(self, query, query_rope, cache, indices): + def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_are_mapped=False): """Adapt native packed rows to the existing quantized SFA operation. Keep compaction stable, including duplicates. The two current lengths are @@ -348,7 +348,11 @@ def _c8_physical_attention(self, query, query_rope, cache, indices): raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") if not math.isfinite(scale) or scale <= 0: raise ValueError("HiSparse C8 attention scale must be finite and positive") - mapped = compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) + mapped = ( + indices + if indices_are_mapped + else compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) + ) packed_query = torch.cat((query, query_rope), dim=-1) query_ends.fill_(rows) key_length.fill_(view.shape[0]) @@ -376,9 +380,11 @@ def _c8_physical_attention(self, query, query_rope, cache, indices): result = result[0] return torch.where((mapped >= 0).any(dim=-1)[:, None, None], result, 0) - def _attention_physical(self, query, query_rope, packed_cache, physical_topk): + def _attention_physical(self, query, query_rope, packed_cache, physical_topk, *, indices_are_mapped=False): if self.enable_sparse_sfa_c8: - return self._c8_physical_attention(query, query_rope, packed_cache, physical_topk) + return self._c8_physical_attention( + query, query_rope, packed_cache, physical_topk, indices_are_mapped=indices_are_mapped + ) from vllm_ascend.hisparse.cache_ops import pack_attention_rows latent, rope, block_table, indices, query_ends, seq_lens = self._workspace(query.dtype) @@ -437,8 +443,11 @@ def _execute_sparse_flash_attention_process( logical = topk_indices.squeeze(1) if topk_indices.ndim == 3 else topk_indices assert metadata.req_id_per_token is not None assert metadata.query_start_loc is not None - rows = torch.arange(logical.shape[0], dtype=torch.int32, device=logical.device) - logical = torch.where(rows[:, None] < metadata.query_start_loc[-1], logical, -1) + if metadata.num_prefills: + rows = torch.arange(logical.shape[0], dtype=torch.int32, device=logical.device) + logical = torch.where(rows[:, None] < metadata.query_start_loc[-1], logical, -1) + # Decode's resolver already masks graph padding from this same device + # scalar. Keep that validation at the producer of physical indices. self.index_group.set_logical_topk_ready(self.index_group_index) outputs = [] num_decode = metadata.num_decode_tokens @@ -461,6 +470,7 @@ def _execute_sparse_flash_attention_process( q_pe[:num_decode] if q_pe is not None else None, self.index_group.physical_kv_cache(self.index_group_index), physical, + indices_are_mapped=cache.runtime.cache_ops.attention_capacity > 0, ) ) if metadata.num_prefills: diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 2c29daa2010a..86f6b4eea6c4 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -112,7 +112,27 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: raise RuntimeError("Bind the HiSparse leader before follower caches.") workspace = leader_ops.workspace mapping = runtime.registered_host_pool._hisparse_registration - return HiSparseNPUOps(mapping, workspace) + attention_capacity = 0 + attention_block_stride = None + if runtime.kv_dtype == torch.float8_e4m3fn: + from vllm_ascend.hisparse.c8 import c8_storage_view + + # The resolver consumes virtual page rows. Attention consumes dense + # addresses relative to its own layer view, including any HMA gaps. + # Qualify the bound geometry once, without adding replay-time work. + view, block, attention_block_stride, _ = c8_storage_view(runtime.hot.cache) + attention_view, _, _, _ = c8_storage_view(runtime.hot.attention_cache) + attention_capacity = view.shape[0] + if block != runtime.hot.block_size or attention_capacity > attention_view.shape[0]: + raise ValueError("HiSparse C8 resolver and Attention views have incompatible geometry") + if runtime is not group.leader and ( + attention_capacity != group.leader.cache_ops.attention_capacity + or attention_block_stride != group.leader.cache_ops.attention_block_stride + ): + raise ValueError("Shared HiSparse C8 indices require matching per-layer page geometry") + return HiSparseNPUOps( + mapping, workspace, attention_capacity=attention_capacity, attention_block_stride=attention_block_stride + ) def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: from .binding import get_hisparse_indexer_specs diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 054dcfc64df7..79e8c529f52b 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -181,9 +181,18 @@ def _check_payload_rows(tensor, name, device, dtype, rows) -> None: class HiSparseNPUOps: """Runtime-owned implementations of the upstream HiSparse operator ABI.""" - def __init__(self, host_mapping: MappedHostTensor, workspace: ResolverWorkspace | None = None) -> None: + def __init__( + self, + host_mapping: MappedHostTensor, + workspace: ResolverWorkspace | None = None, + *, + attention_capacity: int = 0, + attention_block_stride: int | None = None, + ) -> None: self.host_mapping = host_mapping self.workspace = workspace + self.attention_capacity = attention_capacity + self.attention_block_stride = attention_block_stride def hisparse_resolve_residency( self, diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py index 4ef0ad260492..2e6466ce324f 100644 --- a/vllm_ascend/hisparse/resolver.py +++ b/vllm_ascend/hisparse/resolver.py @@ -34,7 +34,7 @@ def _load_resolver(): return namespace -def resolve_tensors(native, arguments, resolved_scratch, errors): +def resolve_tensors(native, arguments, resolved_scratch, errors, *, attention_capacity=0, attention_block_stride=None): """Pass live tensors; C++ derives geometry and protects queued storage. During replay, external tensors remain owned by the runtime/workspace. @@ -66,8 +66,9 @@ def resolve_tensors(native, arguments, resolved_scratch, errors): a["source_block_size"], a["resident_block_size"], a["resident_null_block"], - a["attention_block_stride"], + a["attention_block_stride"] if attention_block_stride is None else attention_block_stride, a["max_union_rows"], + attention_capacity, ) @@ -78,4 +79,11 @@ def resolve(owner, scope): resolved = scope["resolved_global_indices"] if resolved is None: resolved = workspace.first_refs[: rows * topk].view(rows, topk) - resolve_tensors(workspace.native, scope, resolved, workspace.errors[:rows]) + resolve_tensors( + workspace.native, + scope, + resolved, + workspace.errors[:rows], + attention_capacity=owner.attention_capacity, + attention_block_stride=owner.attention_block_stride, + ) From 172d3824ab5177164145a319058d84602bac3b11 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 20:57:33 +0800 Subject: [PATCH 027/120] fix(hisparse): retain native fallback for INT8 projection weights Signed-off-by: zhenwenqi2024 --- .../test_hisparse_foundation_cpu.py | 25 ++++++++++++++++++- vllm_ascend/attention/sfa_hisparse.py | 9 ++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 2a05b66bdc3b..dfaa979fc0ab 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -11,7 +11,11 @@ from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec, SparseCacheRole from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend -from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder, finish_hisparse_kv_update +from vllm_ascend.attention.sfa_hisparse import ( + AscendHiSparseImpl, + AscendHiSparseMetadataBuilder, + finish_hisparse_kv_update, +) from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.device.device_op import DeviceOperator @@ -30,6 +34,25 @@ ) +@pytest.mark.parametrize("mxfp8", [False, True]) +def test_c8_prolog_requires_matching_fp8_weight_scheme(mxfp8): + from vllm_ascend.attention.utils import PreprocessType + from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod, AscendW8A8MXFP8DynamicLinearMethod + + method = AscendW8A8MXFP8DynamicLinearMethod if mxfp8 else AscendW8A8DynamicLinearMethod + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl.enable_sparse_sfa_c8 = True + impl.fused_qkv_a_proj = NS(quant_method=NS(quant_method=method.__new__(method))) + impl.is_kv_consumer, impl.has_indexer = False, False + impl.dcp_group = None + impl.qk_rope_head_dim = 64 + impl.q_a_layernorm = impl.kv_a_layernorm = NS() + impl.q_proj = NS(_chunk_size=0) + impl.enable_sparse_sfa_turboquant = False + expected = PreprocessType.PROLOG_V3 if mxfp8 else None + assert impl._fused_preprocess_type() is expected + + @pytest.mark.parametrize("hisparse", [False, True]) def test_indexer_backend_choice_survives_model_construction_context(hisparse): from vllm.config import CacheConfig, VllmConfig, set_current_vllm_config diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index baf832a10723..c02f1a4e73dd 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -139,7 +139,14 @@ def _fused_preprocess_type(self): # writes owned staging below; the HiSparse handle retains the final # Resident destination, row mapping and Host-mirror completion policy. if self.enable_sparse_sfa_c8: - return super()._fused_preprocess_type() + from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod + + # Native Prolog ties its C8 payload dtype to the weight scheme. + # INT8 weights cannot write this bridge's E4M3FN cache; retain the + # existing native projection/quantizer path for that combination. + method = self._get_layer_quant_method(self.fused_qkv_a_proj) + if type(method) is AscendW8A8MXFP8DynamicLinearMethod: + return super()._fused_preprocess_type() return None def _compose_sfa_kv_cache(self, kv_cache): From 1be712f4723848015947cfd63975bab84fafb90b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 20:58:03 +0800 Subject: [PATCH 028/120] test(hisparse): initialize native ops for focused NPU selections Signed-off-by: zhenwenqi2024 --- tests/e2e/pull_request/one_card/test_hisparse_ops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 3d4cc1020f60..4b564bd4bc4e 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -24,6 +24,9 @@ def _require_mapped_host_device(): if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("HiSparse mapped-host operators currently target Ascend 950") torch.npu.set_device(0) + # Focused -k selections must load their native dependency independently; + # they cannot rely on the private-Host-pool test running first. + AscendHiSparseBackend().check_available() def test_private_host_pool_retains_pinned_storage_for_dma_writeback(): From 66e11b5e5b10e85b9b469f564b56116375f2104c Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 21:03:17 +0800 Subject: [PATCH 029/120] test(hisparse): declare dtype in workspace ownership fixture Signed-off-by: zhenwenqi2024 --- tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index dfaa979fc0ab..dcd82febb0fe 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -167,6 +167,7 @@ def group(): region_stride=hot_rows, max_union_rows=union_rows, device=torch.device("cpu"), + kv_dtype=torch.bfloat16, registered_host_pool=NS(_hisparse_registration=object()), ) owner.leader = leader From fc39fa8d36c05fd80fcce433df88e9ac821126c1 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 21:45:45 +0800 Subject: [PATCH 030/120] docs(hisparse): record current native C8 qualification and limits Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 31 ++++++++++++++----- release_notes/fragments/20926.yaml | 4 ++- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 68127dc9baa5..5199e86da54a 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -73,6 +73,13 @@ Native C8 main rows contain 512 E4M3FN bytes, 64 BF16 RoPE values and four FP32 scales (656 bytes total). Indexer rows contain 128 E4M3FN keys and one FP32 scale. Views preserve the shared-slab page strides and layer offsets. The existing quantizer and Attention operations are reused; payloads are copied as bytes. +With MXFP8 projection weights, C8 reuses SFA's existing Prolog V3 preprocessing +and publishes its staged rows through the HiSparse Resident mapping. Other +supported projection schemes retain the native preprocessing path. + +C8 Decode consumes the resolver's mapped, stably compacted indices directly. +The refill plan retains separate virtual page-row IDs. Prefill keeps its existing +layout bridge, including stable ordering and duplicate selections. Keep completion-before-unregister and graph allocation ownership intact. Overwriting KV still invalidates any old Hot copy. Equal batch shapes do not @@ -81,10 +88,20 @@ or unmerged Core PR is required. ## Validation and performance -The previous Triton-based foundation passed its CPU/NPU component suites and a -TP8/MTP3 full-model smoke. That evidence does not qualify this native-resolver -integration. The fused kernel is reused from the optimized implementation; -this foundation's integration and full-model results are recorded separately. - -Dataset accuracy and same-workload ON/OFF performance are not yet qualified for -this foundation. No throughput, TTFT or TPOT improvement is claimed here. +The native foundation has actual Ascend 950 eager/Graph component and TP8/MTP3 +full-model evidence. The C8 preprocessing and direct-index changes also passed +focused device checks and a fresh full-model ON/OFF comparison. See the PR for +the exact source pins, commands and qualification boundaries. + +In the current GLM-5.2 C4 comparison (32768 input and 8192 output tokens, +TP8/DP1, MTP3, Model Runner V2, prefix caching off and explicit +`weight_nz_mode=0`), both arms completed two rounds with 65536 actual output +token IDs per arm. ON's settled reported miss/H2D deltas were zero. Mean TPOT was +12.21 ms OFF and 14.20 ms ON: an observed increase of 16.29%. The second round +had identical output tokens and MTP acceptance, with an 18.11% TPOT increase; +the first round had different outputs/acceptance. These measurements do not +establish statistical significance or acceptable performance for other workloads. + +Dataset accuracy, the full-model miss path and a TPOT regression within 3% +remain unqualified. Profiling requests are separate from formal performance +timing; reported zero-miss counters do not prove the absence of device DMA. diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index 673f799094a5..591d1fa1def2 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -5,5 +5,7 @@ entries: Added experimental Ascend HiSparse BF16 and native C8 sparse MLA/SFA support using a fused AscendC residency resolver, existing Attention operators and Core lifecycle interfaces. The resolver requires its - compiled extension and does not fall back to Triton. + compiled extension and does not fall back to Triton. Native C8 reuses + supported SFA Prolog preprocessing and consumes the resolver's physical + Decode indices while preserving the separate refill plan. hardware: [A5] From 26c0252f5e564ff390e4488646c321675dea9bf3 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 22:05:56 +0800 Subject: [PATCH 031/120] perf(hisparse): reuse resolver counts for C8 empty-row masking Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_attention.py | 12 ++++++-- vllm_ascend/attention/sfa_hisparse.py | 30 ++++++++++++++++--- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index d1112e49c03c..2682bda3f068 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -111,8 +111,8 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) -@pytest.mark.parametrize("indices_are_mapped", [False, True]) -def test_c8_attention_uses_existing_native_operator(strided, capture, indices_are_mapped): +@pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_counts"]) +def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode): from vllm_ascend.device.device_config import is_950 if not is_950(): @@ -142,6 +142,7 @@ def test_c8_attention_uses_existing_native_operator(strided, capture, indices_ar indices[0] = -1 indices[0, :3] = torch.tensor([3, 3, 7], dtype=torch.int32) indices[-1] = -1 + indices_are_mapped = index_mode != "virtual" def selection(): if not indices_are_mapped: @@ -153,6 +154,7 @@ def selection(): return output selected = selection().to(device) + valid_counts = (indices >= 0).sum(dim=-1, dtype=torch.int32).to(device) if index_mode == "mapped_counts" else None query_end, key_length = (torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) impl.enable_sparse_sfa_c8 = True @@ -161,7 +163,9 @@ def selection(): impl._c8_length_workspace = lambda: (query_end, key_length) def forward(): - return impl._attention_physical(q, qr, cache, selected, indices_are_mapped=indices_are_mapped) + return impl._attention_physical( + q, qr, cache, selected, indices_are_mapped=indices_are_mapped, valid_counts=valid_counts + ) graph = None try: @@ -174,6 +178,8 @@ def forward(): output = forward() indices[0, :4] = torch.tensor([9, 2, 9, 0], dtype=torch.int32) selected.copy_(selection()) + if valid_counts is not None: + valid_counts.copy_((indices >= 0).sum(dim=-1, dtype=torch.int32)) q.neg_() graph.replay() else: diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index c02f1a4e73dd..8e77f1d75ad3 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -328,7 +328,7 @@ def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_s # which would duplicate the write and can overwrite unrelated slab rows. return k_pe, k_nope - def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_are_mapped=False): + def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_are_mapped=False, valid_counts=None): """Adapt native packed rows to the existing quantized SFA operation. Keep compaction stable, including duplicates. The two current lengths are @@ -353,6 +353,13 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a raise ValueError("HiSparse C8 inputs must share one NPU") if any(tensor.shape != (1,) or tensor.dtype != torch.int32 for tensor in (query_ends, key_length)): raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") + if valid_counts is not None and ( + not indices_are_mapped + or valid_counts.shape != (rows,) + or valid_counts.dtype != torch.int32 + or valid_counts.device != query.device + ): + raise ValueError("HiSparse C8 valid counts must describe the native mapped query rows") if not math.isfinite(scale) or scale <= 0: raise ValueError("HiSparse C8 attention scale must be finite and positive") mapped = ( @@ -361,6 +368,9 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a else compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) ) packed_query = torch.cat((query, query_rope), dim=-1) + # Read the resolver's current counts on this consumer stream before + # later groups can reuse their scratch. Keep the empty-row output guard. + has_keys = valid_counts > 0 if valid_counts is not None else (mapped >= 0).any(dim=-1) query_ends.fill_(rows) key_length.fill_(view.shape[0]) result = torch.ops._C_ascend.npu_kv_quant_sparse_flash_attention_vllm( @@ -385,12 +395,19 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a ) if isinstance(result, tuple): result = result[0] - return torch.where((mapped >= 0).any(dim=-1)[:, None, None], result, 0) + return torch.where(has_keys[:, None, None], result, 0) - def _attention_physical(self, query, query_rope, packed_cache, physical_topk, *, indices_are_mapped=False): + def _attention_physical( + self, query, query_rope, packed_cache, physical_topk, *, indices_are_mapped=False, valid_counts=None + ): if self.enable_sparse_sfa_c8: return self._c8_physical_attention( - query, query_rope, packed_cache, physical_topk, indices_are_mapped=indices_are_mapped + query, + query_rope, + packed_cache, + physical_topk, + indices_are_mapped=indices_are_mapped, + valid_counts=valid_counts, ) from vllm_ascend.hisparse.cache_ops import pack_attention_rows @@ -470,7 +487,11 @@ def _execute_sparse_flash_attention_process( logical_topk_indices=logical[:num_decode], block_size=metadata.block_size, num_valid_rows=metadata.query_start_loc[-1:], + return_valid_counts=self.enable_sparse_sfa_c8, ) + valid_counts = None + if self.enable_sparse_sfa_c8: + physical, valid_counts = physical outputs.append( self._attention_physical( ql_nope[:num_decode], @@ -478,6 +499,7 @@ def _execute_sparse_flash_attention_process( self.index_group.physical_kv_cache(self.index_group_index), physical, indices_are_mapped=cache.runtime.cache_ops.attention_capacity > 0, + valid_counts=valid_counts, ) ) if metadata.num_prefills: From a596bbf3b7a6e861c34285d5d0618f37af9beec7 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 22:34:59 +0800 Subject: [PATCH 032/120] perf(hisparse): prepare immutable C8 decode lengths at binding Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_attention.py | 20 ++++++++--- .../test_hisparse_foundation_cpu.py | 36 +++++++++++++++++++ vllm_ascend/attention/sfa_hisparse.py | 32 ++++++++++++----- vllm_ascend/hisparse/backend.py | 15 +++++++- vllm_ascend/hisparse/cache_ops.py | 2 ++ 5 files changed, 92 insertions(+), 13 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 2682bda3f068..d5f4f0585c4c 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -111,7 +111,7 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) -@pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_counts"]) +@pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_counts", "mapped_lengths"]) def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode): from vllm_ascend.device.device_config import is_950 @@ -119,7 +119,7 @@ def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode pytest.skip("This case uses the native A5 E4M3FN C8 payload; A2/A3 use an INT8 format") from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope - from vllm_ascend.hisparse.c8 import pack_c8_components + from vllm_ascend.hisparse.c8 import c8_storage_view, pack_c8_components torch.manual_seed(71) rows, heads, blocks, block, topk = 3, 8, 2, 128, 128 @@ -154,7 +154,19 @@ def selection(): return output selected = selection().to(device) - valid_counts = (indices >= 0).sum(dim=-1, dtype=torch.int32).to(device) if index_mode == "mapped_counts" else None + valid_counts = ( + (indices >= 0).sum(dim=-1, dtype=torch.int32).to(device) + if index_mode in ("mapped_counts", "mapped_lengths") + else None + ) + lengths = ( + ( + torch.arange(rows + 1, dtype=torch.int32, device=device)[rows:], + torch.full((1,), c8_storage_view(cache)[0].shape[0], dtype=torch.int32, device=device), + ) + if index_mode == "mapped_lengths" + else None + ) query_end, key_length = (torch.empty(1, dtype=torch.int32, device=device) for _ in range(2)) impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) impl.enable_sparse_sfa_c8 = True @@ -164,7 +176,7 @@ def selection(): def forward(): return impl._attention_physical( - q, qr, cache, selected, indices_are_mapped=indices_are_mapped, valid_counts=valid_counts + q, qr, cache, selected, indices_are_mapped=indices_are_mapped, valid_counts=valid_counts, lengths=lengths ) graph = None diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index dcd82febb0fe..bc0f5a784844 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -338,6 +338,42 @@ def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): assert backing[offset + 2 * stride_rows * width + 1] == 123 +@pytest.mark.parametrize("pages,stride_rows,max_rows", [(2, 128, 9), (3, 136, 37)]) +def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind(monkeypatch, pages, stride_rows, max_rows): + from vllm_ascend.hisparse import resolver + + monkeypatch.setattr(resolver, "_load_resolver", lambda: NS(hisparse_resolver_scratch_capacity=lambda: 175616)) + backend = AscendHiSparseBackend() + group = NS(copy_stream=NS(device=torch.device("cpu"))) + runtime = NS( + index_group=group, + max_swap_rows=max_rows, + region_stride=128, + max_union_rows=64, + device=torch.device("cpu"), + kv_dtype=torch.float8_e4m3fn, + registered_host_pool=NS(_hisparse_registration=object()), + ) + group.leader = runtime + + def bind(num_pages): + backing = torch.empty(num_pages * stride_rows * 656, dtype=torch.float8_e4m3fn) + cache = backing.as_strided((num_pages, 128, 656), (stride_rows * 656, 656, 1)) + runtime.hot = NS(cache=cache, attention_cache=cache, block_size=128) + return backend.create_cache_ops(runtime) + + first = bind(pages) + query, key = first.attention_lengths + for rows in (1, max_rows, 3, 1): + torch.testing.assert_close(query[rows : rows + 1], torch.tensor([rows], dtype=torch.int32)) + expected = (pages - 1) * stride_rows + 128 + torch.testing.assert_close(key, torch.tensor([expected], dtype=torch.int32)) + second = bind(pages + 1) + torch.testing.assert_close(second.attention_lengths[1], torch.tensor([expected + stride_rows], dtype=torch.int32)) + torch.testing.assert_close(key, torch.tensor([expected], dtype=torch.int32)) + assert key.data_ptr() != second.attention_lengths[1].data_ptr() + + def test_c8_compaction_preserves_duplicate_order_and_empty_rows(): indices = torch.tensor([[7, -1, 2, 7, -1, 0], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) expected = torch.tensor([[7, 2, 7, 0, -1, -1], [-1, -1, -1, -1, -1, -1]], dtype=torch.int32) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 8e77f1d75ad3..9ec3fa7ee449 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -328,16 +328,18 @@ def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_s # which would duplicate the write and can overwrite unrelated slab rows. return k_pe, k_nope - def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_are_mapped=False, valid_counts=None): + def _c8_physical_attention( + self, query, query_rope, cache, indices, *, indices_are_mapped=False, valid_counts=None, lengths=None + ): """Adapt native packed rows to the existing quantized SFA operation. - Keep compaction stable, including duplicates. The two current lengths are - device writes captured on the consumer stream; no Host read or stream/event - policy is introduced here. + Keep compaction stable, including duplicates. Decode lengths describe + binding/graph geometry; Prefill retains current device writes. Neither + path introduces a Host read or a new stream/event policy. """ from vllm_ascend.hisparse.c8 import c8_storage_view, compact_c8_indices, map_c8_physical_indices - query_ends, key_length = self._c8_length_workspace() + query_ends, key_length = self._c8_length_workspace() if lengths is None else lengths scale = self.scale view, block, stride_rows, capacity = c8_storage_view(cache) rows = query.shape[0] @@ -371,8 +373,9 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a # Read the resolver's current counts on this consumer stream before # later groups can reuse their scratch. Keep the empty-row output guard. has_keys = valid_counts > 0 if valid_counts is not None else (mapped >= 0).any(dim=-1) - query_ends.fill_(rows) - key_length.fill_(view.shape[0]) + if lengths is None: + query_ends.fill_(rows) + key_length.fill_(view.shape[0]) result = torch.ops._C_ascend.npu_kv_quant_sparse_flash_attention_vllm( query=packed_query, key=view, @@ -398,7 +401,15 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a return torch.where(has_keys[:, None, None], result, 0) def _attention_physical( - self, query, query_rope, packed_cache, physical_topk, *, indices_are_mapped=False, valid_counts=None + self, + query, + query_rope, + packed_cache, + physical_topk, + *, + indices_are_mapped=False, + valid_counts=None, + lengths=None, ): if self.enable_sparse_sfa_c8: return self._c8_physical_attention( @@ -408,6 +419,7 @@ def _attention_physical( physical_topk, indices_are_mapped=indices_are_mapped, valid_counts=valid_counts, + lengths=lengths, ) from vllm_ascend.hisparse.cache_ops import pack_attention_rows @@ -490,8 +502,11 @@ def _execute_sparse_flash_attention_process( return_valid_counts=self.enable_sparse_sfa_c8, ) valid_counts = None + lengths = None if self.enable_sparse_sfa_c8: physical, valid_counts = physical + query_lengths, key_length = cache.runtime.cache_ops.attention_lengths + lengths = query_lengths[num_decode : num_decode + 1], key_length outputs.append( self._attention_physical( ql_nope[:num_decode], @@ -500,6 +515,7 @@ def _execute_sparse_flash_attention_process( physical, indices_are_mapped=cache.runtime.cache_ops.attention_capacity > 0, valid_counts=valid_counts, + lengths=lengths, ) ) if metadata.num_prefills: diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 86f6b4eea6c4..d708d118f9b3 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -114,6 +114,7 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: mapping = runtime.registered_host_pool._hisparse_registration attention_capacity = 0 attention_block_stride = None + attention_lengths = None if runtime.kv_dtype == torch.float8_e4m3fn: from vllm_ascend.hisparse.c8 import c8_storage_view @@ -130,8 +131,20 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: or attention_block_stride != group.leader.cache_ops.attention_block_stride ): raise ValueError("Shared HiSparse C8 indices require matching per-layer page geometry") + # Decode query shapes are bounded by the resolver workspace. These + # values describe immutable geometry, not the current valid rows. + # Prepare them before capture and retain them with this binding; + # a rebind gets its own table and current physical KV span. + attention_lengths = ( + torch.arange(workspace.max_rows + 1, dtype=torch.int32, device=runtime.device), + torch.full((1,), attention_view.shape[0], dtype=torch.int32, device=runtime.device), + ) return HiSparseNPUOps( - mapping, workspace, attention_capacity=attention_capacity, attention_block_stride=attention_block_stride + mapping, + workspace, + attention_capacity=attention_capacity, + attention_block_stride=attention_block_stride, + attention_lengths=attention_lengths, ) def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 79e8c529f52b..6433f93178b3 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -188,11 +188,13 @@ def __init__( *, attention_capacity: int = 0, attention_block_stride: int | None = None, + attention_lengths: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> None: self.host_mapping = host_mapping self.workspace = workspace self.attention_capacity = attention_capacity self.attention_block_stride = attention_block_stride + self.attention_lengths = attention_lengths def hisparse_resolve_residency( self, From 6e4e251c2468b81e88bc94ba03fbcde3ddfc8813 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 22:40:38 +0800 Subject: [PATCH 033/120] perf(hisparse): reuse native C8 empty-query output handling Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_attention.py | 27 ++++++++++--------- vllm_ascend/attention/sfa_hisparse.py | 24 +++-------------- 2 files changed, 18 insertions(+), 33 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index d5f4f0585c4c..c28fb7997933 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -111,8 +111,9 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) -@pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_counts", "mapped_lengths"]) -def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode): +@pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_lengths"]) +@pytest.mark.parametrize("empty_batch", [False, True]) +def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode, empty_batch): from vllm_ascend.device.device_config import is_950 if not is_950(): @@ -154,11 +155,6 @@ def selection(): return output selected = selection().to(device) - valid_counts = ( - (indices >= 0).sum(dim=-1, dtype=torch.int32).to(device) - if index_mode in ("mapped_counts", "mapped_lengths") - else None - ) lengths = ( ( torch.arange(rows + 1, dtype=torch.int32, device=device)[rows:], @@ -175,9 +171,7 @@ def selection(): impl._c8_length_workspace = lambda: (query_end, key_length) def forward(): - return impl._attention_physical( - q, qr, cache, selected, indices_are_mapped=indices_are_mapped, valid_counts=valid_counts, lengths=lengths - ) + return impl._attention_physical(q, qr, cache, selected, indices_are_mapped=indices_are_mapped, lengths=lengths) graph = None try: @@ -188,14 +182,19 @@ def forward(): graph = torch.npu.NPUGraph() with torch.npu.graph(graph): output = forward() - indices[0, :4] = torch.tensor([9, 2, 9, 0], dtype=torch.int32) + if empty_batch: + indices.fill_(-1) + else: + indices[0, :4] = torch.tensor([9, 2, 9, 0], dtype=torch.int32) selected.copy_(selection()) - if valid_counts is not None: - valid_counts.copy_((indices >= 0).sum(dim=-1, dtype=torch.int32)) q.neg_() graph.replay() else: output = forward() + if empty_batch: + indices.fill_(-1) + selected.copy_(selection()) + output = forward() torch.npu.synchronize() scale = scales.contiguous().view(torch.float32).cpu().reshape(-1, 4).repeat_interleave(128, dim=-1) latent = nope.float().cpu().reshape(-1, 512) * scale @@ -210,6 +209,8 @@ def forward(): expected = torch.einsum("thk,tkd->thd", probability, keys) torch.testing.assert_close(output.float().cpu(), expected, atol=0.01, rtol=0.05) assert torch.count_nonzero(output[-1].cpu()) == 0 + if empty_batch: + assert torch.count_nonzero(output.cpu()) == 0 finally: torch.npu.synchronize() if graph is not None: diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 9ec3fa7ee449..3c991544ab5b 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -328,9 +328,7 @@ def _store_parallel_kv(self, k_pe, k_nope, knope_scale, kv_cache, slot_mapping_s # which would duplicate the write and can overwrite unrelated slab rows. return k_pe, k_nope - def _c8_physical_attention( - self, query, query_rope, cache, indices, *, indices_are_mapped=False, valid_counts=None, lengths=None - ): + def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_are_mapped=False, lengths=None): """Adapt native packed rows to the existing quantized SFA operation. Keep compaction stable, including duplicates. Decode lengths describe @@ -355,13 +353,6 @@ def _c8_physical_attention( raise ValueError("HiSparse C8 inputs must share one NPU") if any(tensor.shape != (1,) or tensor.dtype != torch.int32 for tensor in (query_ends, key_length)): raise ValueError("HiSparse C8 lengths must be int32 scalar tensors") - if valid_counts is not None and ( - not indices_are_mapped - or valid_counts.shape != (rows,) - or valid_counts.dtype != torch.int32 - or valid_counts.device != query.device - ): - raise ValueError("HiSparse C8 valid counts must describe the native mapped query rows") if not math.isfinite(scale) or scale <= 0: raise ValueError("HiSparse C8 attention scale must be finite and positive") mapped = ( @@ -370,9 +361,6 @@ def _c8_physical_attention( else compact_c8_indices(map_c8_physical_indices(indices, block, stride_rows, capacity)) ) packed_query = torch.cat((query, query_rope), dim=-1) - # Read the resolver's current counts on this consumer stream before - # later groups can reuse their scratch. Keep the empty-row output guard. - has_keys = valid_counts > 0 if valid_counts is not None else (mapped >= 0).any(dim=-1) if lengths is None: query_ends.fill_(rows) key_length.fill_(view.shape[0]) @@ -398,7 +386,9 @@ def _c8_physical_attention( ) if isinstance(result, tuple): result = result[0] - return torch.where(has_keys[:, None, None], result, 0) + # Native C8 clears empty query/head ranges, including graph padding. + # Avoid rescanning indices and allocating a second masked output. + return result def _attention_physical( self, @@ -408,7 +398,6 @@ def _attention_physical( physical_topk, *, indices_are_mapped=False, - valid_counts=None, lengths=None, ): if self.enable_sparse_sfa_c8: @@ -418,7 +407,6 @@ def _attention_physical( packed_cache, physical_topk, indices_are_mapped=indices_are_mapped, - valid_counts=valid_counts, lengths=lengths, ) from vllm_ascend.hisparse.cache_ops import pack_attention_rows @@ -499,12 +487,9 @@ def _execute_sparse_flash_attention_process( logical_topk_indices=logical[:num_decode], block_size=metadata.block_size, num_valid_rows=metadata.query_start_loc[-1:], - return_valid_counts=self.enable_sparse_sfa_c8, ) - valid_counts = None lengths = None if self.enable_sparse_sfa_c8: - physical, valid_counts = physical query_lengths, key_length = cache.runtime.cache_ops.attention_lengths lengths = query_lengths[num_decode : num_decode + 1], key_length outputs.append( @@ -514,7 +499,6 @@ def _execute_sparse_flash_attention_process( self.index_group.physical_kv_cache(self.index_group_index), physical, indices_are_mapped=cache.runtime.cache_ops.attention_capacity > 0, - valid_counts=valid_counts, lengths=lengths, ) ) From 52647441ac0fe71c3a1e3bca268bcc4be69e763e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 5 Nov 2026 23:19:43 +0800 Subject: [PATCH 034/120] docs(hisparse): record native C8 output qualification and remaining regression Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 18 ++++++++++++++---- release_notes/fragments/20926.yaml | 4 +++- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 5199e86da54a..473a4fabc2ee 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -80,6 +80,10 @@ supported projection schemes retain the native preprocessing path. C8 Decode consumes the resolver's mapped, stably compacted indices directly. The refill plan retains separate virtual page-row IDs. Prefill keeps its existing layout bridge, including stable ordering and duplicate selections. +Decode's padded query-shape lengths and physical KV extent are prepared before +capture and owned by the current binding. Per-round request mappings and valid +rows remain current. Native C8 Attention already clears empty query outputs; +the adapter returns that output without another index scan and output mask. Keep completion-before-unregister and graph allocation ownership intact. Overwriting KV still invalidates any old Hot copy. Equal batch shapes do not @@ -97,10 +101,16 @@ In the current GLM-5.2 C4 comparison (32768 input and 8192 output tokens, TP8/DP1, MTP3, Model Runner V2, prefix caching off and explicit `weight_nz_mode=0`), both arms completed two rounds with 65536 actual output token IDs per arm. ON's settled reported miss/H2D deltas were zero. Mean TPOT was -12.21 ms OFF and 14.20 ms ON: an observed increase of 16.29%. The second round -had identical output tokens and MTP acceptance, with an 18.11% TPOT increase; -the first round had different outputs/acceptance. These measurements do not -establish statistical significance or acceptable performance for other workloads. +11.96 ms OFF and 13.89 ms ON: an observed increase of 16.09%. Outputs and MTP +work differed between these arms, so this mean is not an isolated measure of +backend overhead. A preceding paired round with identical outputs and MTP work +still showed an 18.11% increase. + +Against the preceding direct-index ON cohort, removing redundant C8 output +masking and Decode length writes reduced observed mean ON TPOT from 14.20 to +13.89 ms (2.19%). Complete input/output-sequence multisets and aggregate MTP +work matched across those two ON cohorts. Two rounds do not establish +statistical significance or acceptable performance for other workloads. Dataset accuracy, the full-model miss path and a TPOT regression within 3% remain unqualified. Profiling requests are separate from formal performance diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index 591d1fa1def2..1e2f625a5a7e 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -7,5 +7,7 @@ entries: operators and Core lifecycle interfaces. The resolver requires its compiled extension and does not fall back to Triton. Native C8 reuses supported SFA Prolog preprocessing and consumes the resolver's physical - Decode indices while preserving the separate refill plan. + Decode indices while preserving the separate refill plan. Decode lengths + are prepared with the cache binding and native empty-query clearing is + reused without redundant output masking. hardware: [A5] From 1861b61ad81bcd9736f15136084dde48018638a9 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 00:22:08 +0800 Subject: [PATCH 035/120] perf(hisparse): bound compact refill programs for current miss counts Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 13 ++-- release_notes/fragments/20926.yaml | 4 +- .../one_card/test_hisparse_ops.py | 64 +++++++++++++++++ vllm_ascend/hisparse/cache_ops.py | 4 +- vllm_ascend/ops/triton/hisparse.py | 68 +++++++++++++------ 5 files changed, 123 insertions(+), 30 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 473a4fabc2ee..7e4d9d086170 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -101,16 +101,17 @@ In the current GLM-5.2 C4 comparison (32768 input and 8192 output tokens, TP8/DP1, MTP3, Model Runner V2, prefix caching off and explicit `weight_nz_mode=0`), both arms completed two rounds with 65536 actual output token IDs per arm. ON's settled reported miss/H2D deltas were zero. Mean TPOT was -11.96 ms OFF and 13.89 ms ON: an observed increase of 16.09%. Outputs and MTP +12.30 ms OFF and 13.78 ms ON: an observed increase of 12.03%. Outputs and MTP work differed between these arms, so this mean is not an isolated measure of backend overhead. A preceding paired round with identical outputs and MTP work still showed an 18.11% increase. -Against the preceding direct-index ON cohort, removing redundant C8 output -masking and Decode length writes reduced observed mean ON TPOT from 14.20 to -13.89 ms (2.19%). Complete input/output-sequence multisets and aggregate MTP -work matched across those two ON cohorts. Two rounds do not establish -statistical significance or acceptable performance for other workloads. +In separate ON comparisons with identical inputs, outputs and MTP work, native C8 output and +length reuse reduced observed mean TPOT by 2.19% across two rounds. Bounding +compact-refill programs reduced TPOT by 1.93% in one matching round. These gains +are not additive and do not establish statistical significance or acceptable +performance for other workloads. Compact refill retains current device counts, +page-strided byte copies and the existing completion dependencies. Dataset accuracy, the full-model miss path and a TPOT regression within 3% remain unqualified. Profiling requests are separate from formal performance diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index 1e2f625a5a7e..2edffa770acd 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -9,5 +9,7 @@ entries: supported SFA Prolog preprocessing and consumes the resolver's physical Decode indices while preserving the separate refill plan. Decode lengths are prepared with the cache binding and native empty-query clearing is - reused without redundant output masking. + reused without redundant output masking. Compact refill bounds the number + of per-query programs while retaining current counts and scalar byte + copies for both native C8 and BF16 payloads. hardware: [A5] diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 4b564bd4bc4e..c65dedf90ecc 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -606,6 +606,70 @@ def forward(): mapping.close() +@pytest.mark.parametrize("rows,topk,row_bytes", [(1, 17, 656), (4, 2048, 656), (16, 2048, 1152)]) +@pytest.mark.parametrize("capture", [False, True]) +def test_compact_refill_tiles_current_counts_and_preserves_page_gaps(rows, topk, row_bytes, capture): + block = 128 + blocks = (rows * topk + block - 1) // block + 1 + host = allocate_mapped_host((blocks * block, row_bytes), torch.uint8) + mapping = host._hisparse_registration + host.copy_((torch.arange(host.shape[0], dtype=torch.int32) % 251).to(torch.uint8)[:, None].expand_as(host)) + backing = torch.full((blocks, 2, block, row_bytes), 203, dtype=torch.uint8, device="npu") + hot = backing[:, 1] + ops = HiSparseNPUOps(mapping) + original_source = torch.arange(rows * topk, dtype=torch.int32).view(rows, topk) + original_destination = original_source + block + original_source[:, 4] = -1 + original_source[:, 6] = host.shape[0] + 7 + original_destination[:, 2] = -1 + original_destination[:, 5] = blocks * block + 3 + source, destination = original_source.to("npu"), original_destination.to("npu") + counts = torch.zeros(rows, dtype=torch.int32, device="npu") + + def forward(): + ops.hisparse_gather_compact(host, hot, source, destination, counts) + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + forward() + for count in (0, 1, 7, 8, 9, 15, 16, 17, topk, 0): + backing.fill_(203) + current = torch.full((rows,), count, dtype=torch.int32) + if rows > 1: + current[1] = -3 + if rows > 2 and count == topk: + current[2] = topk + 3 + counts.copy_(current) + source.copy_(original_source if count else torch.full_like(original_source, -1)) + destination.copy_(original_destination if count else torch.full_like(original_destination, -1)) + if graph is None: + forward() + else: + graph.replay() + torch.npu.synchronize() + expected = torch.full(backing.shape, 203, dtype=torch.uint8) + for row in range(rows): + length = min(topk, max(0, int(current[row]))) + src = original_source[row, :length].long() + dst = original_destination[row, :length].long() + valid = (src >= 0) & (dst >= 0) & (dst < blocks * block) + payload = torch.zeros((length, row_bytes), dtype=torch.uint8) + from_host = valid & (src < host.shape[0]) + payload[from_host] = host[src[from_host]] + expected[dst[valid] // block, 1, dst[valid] % block] = payload[valid] + torch.testing.assert_close(backing.cpu(), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + mapping.close() + + def test_plan_gather_and_written_slot_invalidation(): host = allocate_mapped_host((32, 16), torch.float32) mapping = host._hisparse_registration diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 6433f93178b3..aa613dc72290 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -366,7 +366,8 @@ def _gather( raise ValueError("HiSparse attention block stride must cover a physical block") if rows == 0: return - gather_rows_kernel[(rows, topk)]( + parallel_columns = min(topk, 128) if counts is not None and attention is None else 0 + gather_rows_kernel[(rows, parallel_columns or topk)]( self.host_mapping.address_of(host_cache), hot_cache.view(torch.uint8), source, @@ -386,6 +387,7 @@ def _gather( HAS_ATTENTION=attention is not None, HAS_STATE=state is not None, BLOCK=min(1024, triton.next_power_of_2(row_bytes)), + PARALLEL_COLUMNS=parallel_columns, multibuffer=False, ) diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 09d7f219e5cc..88901a903131 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -29,33 +29,57 @@ def gather_rows_kernel( HAS_ATTENTION: tl.constexpr, HAS_STATE: tl.constexpr, BLOCK: tl.constexpr, + PARALLEL_COLUMNS: tl.constexpr = 0, ): row = tl.program_id(0) - col = tl.program_id(1) offset = tl.arange(0, BLOCK) - index = row * TOP_K + col - dst = tl.load(destination + index) padding = tl.load(state + row) < 0 if HAS_STATE else False - if HAS_ATTENTION: - value = tl.where((dst >= 0) & ~padding, dst // HOT_BLOCK * ATTENTION_BLOCK + dst % HOT_BLOCK, -1) - tl.store(attention + index, value) - if COMPACT: - copy = col < tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) + if PARALLEL_COLUMNS > 0: + tl.static_assert(COMPACT and not HAS_ATTENTION) + count = tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) + # Empty compact plans do not inspect stale source/destination entries. + # Keep enough independent programs for small misses and retain the + # original contiguous byte-copy shape instead of a wide 2-D payload. + if ~padding: + for col in range(tl.program_id(1), count, PARALLEL_COLUMNS): + index = row * TOP_K + col + src = tl.load(source + index) + dst = tl.load(destination + index) + valid = (src >= 0) & (dst >= 0) & (dst < HOT_ROWS) + if valid: + host = host_address.to(tl.pointer_type(tl.uint8)) + destination_offset = dst.to(tl.int64) // HOT_BLOCK * HOT_STRIDE_BYTES + dst % HOT_BLOCK * ROW_BYTES + for start in range(0, ROW_BYTES, BLOCK): + payload = tl.load( + host + src.to(tl.int64) * ROW_BYTES + start + offset, + mask=(src < HOST_ROWS) & (start + offset < ROW_BYTES), + other=0, + ) + tl.store(hot + destination_offset + start + offset, payload, mask=start + offset < ROW_BYTES) else: - copy = tl.load(mask_ptr + index) != 0 - if copy & ~padding: - src = tl.load(source + index) - valid = (src >= 0) & (dst >= 0) & (dst < HOT_ROWS) - if valid: - host = host_address.to(tl.pointer_type(tl.uint8)) - destination_offset = dst.to(tl.int64) // HOT_BLOCK * HOT_STRIDE_BYTES + dst % HOT_BLOCK * ROW_BYTES - for start in range(0, ROW_BYTES, BLOCK): - payload = tl.load( - host + src.to(tl.int64) * ROW_BYTES + start + offset, - mask=(src < HOST_ROWS) & (start + offset < ROW_BYTES), - other=0, - ) - tl.store(hot + destination_offset + start + offset, payload, mask=start + offset < ROW_BYTES) + col = tl.program_id(1) + index = row * TOP_K + col + dst = tl.load(destination + index) + if HAS_ATTENTION: + value = tl.where((dst >= 0) & ~padding, dst // HOT_BLOCK * ATTENTION_BLOCK + dst % HOT_BLOCK, -1) + tl.store(attention + index, value) + if COMPACT: + copy = col < tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) + else: + copy = tl.load(mask_ptr + index) != 0 + if copy & ~padding: + src = tl.load(source + index) + valid = (src >= 0) & (dst >= 0) & (dst < HOT_ROWS) + if valid: + host = host_address.to(tl.pointer_type(tl.uint8)) + destination_offset = dst.to(tl.int64) // HOT_BLOCK * HOT_STRIDE_BYTES + dst % HOT_BLOCK * ROW_BYTES + for start in range(0, ROW_BYTES, BLOCK): + payload = tl.load( + host + src.to(tl.int64) * ROW_BYTES + start + offset, + mask=(src < HOST_ROWS) & (start + offset < ROW_BYTES), + other=0, + ) + tl.store(hot + destination_offset + start + offset, payload, mask=start + offset < ROW_BYTES) @triton.jit From 03cc8b37be3dc363497261836d55b0c9e8eb9e5d Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 06:14:07 +0800 Subject: [PATCH 036/120] refactor(hisparse): omit unused C8 softmax LSE outputs The HiSparse C8 adapter consumes only the Attention result. Keep the existing native layout and empty-query behavior while disabling the unused LSE outputs. Actual eager/Graph parity checks cover changing selections, strided pages and large physical extents; full-model results remain separate. Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 3c991544ab5b..97cd321031e4 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -382,7 +382,7 @@ def _c8_physical_attention(self, query, query_rope, cache, indices, *, indices_a key_quant_mode=2, value_quant_mode=2, rope_head_dim=64, - return_softmax_lse=True, + return_softmax_lse=False, ) if isinstance(result, tuple): result = result[0] From 0f158d99a92caeb2db2ba1b2b780e2756e94182b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 06:43:20 +0800 Subject: [PATCH 037/120] refactor(hisparse): reuse Core device request mapping Private candidate: replace the arange/searchsorted/copy/clamp chain with the existing merged Core metadata operation. Preserve Ascend padded input extents and independent MTP graph-owner buffers. No Attention or resolver kernel changes; actual NPU and full-model qualification remain required before public promotion. Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 97cd321031e4..2491604de4bb 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -8,6 +8,7 @@ import torch import torch_npu from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.worker.workspace import current_workspace_manager from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder @@ -59,14 +60,18 @@ def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): def _build(self, common_attn_metadata, draft_index=None, pcp_context=None): metadata = super()._build(common_attn_metadata, draft_index, pcp_context) metadata.num_reqs = common_attn_metadata.num_reqs - rows = torch.arange(metadata.num_input_tokens, device=self.device, dtype=torch.int32) if draft_index not in self._request_row_buffers: self._request_row_buffers[draft_index] = torch.empty( self._row_capacity, device=self.device, dtype=torch.int32 ) - metadata.req_id_per_token = self._request_row_buffers[draft_index][: metadata.num_input_tokens] - metadata.req_id_per_token.copy_(torch.searchsorted(metadata.query_start_loc[1:].contiguous(), rows, right=True)) - metadata.req_id_per_token.clamp_(max=max(metadata.num_reqs - 1, 0)) + # Use Core's device-boundary mapping with Ascend's padded input extent. + # Keep separate backing storage for each captured MTP graph owner. + metadata.req_id_per_token = compute_token_to_req_indices( + metadata.query_start_loc, + self._request_row_buffers[draft_index], + int(common_attn_metadata.query_start_loc_cpu[-1]), + metadata.num_input_tokens, + ) if metadata.num_prefills: seq_lens_cpu = metadata.seq_lens_cpu if seq_lens_cpu is None: From 7fa1df003e8b6678f1c5b5ab71c9bf8660d99764 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 07:12:09 +0800 Subject: [PATCH 038/120] perf(hisparse): reuse indexer store with page strides Private unqualified candidate: extend the existing key/scale writer with independent cache-page strides while retaining the dense defaults. The HiSparse FP8 adapter supplies byte-preserving keys and FP32 scales at the existing write point, removing materialized address tensors and separate scatters. Native quantization, Attention and Core are unchanged. Reuse the existing real-NPU page-gap/bit-pattern tests and dense INT8/FP16 tests before any model or public qualification. Signed-off-by: zhenwenqi2024 --- tests/ut/attention/test_sfa_dcp.py | 7 ++++- vllm_ascend/attention/indexer.py | 1 - vllm_ascend/hisparse/binding.py | 32 +++++++++------------ vllm_ascend/ops/triton/sfa_indexer_store.py | 32 +++++++++++++++++---- 4 files changed, 47 insertions(+), 25 deletions(-) diff --git a/tests/ut/attention/test_sfa_dcp.py b/tests/ut/attention/test_sfa_dcp.py index ba65acafc2f4..7aa9a6e09693 100644 --- a/tests/ut/attention/test_sfa_dcp.py +++ b/tests/ut/attention/test_sfa_dcp.py @@ -557,7 +557,12 @@ def legacy(keys, stats, slots, **kwargs): }, ), ("dcp/sfa_dcp_query.py", "_unpack_query_fragments", {"NOPE_DIM", "ROPE_DIM"}, {"tokens", "programs", "heads"}), - ("sfa_indexer_store.py", "_store_indexer_key_scale", {"KEY_WIDTH"}, {"capacity"}), + ( + "sfa_indexer_store.py", + "_store_indexer_key_scale", + {"KEY_WIDTH", "BLOCK_SIZE"}, + {"capacity", "key_block_stride", "scale_block_stride"}, + ), ] diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 33493544b70e..3ea0f6cf9159 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -546,7 +546,6 @@ def write_cache(self, k_li, k_li_scale, slot_mapping) -> None: k_li, k_li_scale, slot_mapping.reshape(-1), - DeviceOperator.scatter_cache, ) diff --git a/vllm_ascend/hisparse/binding.py b/vllm_ascend/hisparse/binding.py index 1f98b844e81d..0b17deb456dc 100644 --- a/vllm_ascend/hisparse/binding.py +++ b/vllm_ascend/hisparse/binding.py @@ -119,34 +119,30 @@ def _quantized_indexer_views(cache: torch.Tensor, spec: AscendSFAIndexerCacheSpe return key, scale -def write_hisparse_indexer_c8(key, scale, values, scales, slots, scatter) -> None: +def write_hisparse_indexer_c8(key, scale, values, scales, slots) -> None: """Write byte-preserving K and FP32 scales using this layer's page stride. Derive aliases from the current cache tuple, so profiling-cache replacement and teardown cannot leave a second set of persistent cache references. """ - pages, block, heads, width = key.shape - stride = key.stride(0) + from vllm_ascend.ops.triton.sfa_indexer_store import store_indexer_key_scale + + block, width = key.shape[1], key.shape[-1] count = slots.numel() if values.numel() != count * width or scales.numel() != count: raise ValueError("HiSparse FP8 indexer writes require matching token/slot counts.") if values.dtype != key.dtype or scales.dtype != torch.float32: raise ValueError("HiSparse FP8 indexer writes require E4M3FN data and FP32 scales.") - key_rows = ((pages - 1) * stride + block * width) // width - scale_rows = (pages - 1) * scale.stride(0) + block - raw_key = torch.as_strided(key, (0,), (1,), storage_offset=0).view(torch.int8) - key_write = torch.as_strided(raw_key, (key_rows, width), (width, 1), storage_offset=key.storage_offset()) - scale_write = torch.as_strided(scale, (scale_rows, 1), (1, 1)) - # Ascend V2 supplies int32 slots. Page-strided scale addresses can exceed - # INT32_MAX even when logical slots fit; widen before multiplying strides. - logical = slots.reshape(-1).to(torch.int64) - page = logical // block - row = logical % block - # Preserve the native -1 padding sentinel; never turn it into a valid row. - key_slots = torch.where(logical >= 0, page * (stride // width) + row, -1) - scale_slots = torch.where(logical >= 0, page * scale.stride(0) + row, -1) - scatter(key_write, key_slots.view(-1, 1), values.contiguous().view(torch.int8).reshape(-1, width)) - scatter(scale_write, scale_slots.view(-1, 1), scales.reshape(-1, 1)) + store_indexer_key_scale( + values.contiguous().view(torch.int8).reshape(-1, width), + scales.reshape(-1), + slots.reshape(-1), + key.view(torch.int8), + scale, + block_size=block, + key_block_stride=key.stride(0), + scale_block_stride=scale.stride(0), + ) def _bind_indexer_views( diff --git a/vllm_ascend/ops/triton/sfa_indexer_store.py b/vllm_ascend/ops/triton/sfa_indexer_store.py index 9c2334623983..8c6f8c301b6f 100644 --- a/vllm_ascend/ops/triton/sfa_indexer_store.py +++ b/vllm_ascend/ops/triton/sfa_indexer_store.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Store native-quantized indexer keys and FP16 scales in one kernel. +"""Store native-quantized indexer keys and scales in one kernel. Keep native quantization and its INT8 tie behavior. Only final scale conversion and the two scatters are fused. Valid cache slots must be unique, as provided @@ -10,7 +10,7 @@ from vllm.triton_utils import tl, triton -@triton.jit(do_not_specialize=["capacity"]) +@triton.jit(do_not_specialize=["capacity", "key_block_stride", "scale_block_stride"]) def _store_indexer_key_scale( quant_ptr, scales_ptr, @@ -18,7 +18,10 @@ def _store_indexer_key_scale( cache_ptr, scale_cache_ptr, capacity, + key_block_stride, + scale_block_stride, KEY_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, ): row = tl.program_id(0) column = tl.arange(0, KEY_WIDTH) @@ -26,8 +29,10 @@ def _store_indexer_key_scale( scale = tl.load(scales_ptr + row) slot = tl.load(slots_ptr + row).to(tl.int64) valid = (slot >= 0) & (slot < capacity) - tl.store(cache_ptr + slot * KEY_WIDTH + column, key, valid) - tl.store(scale_cache_ptr + slot, scale, valid) + page = slot // BLOCK_SIZE + offset = slot % BLOCK_SIZE + tl.store(cache_ptr + page * key_block_stride + offset * KEY_WIDTH + column, key, valid) + tl.store(scale_cache_ptr + page * scale_block_stride + offset, scale, valid) def store_indexer_key_scale( @@ -36,9 +41,26 @@ def store_indexer_key_scale( slots: torch.Tensor, cache: torch.Tensor, scale_cache: torch.Tensor, + *, + block_size: int = 1, + key_block_stride: int | None = None, + scale_block_stride: int | None = None, ) -> None: """Publish at the existing cache-write point, after the connector wait.""" rows = slots.numel() key_width = quant.shape[-1] + key_block_stride = key_width if key_block_stride is None else key_block_stride + scale_block_stride = 1 if scale_block_stride is None else scale_block_stride # Rows only set the launch grid; 12 was a test setting, not a kernel limit. - _store_indexer_key_scale[(rows,)](quant, scales, slots, cache, scale_cache, cache.numel() // key_width, key_width) + _store_indexer_key_scale[(rows,)]( + quant, + scales, + slots, + cache, + scale_cache, + cache.numel() // key_width, + key_block_stride, + scale_block_stride, + key_width, + block_size, + ) From 5a453209677af6f5f030829e13ee14286bcc1f1c Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 07:14:10 +0800 Subject: [PATCH 039/120] fix(hisparse): preserve indexer scale payload bits Keep FP32 scale storage opaque through int32 aliases at the shared writer boundary, including signed zero and NaN payloads. This changes no quantization arithmetic and needs the existing real-NPU page-stride byte oracle before qualification. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/binding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/hisparse/binding.py b/vllm_ascend/hisparse/binding.py index 0b17deb456dc..3abbe55f33ea 100644 --- a/vllm_ascend/hisparse/binding.py +++ b/vllm_ascend/hisparse/binding.py @@ -135,10 +135,10 @@ def write_hisparse_indexer_c8(key, scale, values, scales, slots) -> None: raise ValueError("HiSparse FP8 indexer writes require E4M3FN data and FP32 scales.") store_indexer_key_scale( values.contiguous().view(torch.int8).reshape(-1, width), - scales.reshape(-1), + scales.view(torch.int32).reshape(-1), slots.reshape(-1), key.view(torch.int8), - scale, + scale.view(torch.int32), block_size=block, key_block_stride=key.stride(0), scale_block_stride=scale.stride(0), From 1efad3415f76198dddeb4c0fb45ed355ac311bd3 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 07:25:53 +0800 Subject: [PATCH 040/120] fix(indexer): derive contiguous page-stride defaults A caller that supplies block_size while omitting strides still has contiguous key/scale pages. Derive those strides from the block size; existing dense block-size-one and explicit HiSparse strides are unchanged. Signed-off-by: zhenwenqi2024 --- vllm_ascend/ops/triton/sfa_indexer_store.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/ops/triton/sfa_indexer_store.py b/vllm_ascend/ops/triton/sfa_indexer_store.py index 8c6f8c301b6f..1f19abc4753c 100644 --- a/vllm_ascend/ops/triton/sfa_indexer_store.py +++ b/vllm_ascend/ops/triton/sfa_indexer_store.py @@ -49,8 +49,8 @@ def store_indexer_key_scale( """Publish at the existing cache-write point, after the connector wait.""" rows = slots.numel() key_width = quant.shape[-1] - key_block_stride = key_width if key_block_stride is None else key_block_stride - scale_block_stride = 1 if scale_block_stride is None else scale_block_stride + key_block_stride = block_size * key_width if key_block_stride is None else key_block_stride + scale_block_stride = block_size if scale_block_stride is None else scale_block_stride # Rows only set the launch grid; 12 was a test setting, not a kernel limit. _store_indexer_key_scale[(rows,)]( quant, From 6ef170d1880966c73bbfea2c3ca1106372739f3e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 07:36:14 +0800 Subject: [PATCH 041/120] test(hisparse): mock the shared indexer store boundary on CPU Signed-off-by: zhenwenqi2024 --- .../test_hisparse_foundation_cpu.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index bc0f5a784844..ee404e2e0c17 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -18,7 +18,6 @@ ) from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec -from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.hisparse import cache_ops as hisparse_ops from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.binding import ( @@ -32,6 +31,7 @@ map_c8_physical_indices, pack_c8_components, ) +from vllm_ascend.ops.triton import sfa_indexer_store @pytest.mark.parametrize("mxfp8", [False, True]) @@ -409,12 +409,18 @@ def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(monkeypatch values = values_raw.view(torch.float8_e4m3fn) scales = torch.arange(6, dtype=torch.float32).reshape(6, 1) + 0.5 - def scatter(destination, indexes, source): - indexes = indexes.flatten() - valid = indexes >= 0 - destination[indexes[valid].long()] = source[valid] - - monkeypatch.setattr(DeviceOperator, "scatter_cache", scatter) + def store(quant, scales, slots, cache, scale_cache, *, block_size, key_block_stride, scale_block_stride): + # Emulate only the device-store boundary; production views and addresses + # are checked against the independent byte oracle below. + assert block_size == 128 and key_block_stride == stride and scale_block_stride == stride // 4 + assert cache.dtype == torch.int8 and scale_cache.dtype == torch.int32 + for i, slot in enumerate(slots.tolist()): + if 0 <= slot < pages * block_size: + page, row = divmod(slot, block_size) + cache[page, row, 0].copy_(quant[i]) + scale_cache[page, row, 0, 0] = scales[i] + + monkeypatch.setattr(sfa_indexer_store, "store_indexer_key_scale", store) writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) torch.nn.Module.__init__(writer) writer.k_cache = NS(kv_cache=caches["indexer"]) From fcdeef2318c00511dd9de5b678af7fa1c1ef5af1 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 08:00:51 +0800 Subject: [PATCH 042/120] perf(hisparse): reuse PTA queue for small host mirror submissions Keep Core descriptor retirement and completion events. Reuse the qualified queued-copy implementation without the mapped-copy kernel, and extend the existing payload test to changing Graph producers and immutable scalar descriptors. Signed-off-by: zhenwenqi2024 --- csrc/torch_binding.cpp | 47 +++++++++++- .../one_card/test_hisparse_ops.py | 76 ++++++++++++++----- vllm_ascend/hisparse/backend.py | 2 +- 3 files changed, 101 insertions(+), 24 deletions(-) diff --git a/csrc/torch_binding.cpp b/csrc/torch_binding.cpp index 045f417d3873..57d71b4d64da 100644 --- a/csrc/torch_binding.cpp +++ b/csrc/torch_binding.cpp @@ -214,7 +214,8 @@ void enqueue_device_print(std::unique_ptr payload, void swap_blocks_batch(const torch::Tensor& src_ptrs, const torch::Tensor& dst_ptrs, const torch::Tensor& sizes, - int64_t direction) { + int64_t direction, + bool prefer_small_async) { TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); @@ -233,8 +234,6 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, const int64_t* dst_data = dst_ptrs.data_ptr(); const int64_t* size_data = sizes.data_ptr(); - aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); - aclrtMemcpyKind memcpy_kind; switch (direction) { case 0: @@ -252,10 +251,50 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, " (expected 0=H2D, 1=D2H, 2=D2D)"); } + constexpr int64_t kSmallAsyncCopyLimit = 16; + if (prefer_small_async && n <= kSmallAsyncCopyLimit && + memcpy_kind != ACL_MEMCPY_DEVICE_TO_DEVICE) { + // Copy scalar descriptors into PTA's queued command, before the + // caller's completion event. Core retains payload owners until that + // event completes; no tensor/Python owners move to the queue thread. + struct Copy { + void* dst; + const void* src; + size_t size; + }; + std::vector copies; + copies.reserve(n); + for (int64_t i = 0; i < n; ++i) { + copies.push_back({reinterpret_cast(dst_data[i]), + reinterpret_cast(src_data[i]), + static_cast(size_data[i])}); + } + const auto task_stream = c10_npu::getCurrentNPUStream(); + auto submit = [copies = std::move(copies), task_stream, memcpy_kind]() -> int { + // Only the ordered PTA consumer may bypass Host queue draining. + const aclrtStream stream = task_stream.stream(false); + for (size_t i = 0; i < copies.size(); ++i) { + const auto& copy = copies[i]; + const aclError ret = aclrtMemcpyAsync( + copy.dst, copy.size, copy.src, copy.size, memcpy_kind, stream); + // Throw on a partial submission; never retry these copies. + TORCH_CHECK(ret == ACL_SUCCESS, + "Queued async memcpy failed at index ", i, + " with error code ", ret); + } + return 0; + }; + at_npu::native::OpCommand::RunOpApiV2("HiSparseSmallHostMirror", submit); + return; + } + + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + // ========================================================================= // path 1: aclrtMemcpyBatchAsync (CANN 8.5+) // ========================================================================= #if defined(CANN_MEMCPY_BATCH_ASYNC) + // Full-layer mirrors and callers that did not opt in keep the batch API. if (memcpy_kind != ACL_MEMCPY_DEVICE_TO_DEVICE) { static_assert(sizeof(void*) == sizeof(int64_t), "void* and int64_t must be the same size"); @@ -2932,7 +2971,7 @@ TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) // swap_blocks_batch takes CPU tensors (int64 pointer/size arrays), not NPU // tensors, so dispatch must be registered on the CPU backend. The function // internally submits async memcpy on the current NPU stream. - ops.def("swap_blocks_batch(Tensor x, Tensor y, Tensor z, int direction) -> ()"); + ops.def("swap_blocks_batch(Tensor x, Tensor y, Tensor z, int direction, bool prefer_small_async=False) -> ()"); ops.impl("swap_blocks_batch", torch::kCPU, &vllm_ascend::swap_blocks_batch); #ifdef VLLM_ASCEND_ENABLE_SPARSE_KV_OFFLOAD diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index c65dedf90ecc..d0839dc7d2ff 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -29,39 +29,77 @@ def _require_mapped_host_device(): AscendHiSparseBackend().check_available() -def test_private_host_pool_retains_pinned_storage_for_dma_writeback(): +@pytest.mark.parametrize("descriptor_count", [0, 1, 2, 4, 16, 17, 312]) +@pytest.mark.parametrize("graph_producer", [False, True]) +def test_private_host_pool_retains_pinned_storage_for_dma_writeback(descriptor_count, graph_producer): from vllm.v1.hisparse import runtime backend = AscendHiSparseBackend() backend.check_available() - payload, registered = runtime.allocate_pinned_host_pool(32768) + capacity = max(32768, (descriptor_count + 1) * 2048) + payload, registered = runtime.allocate_pinned_host_pool(capacity) + graph = None try: - assert payload.numel() == 32768 + assert payload.numel() == capacity assert payload.is_pinned() and registered.is_pinned() assert registered.data_ptr() % 4096 == 0 mapping = registered._hisparse_registration assert mapping.address_of(payload[656:1312]) == mapping.device_ptr + 656 assert payload.untyped_storage().data_ptr() == registered.untyped_storage().data_ptr() - source = torch.arange(32768, dtype=torch.int32, device="npu").to(torch.uint8) + cpu_source = torch.arange(capacity, dtype=torch.int32).to(torch.uint8) + source = cpu_source.to("npu") + staging = source.clone() + if graph_producer: + source.copy_(staging) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + source.copy_(staging) stream = torch.npu.Stream() + source_offsets = [i * 2048 for i in range(descriptor_count)] + destination_offsets = [(descriptor_count - i - 1) * 2048 + 656 for i in range(descriptor_count)] + copy_sizes = [656 * (1 + i % 2) for i in range(descriptor_count)] + src = torch.tensor([source.data_ptr() + i for i in source_offsets], dtype=torch.int64, pin_memory=True) + dst = torch.tensor([payload.data_ptr() + i for i in destination_offsets], dtype=torch.int64, pin_memory=True) + sizes = torch.tensor(copy_sizes, dtype=torch.int64, pin_memory=True) + originals = tuple(t.clone() for t in (src, dst, sizes)) for change in (1, 3, 7): - source.bitwise_xor_(change) - payload.fill_(-1) - src = torch.tensor([source.data_ptr(), source.data_ptr() + 8192], dtype=torch.int64) - dst = torch.tensor([payload.data_ptr() + 656, payload.data_ptr() + 16384], dtype=torch.int64) - sizes = torch.tensor([656, 1312], dtype=torch.int64) - stream.wait_stream(torch.npu.current_stream()) - with stream: - backend.swap_blocks_batch(src, dst, sizes) - complete = torch.npu.Event() - complete.record() - complete.synchronize() - expected = torch.full((32768,), 255, dtype=torch.uint8) - expected[656:1312] = source[:656].cpu() - expected[16384:17696] = source[8192:9504].cpu() - torch.testing.assert_close(payload.view(torch.uint8), expected, rtol=0, atol=0) + cpu_source.bitwise_xor_(change) + staging.copy_(cpu_source, non_blocking=True) + if graph is None: + source.copy_(staging) + else: + graph.replay() + expected = torch.full((capacity,), 255, dtype=torch.uint8) + for source_offset, destination_offset, size in zip(source_offsets, destination_offsets, copy_sizes): + expected[destination_offset : destination_offset + size] = cpu_source[ + source_offset : source_offset + size + ] + for optimized in (False, True): + payload.fill_(-1) + stream.wait_stream(torch.npu.current_stream()) + with stream: + if optimized: + backend.swap_blocks_batch(src, dst, sizes) + if 0 < descriptor_count <= 16: + # The queued command owns scalar descriptors even + # after the caller reuses its CPU metadata bank. + for descriptor in (src, dst, sizes): + descriptor.zero_() + else: + # Existing callers keep the four-argument schema and + # original API selection unless they opt in explicitly. + torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1) + complete = torch.npu.Event() + complete.record() + complete.synchronize() + torch.testing.assert_close(payload.view(torch.uint8), expected, rtol=0, atol=0) + for descriptor, original in zip((src, dst, sizes), originals): + descriptor.copy_(original) finally: torch.npu.synchronize() + if graph is not None: + graph.reset() backend.unpin_tensor(registered) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index d708d118f9b3..367baf21c35e 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -176,4 +176,4 @@ def gather_resident_rows(self, source: torch.Tensor, destination: torch.Tensor, super().gather_resident_rows(source, destination, row_ids) def swap_blocks_batch(self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor) -> None: - torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1) + torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1, prefer_small_async=True) From b2dfb24cfbabeccf7d41ab30066b0cb83b7a1e1a Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 09:03:27 +0800 Subject: [PATCH 043/120] docs(hisparse): record matched zero-miss qualification Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 20 +++++++++++++++---- vllm_ascend/hisparse/backend.py | 2 +- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 7e4d9d086170..3197ecf84c31 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -85,6 +85,13 @@ capture and owned by the current binding. Per-round request mappings and valid rows remain current. Native C8 Attention already clears empty query outputs; the adapter returns that output without another index scan and output mask. +Request-row mapping reuses Core's device implementation, with separate storage +for each captured MTP graph owner and Ascend's padded input extent. Indexer +writes reuse the existing SFA store with independent page strides for keys and +scales, preserving the original payload bits. Small Host-mirror submissions use +PTA's ordered queue with immutable scalar descriptors; full-layer batches keep +the batch-copy API and Core's completion and payload-ownership boundaries. + Keep completion-before-unregister and graph allocation ownership intact. Overwriting KV still invalidates any old Hot copy. Equal batch shapes do not permit reusing stale request mappings or lengths. No upstream vLLM source change @@ -101,10 +108,11 @@ In the current GLM-5.2 C4 comparison (32768 input and 8192 output tokens, TP8/DP1, MTP3, Model Runner V2, prefix caching off and explicit `weight_nz_mode=0`), both arms completed two rounds with 65536 actual output token IDs per arm. ON's settled reported miss/H2D deltas were zero. Mean TPOT was -12.30 ms OFF and 13.78 ms ON: an observed increase of 12.03%. Outputs and MTP -work differed between these arms, so this mean is not an isolated measure of -backend overhead. A preceding paired round with identical outputs and MTP work -still showed an 18.11% increase. +12.313 ms OFF and 13.249 ms ON: an observed increase of 7.61%. All four rounds +had identical complete input/output-sequence multisets and MTP counters +(59.8936% acceptance). Sampled steady Decode throughput decreased by 4.58%. +TPOT includes the mixed Prefill/Decode portion of the requests, so its increase +must not be attributed entirely to steady Decode or to a single operation. In separate ON comparisons with identical inputs, outputs and MTP work, native C8 output and length reuse reduced observed mean TPOT by 2.19% across two rounds. Bounding @@ -113,6 +121,10 @@ are not additive and do not establish statistical significance or acceptable performance for other workloads. Compact refill retains current device counts, page-strided byte copies and the existing completion dependencies. +The request mapping, page-strided Indexer store and small-copy queue passed +focused eager/Graph NPU checks before the current full-model comparison. These +checks do not establish an isolated end-to-end gain for each individual change. + Dataset accuracy, the full-model miss path and a TPOT regression within 3% remain unqualified. Profiling requests are separate from formal performance timing; reported zero-miss counters do not prove the absence of device DMA. diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 367baf21c35e..52efccb46ebf 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -49,7 +49,7 @@ def capture_aware_stream(stream: torch.Stream) -> torch.Stream: class AscendHiSparseBackend(HiSparseBackend): - # Inherit Core's private Host-pool policy; NPU event IPC is not implemented. + # Keep private pools until NPU IPC events interoperate with Core's generic streams. stats_dtype = torch.int64 def check_available(self) -> None: From ca2212feadafa838a63d1af0d5e90254bd1f33c7 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 10:54:10 +0800 Subject: [PATCH 044/120] perf(hisparse): prototype Target mirror overlap with MTP Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/mirror_overlap.py | 149 +++++++++++++++++++++++++ vllm_ascend/worker/v2/model_runner.py | 5 + 2 files changed, 154 insertions(+) create mode 100644 vllm_ascend/hisparse/mirror_overlap.py diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/mirror_overlap.py new file mode 100644 index 000000000000..7a0f00e430f4 --- /dev/null +++ b/vllm_ascend/hisparse/mirror_overlap.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Private MTP experiment: preserve output durability while overlapping Target D2H. + +The common worker's mirror submission, invalidation, descriptor retirement and +next-step/shutdown fences remain unchanged. Only its immediate compute-stream +join is redirected to the existing output-copy stream, then conservatively +joined on compute at the existing post-draft finalize boundary. No new stream, +event, descriptor pool or asynchronous resource owner is introduced. + +This candidate targets the reviewed Core 24420b7 worker/AsyncOutput contract. +It is not qualified for publication or enabled in the public foundation PR. +""" + +from __future__ import annotations + +import torch + + +def _payload_regions(tensor: torch.Tensor) -> list[tuple[int, int]] | None: + """Describe actual contiguous rows, retaining gaps between Resident pages.""" + if tensor.ndim not in (2, 3) or tensor.stride(-1) != 1 or not all(tensor.shape): + return None + width = tensor.shape[-1] + if tensor.stride(-2) != width: + return None + start = tensor.data_ptr() + row_bytes = width * tensor.element_size() + if tensor.ndim == 2: + return [(start, start + tensor.shape[0] * row_bytes)] + page_bytes = tensor.shape[1] * row_bytes + page_stride = tensor.stride(0) * tensor.element_size() + if page_stride < page_bytes: + return None + if page_stride == page_bytes: + return [(start, start + tensor.shape[0] * page_bytes)] + return [(start + page * page_stride, start + page * page_stride + page_bytes) for page in range(tensor.shape[0])] + + +def _separate_layers(caches, draft_layers: set[int]) -> bool: + """Prove Target DMA cannot race the draft's cache reads/writes at cold bind.""" + regions = [] + for index, cache in enumerate(caches): + payload = _payload_regions(cache) + if payload is None: + return False + regions.extend((start, end, index in draft_layers) for start, end in payload) + # A sweep over occupied ranges also handles interleaved per-page layer views. + furthest_end = {False: -1, True: -1} + for start, end, draft in sorted(regions): + if furthest_end[not draft] > start: + return False + furthest_end[draft] = max(furthest_end[draft], end) + return True + + +class MTPMirrorOverlapConnector: + """Wrap this runner's direct connector; delegate all unmodified boundaries.""" + + def __init__(self, connector, worker, output_copy_stream): + self._connector = connector + self._worker = worker + self._output_copy_stream = output_copy_stream + self._pending_event = None + + def __getattr__(self, name): + return getattr(self._connector, name) + + def _join_compute(self) -> None: + if self._pending_event is not None: + torch.accelerator.current_stream().wait_event(self._pending_event) + self._pending_event = None + + def pre_forward(self, *args, **kwargs): + # Also protect a subsequent call if a proposal exited before finalize. + self._join_compute() + return self._connector.pre_forward(*args, **kwargs) + + def finish_forward(self) -> None: + if self._connector._disabled: + return self._connector.finish_forward() + worker = self._worker + compute_stream = torch.accelerator.current_stream() + worker._forward_ready_event.record(compute_stream) + # This is the original submission + written-slot invalidation routine. + worker._finish_mirror_phase(worker._forward_ready_event) + if worker.is_host_writer: + if worker._dma_submitted: + self._pending_event = worker.host_write_event + # AsyncOutput keeps its original count-copy-before-draft order. + # Its existing completion event cannot publish unfinished Host KV. + self._output_copy_stream.wait_event(self._pending_event) + worker._dma_submitted = False + else: + worker.host_write_event.record(compute_stream) + worker._release_completed_dma_descriptors() + + def post_forward(self, finished_req_ids, defer_finalize=False): + if not defer_finalize: + self._join_compute() + return self._connector.post_forward(finished_req_ids, defer_finalize=defer_finalize) + + def finalize_step(self) -> None: + # Core calls this after the last draft, before clearing step metadata. + self._join_compute() + self._connector.finalize_step() + + def no_forward(self, scheduler_output): + self._join_compute() + return self._connector.no_forward(scheduler_output) + + def set_disabled(self, disabled) -> None: + self._join_compute() + self._connector.set_disabled(disabled) + + def reset_capture_state(self) -> None: + self._join_compute() + self._connector.reset_capture_state() + + +def overlap_mtp_host_mirror(connector, speculator, output_copy_stream): + """Fall back to the common connector unless the independent-cache proof holds.""" + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector + from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + from vllm_ascend.hisparse.backend import AscendHiSparseBackend + + if speculator is None or type(connector) is not ActiveKVConnector: + return connector + if type(connector.kv_connector) is not HiSparseConnector: + return connector + worker = connector.kv_connector.connector_worker + if worker is None or not worker._initialized or not isinstance(worker.backend, AscendHiSparseBackend): + return connector + if worker.shared_host_region is not None or not worker.is_host_writer: + return connector + draft_layers = set(worker._draft_layers) + expected_names = set(getattr(speculator, "draft_attn_layer_names", ())) + if not draft_layers or draft_layers == set(range(len(worker.cache_handles))): + return connector + if {worker.cache_layer_names[index] for index in draft_layers} != expected_names: + return connector + if output_copy_stream == torch.accelerator.current_stream(): + return connector + if not _separate_layers(worker.resident_caches, draft_layers): + return connector + if not _separate_layers(worker.host_caches, draft_layers): + return connector + return MTPMirrorOverlapConnector(connector, worker, output_copy_stream) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index b1370d8b4349..1bde7f4ce779 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -573,6 +573,11 @@ def initialize_kv_cache( self._fill_c8_mxfp_v_scale_caches() + # Private candidate only: bind after Core seals all cache/stream owners. + from vllm_ascend.hisparse.mirror_overlap import overlap_mtp_host_mirror + + self.kv_connector = overlap_mtp_host_mirror(self.kv_connector, self.speculator, self.output_copy_stream) + def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" self._c8_mxfp_v_scale_caches = [] From 7bc7b5daff742127a93b0efdaca9b3c5f0e849a5 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 17:03:59 +0800 Subject: [PATCH 045/120] perf(hisparse): batch row mirror DMA descriptor preparation Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 2 + vllm_ascend/hisparse/row_dma.py | 138 ++++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+) create mode 100644 vllm_ascend/hisparse/row_dma.py diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 52efccb46ebf..860257525289 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -54,6 +54,7 @@ class AscendHiSparseBackend(HiSparseBackend): def check_available(self) -> None: from vllm_ascend.hisparse.resolver import _load_resolver + from vllm_ascend.hisparse.row_dma import install as install_row_dma from vllm_ascend.patch.worker.patch_hisparse_host_pool import install from vllm_ascend.utils import bootstrap_custom_op_env @@ -72,6 +73,7 @@ def check_available(self) -> None: if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") _load_resolver() + install_row_dma(AscendHiSparseBackend) def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Tensor]: from vllm_ascend.hisparse.cache_ops import allocate_mapped_host diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py new file mode 100644 index 000000000000..69fbf8b4d3a4 --- /dev/null +++ b/vllm_ascend/hisparse/row_dma.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Private experiment for the frozen Core HiSparse row-mirror worker.""" + +from dataclasses import dataclass +from functools import wraps +from hashlib import sha256 +from pathlib import Path + +import numpy as np + +_WORKER_SHA256 = "864e2c7ce4e2cff320b4f68f8dccc9691adb829f0e50fdf6a27dc9c7885040e5" +# Large prefill batches keep the scalar path to bound temporary memory. +_MAX_DESCRIPTORS = 32768 + + +@dataclass(frozen=True) +class _Geometry: + sources: tuple + destinations: tuple + block_size: int + source_columns: np.ndarray + source_addresses: np.ndarray + source_page_bytes: np.ndarray + source_pages: np.ndarray + destination_addresses: np.ndarray + destination_rows: np.ndarray + row_bytes: np.ndarray + + @classmethod + def bind(cls, worker): + sources, destinations = worker.resident_caches, worker.host_caches + block_size = worker.kernel_block_size + facts = [] + for handle, source, destination in zip(worker.cache_handles, sources, destinations, strict=True): + column = handle.runtime.resident_source_index + if column < 0: + return None + item_size = source.element_size() + facts.append( + ( + column, + source.data_ptr(), + source.stride(0) * item_size, + source.shape[0], + destination.data_ptr(), + destination.shape[0], + source.shape[-1] * item_size, + ) + ) + if not facts: + return None + columns = np.ascontiguousarray(np.asarray(facts, dtype=np.int64).T) + columns.flags.writeable = False + # Core _init_dma already validated the rows and bound the cache owners. + return cls(sources, destinations, block_size, *columns) + + +def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: + if not layer_indices or not worker._row_mirror_num_rows or not worker.is_host_writer: + return True + num_layers = len(layer_indices) + num_mirrors = len(worker._row_mirrors) + count = num_mirrors * num_layers + if ( + geometry is None + or num_layers < 2 + or count > _MAX_DESCRIPTORS + or geometry.sources is not worker.resident_caches + or geometry.destinations is not worker.host_caches + or geometry.block_size != worker.kernel_block_size + ): + return False + + layers = np.asarray(layer_indices, dtype=np.intp) + starts = worker._row_mirror_source_starts + columns = geometry.source_columns[layers] + if np.any(columns >= starts.shape[1]): + return False + pages, offsets = np.divmod(starts[:, columns], geometry.block_size) + row_counts = worker._row_mirror_counts[:, None] + destinations = worker._row_mirror_destination_starts[:, None] + # Delegate invalid inputs to the original path for its checks and errors. + if ( + np.any(pages < 0) + or np.any(pages >= geometry.source_pages[layers]) + or np.any(offsets + row_counts > geometry.block_size) + or np.any(destinations < 0) + or np.any(destinations + row_counts > geometry.destination_rows[layers]) + ): + return False + + descriptors = worker._acquire_dma_descriptors(count) + shape = (num_mirrors, num_layers) + src = descriptors.src_np[:count].reshape(shape) + dst = descriptors.dst_np[:count].reshape(shape) + row_bytes = geometry.row_bytes[layers] + # The native submission consumes mirror-major descriptors in this order. + np.multiply(pages, geometry.source_page_bytes[layers], out=src) + src += geometry.source_addresses[layers] + src += offsets * row_bytes + np.multiply(destinations, row_bytes, out=dst) + dst += geometry.destination_addresses[layers] + descriptors.sizes_np[:count].reshape(shape)[:] = row_counts * row_bytes + worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) + return True + + +def _install_worker(worker_type, backend_type) -> None: + original_init = worker_type._init_dma + original_enqueue = worker_type._enqueue_row_dma + if getattr(original_enqueue, "_ascend_row_dma", False): + return + + @wraps(original_init) + def init(worker): + original_init(worker) + worker._ascend_row_dma_geometry = None + if isinstance(worker.backend, backend_type) and worker.is_host_writer and worker.shared_host_region is None: + worker._ascend_row_dma_geometry = _Geometry.bind(worker) + + @wraps(original_enqueue) + def enqueue(worker, layer_indices, ready_event=None): + geometry = getattr(worker, "_ascend_row_dma_geometry", None) + if geometry is None or not _try_enqueue(worker, geometry, layer_indices, ready_event): + return original_enqueue(worker, layer_indices, ready_event=ready_event) + + enqueue._ascend_row_dma = True + worker_type._init_dma = init + worker_type._enqueue_row_dma = enqueue + + +def install(backend_type) -> None: + """Only adapt the exact Core implementation reviewed for this experiment.""" + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker + + if sha256(Path(worker.__file__).read_bytes()).hexdigest() == _WORKER_SHA256: + _install_worker(worker.HiSparseConnectorWorker, backend_type) From d9e19fcbe29f1f0e2cff548a6cd1fb8d9bdbe82a Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 18:33:05 +0800 Subject: [PATCH 046/120] perf(hisparse): overlap Host mirror with Target sampling Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/mirror_overlap.py | 189 ++++++++++++------------- vllm_ascend/worker/v2/model_runner.py | 15 +- 2 files changed, 102 insertions(+), 102 deletions(-) diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/mirror_overlap.py index 7a0f00e430f4..619eda946f6f 100644 --- a/vllm_ascend/hisparse/mirror_overlap.py +++ b/vllm_ascend/hisparse/mirror_overlap.py @@ -1,149 +1,138 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Private MTP experiment: preserve output durability while overlapping Target D2H. - -The common worker's mirror submission, invalidation, descriptor retirement and -next-step/shutdown fences remain unchanged. Only its immediate compute-stream -join is redirected to the existing output-copy stream, then conservatively -joined on compute at the existing post-draft finalize boundary. No new stream, -event, descriptor pool or asynchronous resource owner is introduced. - -This candidate targets the reviewed Core 24420b7 worker/AsyncOutput contract. -It is not qualified for publication or enabled in the public foundation PR. -""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Private decode experiment for Core 24420b7's sampling boundary.""" from __future__ import annotations import torch -def _payload_regions(tensor: torch.Tensor) -> list[tuple[int, int]] | None: - """Describe actual contiguous rows, retaining gaps between Resident pages.""" - if tensor.ndim not in (2, 3) or tensor.stride(-1) != 1 or not all(tensor.shape): - return None - width = tensor.shape[-1] - if tensor.stride(-2) != width: - return None - start = tensor.data_ptr() - row_bytes = width * tensor.element_size() - if tensor.ndim == 2: - return [(start, start + tensor.shape[0] * row_bytes)] - page_bytes = tensor.shape[1] * row_bytes - page_stride = tensor.stride(0) * tensor.element_size() - if page_stride < page_bytes: - return None - if page_stride == page_bytes: - return [(start, start + tensor.shape[0] * page_bytes)] - return [(start + page * page_stride, start + page * page_stride + page_bytes) for page in range(tensor.shape[0])] - - -def _separate_layers(caches, draft_layers: set[int]) -> bool: - """Prove Target DMA cannot race the draft's cache reads/writes at cold bind.""" - regions = [] - for index, cache in enumerate(caches): - payload = _payload_regions(cache) - if payload is None: - return False - regions.extend((start, end, index in draft_layers) for start, end in payload) - # A sweep over occupied ranges also handles interleaved per-page layer views. - furthest_end = {False: -1, True: -1} - for start, end, draft in sorted(regions): - if furthest_end[not draft] > start: - return False - furthest_end[draft] = max(furthest_end[draft], end) - return True - - -class MTPMirrorOverlapConnector: - """Wrap this runner's direct connector; delegate all unmodified boundaries.""" - - def __init__(self, connector, worker, output_copy_stream): +class SamplingMirrorOverlapConnector: + """Let Target logits and sampling run before joining the Host mirror.""" + + def __init__(self, connector, worker, output_copy_stream, logger, rank): self._connector = connector self._worker = worker self._output_copy_stream = output_copy_stream + self._logger = logger + self._rank = rank self._pending_event = None + self._allow_step = False + self._eligible_forward = False + self._batch_size = 0 + self._logged_batch_sizes: set[int] = set() def __getattr__(self, name): return getattr(self._connector, name) - def _join_compute(self) -> None: + def join_pending(self) -> None: if self._pending_event is not None: torch.accelerator.current_stream().wait_event(self._pending_event) self._pending_event = None - def pre_forward(self, *args, **kwargs): - # Also protect a subsequent call if a proposal exited before finalize. - self._join_compute() - return self._connector.pre_forward(*args, **kwargs) + def begin_execution(self, allow_step: bool) -> None: + # The runner calls this before preemption or any scheduler state writes. + self.join_pending() + self._allow_step = allow_step + self._eligible_forward = False + + def pre_forward(self, scheduler_output, input_batch=None, attn_metadata=None): + self._eligible_forward = ( + self._allow_step + and input_batch is not None + and attn_metadata is not None + and not input_batch.has_prefill + and 0 < input_batch.num_tokens == input_batch.num_tokens_after_padding + ) + if self._eligible_forward: + self._batch_size = input_batch.num_reqs + return self._connector.pre_forward(scheduler_output, input_batch, attn_metadata=attn_metadata) def finish_forward(self) -> None: - if self._connector._disabled: + eligible = self._eligible_forward + self._eligible_forward = False + if not eligible or self._connector._disabled: return self._connector.finish_forward() + worker = self._worker compute_stream = torch.accelerator.current_stream() worker._forward_ready_event.record(compute_stream) - # This is the original submission + written-slot invalidation routine. worker._finish_mirror_phase(worker._forward_ready_event) - if worker.is_host_writer: - if worker._dma_submitted: - self._pending_event = worker.host_write_event - # AsyncOutput keeps its original count-copy-before-draft order. - # Its existing completion event cannot publish unfinished Host KV. - self._output_copy_stream.wait_event(self._pending_event) - worker._dma_submitted = False - else: - worker.host_write_event.record(compute_stream) + if worker._dma_submitted: + self._pending_event = worker.host_write_event + # Preserve the Host durability fence on AsyncOutput's copy stream. + self._output_copy_stream.wait_event(self._pending_event) + worker._dma_submitted = False + if self._batch_size not in self._logged_batch_sizes: + self._logged_batch_sizes.add(self._batch_size) + self._logger.info("HiSparse MTP mirror overlap used: rank=%d requests=%d", self._rank, self._batch_size) + else: + worker.host_write_event.record(compute_stream) worker._release_completed_dma_descriptors() def post_forward(self, finished_req_ids, defer_finalize=False): - if not defer_finalize: - self._join_compute() + # In the audited runner, sampling and output-copy submission precede + # this call. Draft preparation and all draft KV access follow it. + self.join_pending() return self._connector.post_forward(finished_req_ids, defer_finalize=defer_finalize) - def finalize_step(self) -> None: - # Core calls this after the last draft, before clearing step metadata. - self._join_compute() - self._connector.finalize_step() - def no_forward(self, scheduler_output): - self._join_compute() + self.begin_execution(False) return self._connector.no_forward(scheduler_output) - def set_disabled(self, disabled) -> None: - self._join_compute() + def set_disabled(self, disabled: bool) -> None: + self.begin_execution(False) self._connector.set_disabled(disabled) def reset_capture_state(self) -> None: - self._join_compute() + self.begin_execution(False) self._connector.reset_capture_state() -def overlap_mtp_host_mirror(connector, speculator, output_copy_stream): - """Fall back to the common connector unless the independent-cache proof holds.""" +def overlap_mtp_host_mirror(runner): + """Bind only the audited single-stage GLM MTP sampling path.""" from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + from vllm.logger import init_logger from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector from vllm_ascend.hisparse.backend import AscendHiSparseBackend - if speculator is None or type(connector) is not ActiveKVConnector: - return connector - if type(connector.kv_connector) is not HiSparseConnector: + connector = runner.kv_connector + config = runner.vllm_config.parallel_config + if ( + type(connector) is not ActiveKVConnector + or type(connector.kv_connector) is not HiSparseConnector + or runner.speculator is None + or runner.speculative_config.method != "mtp" + or runner.model_config.architecture != "GlmMoeDsaForCausalLM" + or config.pipeline_parallel_size != 1 + or config.data_parallel_size != 1 + or config.decode_context_parallel_size != 1 + or runner.pcp_manager is not None + or runner.is_pooling_model + or runner.supports_mm_inputs + or runner.aux_output_connector is not None + or runner.kvpp.scheduler is not None + or getattr(runner, "pd_dspark_aux_layer_ids", ()) + ): return connector worker = connector.kv_connector.connector_worker - if worker is None or not worker._initialized or not isinstance(worker.backend, AscendHiSparseBackend): - return connector - if worker.shared_host_region is not None or not worker.is_host_writer: - return connector - draft_layers = set(worker._draft_layers) - expected_names = set(getattr(speculator, "draft_attn_layer_names", ())) - if not draft_layers or draft_layers == set(range(len(worker.cache_handles))): - return connector - if {worker.cache_layer_names[index] for index in draft_layers} != expected_names: - return connector - if output_copy_stream == torch.accelerator.current_stream(): - return connector - if not _separate_layers(worker.resident_caches, draft_layers): + if ( + worker is None + or not worker._initialized + or not isinstance(worker.backend, AscendHiSparseBackend) + or worker.shared_host_region is not None + or not worker.is_host_writer + ): return connector - if not _separate_layers(worker.host_caches, draft_layers): + output = runner.output_copy_stream + output_stream = torch.Stream( + stream_id=output.stream_id, device_index=output.device_index, device_type=output.device_type + ) + if output_stream == torch.accelerator.current_stream(): return connector - return MTPMirrorOverlapConnector(connector, worker, output_copy_stream) + logger = init_logger(__name__) + rank = get_tensor_model_parallel_rank() + logger.info("HiSparse MTP mirror overlap enabled: rank=%d sampling only; join before draft", rank) + return SamplingMirrorOverlapConnector(connector, worker, output_stream, logger, rank) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 1bde7f4ce779..eecce6a9fc89 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -393,7 +393,11 @@ def sample_tokens(self, grammar_output): if pcp_manager is not None and isinstance(pcp_manager, AscendPCPManager): pcp_manager._sampling_hidden_restored = False self._restore_replicated_draft_target_states() - output = super().sample_tokens(grammar_output) + try: + output = super().sample_tokens(grammar_output) + finally: + if (mirror := getattr(self, "_sampling_mirror", None)) is not None: + mirror.join_pending() # Close the per-step msprobe dump cycle on the normal last-PP-rank # path (v1 parity: v1 finalizes at the end of sample_tokens). self._finalize_dump_data() @@ -576,7 +580,9 @@ def initialize_kv_cache( # Private candidate only: bind after Core seals all cache/stream owners. from vllm_ascend.hisparse.mirror_overlap import overlap_mtp_host_mirror - self.kv_connector = overlap_mtp_host_mirror(self.kv_connector, self.speculator, self.output_copy_stream) + overlap_connector = overlap_mtp_host_mirror(self) + self._sampling_mirror = overlap_connector if overlap_connector is not self.kv_connector else None + self.kv_connector = overlap_connector def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" @@ -668,6 +674,9 @@ def execute_model( valid_dummy_state_slots: bool = False, randomize_inputs: bool = False, ): + mirror = getattr(self, "_sampling_mirror", None) + if mirror is not None: + mirror.begin_execution(not dummy_run and not is_profile) self._cpp_execution_time_ms = None profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config execution_start_time = _start_profiling_chunk_timing( @@ -709,6 +718,8 @@ def execute_model( ) forward_failed = False finally: + if forward_failed and mirror is not None: + mirror.join_pending() finish_execution = getattr(self.model_state, "finish_execution", None) if finish_execution is not None: finish_execution(failed=forward_failed) From 42b3aeffda9e369b494314be5d1da07e7f580093 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 19:06:12 +0800 Subject: [PATCH 047/120] perf(hisparse): stage Target rows before overlapping draft Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/mirror_overlap.py | 63 +++++++-- vllm_ascend/hisparse/row_dma.py | 4 +- vllm_ascend/hisparse/target_snapshot.py | 174 ++++++++++++++++++++++++ 3 files changed, 232 insertions(+), 9 deletions(-) create mode 100644 vllm_ascend/hisparse/target_snapshot.py diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/mirror_overlap.py index 619eda946f6f..3dbcd4f7ee7a 100644 --- a/vllm_ascend/hisparse/mirror_overlap.py +++ b/vllm_ascend/hisparse/mirror_overlap.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Private decode experiment for Core 24420b7's sampling boundary.""" +"""Private decode experiment for Core 24420b7's sampling and draft boundaries.""" from __future__ import annotations @@ -10,17 +10,21 @@ class SamplingMirrorOverlapConnector: """Let Target logits and sampling run before joining the Host mirror.""" - def __init__(self, connector, worker, output_copy_stream, logger, rank): + def __init__(self, connector, worker, output_copy_stream, logger, rank, snapshot=None): self._connector = connector self._worker = worker self._output_copy_stream = output_copy_stream self._logger = logger self._rank = rank + self._snapshot = snapshot + self._overlap_draft = False self._pending_event = None self._allow_step = False self._eligible_forward = False self._batch_size = 0 self._logged_batch_sizes: set[int] = set() + self._logged_snapshot_sizes: set[int] = set() + self._logged_direct_sizes: set[int] = set() def __getattr__(self, name): return getattr(self._connector, name) @@ -29,6 +33,7 @@ def join_pending(self) -> None: if self._pending_event is not None: torch.accelerator.current_stream().wait_event(self._pending_event) self._pending_event = None + self._overlap_draft = False def begin_execution(self, allow_step: bool) -> None: # The runner calls this before preemption or any scheduler state writes. @@ -39,6 +44,7 @@ def begin_execution(self, allow_step: bool) -> None: def pre_forward(self, scheduler_output, input_batch=None, attn_metadata=None): self._eligible_forward = ( self._allow_step + and scheduler_output.has_sync_kv_loads and input_batch is not None and attn_metadata is not None and not input_batch.has_prefill @@ -57,12 +63,45 @@ def finish_forward(self) -> None: worker = self._worker compute_stream = torch.accelerator.current_stream() worker._forward_ready_event.record(compute_stream) - worker._finish_mirror_phase(worker._forward_ready_event) + snapshot = self._snapshot + pending_before = 0 + if snapshot is not None: + pending_before = len(worker._pending_dma_descriptors) + snapshot.submitted_completion = None + snapshot.armed = ( + not worker._dma_submitted and not worker._per_layer_mirrored and not worker._submitted_mirror_layers + ) + try: + worker._finish_mirror_phase(worker._forward_ready_event) + except BaseException: + # The runner's failure boundary joins an already-recorded copy. + if worker._dma_submitted: + self._pending_event = worker.host_write_event + raise + finally: + if snapshot is not None: + snapshot.armed = False if worker._dma_submitted: self._pending_event = worker.host_write_event # Preserve the Host durability fence on AsyncOutput's copy stream. self._output_copy_stream.wait_event(self._pending_event) worker._dma_submitted = False + if snapshot is not None and snapshot.submitted_completion is not None: + pending = worker._pending_dma_descriptors + # Direct fallback or an extra transfer keeps the sampling-only fence. + self._overlap_draft = ( + len(pending) == pending_before + 1 and pending[-1][0] is snapshot.submitted_completion + ) + if self._overlap_draft and self._batch_size not in self._logged_snapshot_sizes: + self._logged_snapshot_sizes.add(self._batch_size) + self._logger.info( + "HiSparse Target snapshot used: rank=%d requests=%d", self._rank, self._batch_size + ) + if snapshot is not None and not self._overlap_draft and self._batch_size not in self._logged_direct_sizes: + self._logged_direct_sizes.add(self._batch_size) + self._logger.info( + "HiSparse Target snapshot fallback: rank=%d requests=%d", self._rank, self._batch_size + ) if self._batch_size not in self._logged_batch_sizes: self._logged_batch_sizes.add(self._batch_size) self._logger.info("HiSparse MTP mirror overlap used: rank=%d requests=%d", self._rank, self._batch_size) @@ -71,11 +110,14 @@ def finish_forward(self) -> None: worker._release_completed_dma_descriptors() def post_forward(self, finished_req_ids, defer_finalize=False): - # In the audited runner, sampling and output-copy submission precede - # this call. Draft preparation and all draft KV access follow it. - self.join_pending() + if not defer_finalize or not self._overlap_draft: + self.join_pending() return self._connector.post_forward(finished_req_ids, defer_finalize=defer_finalize) + def finalize_step(self): + self.join_pending() + return self._connector.finalize_step() + def no_forward(self, scheduler_output): self.begin_execution(False) return self._connector.no_forward(scheduler_output) @@ -97,6 +139,7 @@ def overlap_mtp_host_mirror(runner): from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector from vllm_ascend.hisparse.backend import AscendHiSparseBackend + from vllm_ascend.hisparse.target_snapshot import TargetMirrorSnapshot connector = runner.kv_connector config = runner.vllm_config.parallel_config @@ -134,5 +177,9 @@ def overlap_mtp_host_mirror(runner): return connector logger = init_logger(__name__) rank = get_tensor_model_parallel_rank() - logger.info("HiSparse MTP mirror overlap enabled: rank=%d sampling only; join before draft", rank) - return SamplingMirrorOverlapConnector(connector, worker, output_stream, logger, rank) + snapshot = TargetMirrorSnapshot.bind(worker, runner.speculator) + worker._ascend_target_snapshot = snapshot + logger.info("HiSparse MTP mirror overlap enabled: rank=%d", rank) + if snapshot is not None: + logger.info("HiSparse Target snapshot enabled: rank=%d", rank) + return SamplingMirrorOverlapConnector(connector, worker, output_stream, logger, rank, snapshot) diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index 69fbf8b4d3a4..f7523c41c9c9 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -102,7 +102,9 @@ def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: np.multiply(destinations, row_bytes, out=dst) dst += geometry.destination_addresses[layers] descriptors.sizes_np[:count].reshape(shape)[:] = row_counts * row_bytes - worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) + snapshot = getattr(worker, "_ascend_target_snapshot", None) + if snapshot is None or not snapshot.try_submit(worker, descriptors, count, layer_indices, ready_event): + worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) return True diff --git a/vllm_ascend/hisparse/target_snapshot.py b/vllm_ascend/hisparse/target_snapshot.py new file mode 100644 index 000000000000..ece4ed6c5ac1 --- /dev/null +++ b/vllm_ascend/hisparse/target_snapshot.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Private Target snapshot experiment for the frozen HiSparse worker.""" + +from dataclasses import dataclass + +import numpy as np +import torch + + +def _host_disjoint(caches, draft_layers): + if any(cache.ndim != 2 or not cache.is_contiguous() or cache.numel() == 0 for cache in caches): + return False + spans = [(cache.data_ptr(), cache.data_ptr() + cache.numel() * cache.element_size()) for cache in caches] + return all( + left_end <= right_start or right_end <= left_start + for index, (left_start, left_end) in enumerate(spans) + if index in draft_layers + for other, (right_start, right_end) in enumerate(spans) + if other not in draft_layers + ) + + +def _draft_groups_disjoint(handles, draft_layers): + draft_runtimes = {id(handles[index].runtime) for index in draft_layers} + # A draft-led resolver can refill every follower from Host in one call. + for index in draft_layers: + group = handles[index].runtime.index_group + if id(group.leader) not in draft_runtimes or any( + id(follower) not in draft_runtimes for follower in group.followers + ): + return False + return True + + +def _index_plan(addresses, sizes, base, width, source_rows, limit): + if not len(addresses): + return None + first, alignment = np.divmod(addresses - base, width) + counts, size_remainder = np.divmod(sizes, width) + if ( + np.any(alignment) + or np.any(size_remainder) + or np.any(counts <= 0) + or np.any(first < 0) + or np.any(first + counts > source_rows) + ): + return None + ends = np.cumsum(counts) + total = int(ends[-1]) + if total > limit: + return None + starts = ends - counts + rows = np.repeat(first - starts, counts) + np.arange(total, dtype=np.int64) + return starts, rows + + +@dataclass +class _Scratch: + cpu_indices: torch.Tensor + device_indices: torch.Tensor + payload: torch.Tensor + ready: torch.Event + completion: torch.Event | None = None + in_flight: bool = False + + def available(self): + return not self.in_flight or (self.completion is not None and self.completion.query()) + + +class TargetMirrorSnapshot: + """Stage exact Target source bytes before draft can overwrite shared HMA.""" + + def __init__(self, worker, layers, width, capacity): + backing = worker.hot_backing + self._source = backing.view(torch.uint8).reshape(-1, width) + self._base = backing.data_ptr() + self._width = width + self._layers = layers + self._sources = worker.resident_caches + self._destinations = worker.host_caches + self._capacity = capacity + self._slots = tuple( + _Scratch( + torch.empty(capacity, dtype=torch.int64, pin_memory=True), + torch.empty(capacity, dtype=torch.int64, device=backing.device), + torch.empty((capacity, width), dtype=torch.uint8, device=backing.device), + torch.Event(), + ) + for _ in range(2) + ) + self.armed = False + self.submitted_completion = None + + @classmethod + def bind(cls, worker, speculator): + geometry = getattr(worker, "_ascend_row_dma_geometry", None) + draft = set(worker._draft_layers) + names = worker.cache_layer_names + expected = set(getattr(speculator, "draft_attn_layer_names", ())) & set(names) + if geometry is None or not draft or {names[i] for i in draft} != expected: + return None + layers = tuple(index for index in range(len(names)) if index not in draft) + if ( + not layers + or not _draft_groups_disjoint(worker.cache_handles, draft) + or not _host_disjoint(worker.host_caches, draft) + ): + return None + widths = {int(geometry.row_bytes[i]) for i in layers} + if len(widths) != 1: + return None + width = widths.pop() + backing = worker.hot_backing + base = backing.data_ptr() + if ( + width <= 0 + or backing.device.type != "npu" + or backing.dtype not in (torch.int8, torch.uint8) + or backing.ndim != 1 + or not backing.is_contiguous() + or backing.storage_offset() != 0 + or backing.numel() % width + or any( + cache.untyped_storage().data_ptr() != base + or (cache.data_ptr() - base) % width + or geometry.source_page_bytes[index] % width + for index in layers + for cache in (worker.resident_caches[index],) + ) + ): + return None + config = worker.vllm_config + capacity = config.scheduler_config.max_num_seqs * (config.num_lookahead_tokens + 1) * len(layers) + return cls(worker, layers, width, capacity) + + def try_submit(self, worker, descriptors, count, layer_indices, ready_event): + if ( + not self.armed + or self.submitted_completion is not None + or ready_event is not worker._forward_ready_event + or tuple(layer_indices) != self._layers + or worker.resident_caches is not self._sources + or worker.host_caches is not self._destinations + ): + return False + slot = next((item for item in self._slots if item.available()), None) + if slot is None: + return False + plan = _index_plan( + descriptors.src_np[:count], + descriptors.sizes_np[:count], + self._base, + self._width, + self._source.shape[0], + self._capacity, + ) + if plan is None: + return False + starts, rows = plan + size = len(rows) + # Owners remain pinned if any submission fails. Never retry a partial + # enqueue using the original device addresses. + slot.in_flight = True + slot.completion = None + slot.cpu_indices[:size].numpy()[:] = rows + slot.device_indices[:size].copy_(slot.cpu_indices[:size], non_blocking=True) + torch.index_select(self._source, 0, slot.device_indices[:size], out=slot.payload[:size]) + slot.ready.record() + descriptors.src_np[:count] = slot.payload.data_ptr() + starts * self._width + worker._submit_dma_descriptors(descriptors, count, ready_event=slot.ready) + slot.completion = worker._pending_dma_descriptors[-1][0] + self.submitted_completion = slot.completion + return True From 23949ef40670afba0ac2f60076467382529d895e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 19:38:41 +0800 Subject: [PATCH 048/120] perf(hisparse): reuse Target RoPE metadata within decode forwards Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/indexer.py | 9 ++++- vllm_ascend/attention/sfa_hisparse.py | 12 +++++- vllm_ascend/attention/sfa_v1.py | 15 ++++++- vllm_ascend/ops/rotary_embedding.py | 40 +++++++++++++++++++ vllm_ascend/worker/v2/attn_utils.py | 34 ++++++++++++++++ vllm_ascend/worker/v2/model_states/default.py | 13 ++++-- 6 files changed, 114 insertions(+), 9 deletions(-) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 3ea0f6cf9159..158ebf0c978b 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -36,7 +36,7 @@ ) from vllm_ascend.attention.utils import split_decodes_and_prefills, uses_pcp_sharded_layout from vllm_ascend.device.device_op import DeviceOperator -from vllm_ascend.ops.rotary_embedding import get_cos_and_sin_mla +from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache, get_cos_and_sin_mla from vllm_ascend.ops.triton.rope import rope_forward_triton_siso from vllm_ascend.ops.triton.sfa_indexer_store import store_indexer_key_scale from vllm_ascend.utils import ( @@ -928,6 +928,8 @@ def _build( buffer_key: object, use_cached_rope: bool, copy_rope: bool, + *, + shared_mla_rope: MLARopeMetadataCache | None = None, **kwargs, ) -> AscendSFAIndexerMetadata: num_reqs = common_attn_metadata.num_reqs @@ -953,7 +955,10 @@ def _build( input_positions = common_attn_metadata.positions[:num_input_tokens].long() block_size = self.kernel_block_size - cos, sin = get_cos_and_sin_mla(input_positions, use_cache=use_cached_rope) + if shared_mla_rope is not None and use_cached_rope: + cos, sin = shared_mla_rope.get(common_attn_metadata.positions, input_positions) + else: + cos, sin = get_cos_and_sin_mla(input_positions, use_cache=use_cached_rope) cos = cos[:num_input_tokens] sin = sin[:num_input_tokens] cum_query_lens = common_attn_metadata.query_start_loc[1 : num_reqs + 1] diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 2491604de4bb..07f077d845f8 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -14,6 +14,7 @@ from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder from vllm_ascend.attention.utils import PreprocessType from vllm_ascend.hisparse.prefill import AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan +from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache # Bound the temporary D2D packing workspace independently of prefill length. _ATTENTION_CHUNK_ROWS = 32 @@ -57,8 +58,15 @@ def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): vllm_config.compilation_config.max_cudagraph_capture_size or 0, ) - def _build(self, common_attn_metadata, draft_index=None, pcp_context=None): - metadata = super()._build(common_attn_metadata, draft_index, pcp_context) + def _build( + self, + common_attn_metadata, + draft_index=None, + pcp_context=None, + *, + shared_mla_rope: MLARopeMetadataCache | None = None, + ): + metadata = super()._build(common_attn_metadata, draft_index, pcp_context, shared_mla_rope=shared_mla_rope) metadata.num_reqs = common_attn_metadata.num_reqs if draft_index not in self._request_row_buffers: self._request_row_buffers[draft_index] = torch.empty( diff --git a/vllm_ascend/attention/sfa_v1.py b/vllm_ascend/attention/sfa_v1.py index c8b4fc01121c..9b328493d034 100644 --- a/vllm_ascend/attention/sfa_v1.py +++ b/vllm_ascend/attention/sfa_v1.py @@ -56,7 +56,7 @@ OFFLOAD_KV_CACHE_TUPLE_LEN, OFFLOAD_V_CACHE_NPU_INDEX, ) -from vllm_ascend.ops.rotary_embedding import get_cos_and_sin_mla +from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache, get_cos_and_sin_mla from vllm_ascend.quantization.methods import ( AscendW8A8DynamicLinearMethod, AscendW8A8LinearMethod, @@ -588,12 +588,19 @@ def build( common_prefix_len: int, common_attn_metadata: AscendCommonAttentionMetadata, fast_build: bool = False, + *, + shared_mla_rope: MLARopeMetadataCache | None = None, **kwargs, ) -> AscendSFAMetadata: # common_prefix_len / fast_build are unused; kept for API compatibility. metadata = self._build_with_metadata_view( common_attn_metadata, - lambda: self._build(common_attn_metadata, draft_index=None, pcp_context=common_attn_metadata.pcp_context), + lambda: self._build( + common_attn_metadata, + draft_index=None, + pcp_context=common_attn_metadata.pcp_context, + shared_mla_rope=shared_mla_rope, + ), ) if ( self.use_pcp @@ -647,6 +654,8 @@ def _build( common_attn_metadata: AscendCommonAttentionMetadata, draft_index: int | None = None, pcp_context: Any | None = None, + *, + shared_mla_rope: MLARopeMetadataCache | None = None, ) -> AscendSFAMetadata: num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens @@ -676,6 +685,8 @@ def _build( if self.nope: cos, sin = None, None + elif shared_mla_rope is not None and draft_index is None: + cos, sin = shared_mla_rope.get(common_attn_metadata.positions, input_positions) else: cos, sin = get_cos_and_sin_mla(input_positions, use_cache=(draft_index is None)) diff --git a/vllm_ascend/ops/rotary_embedding.py b/vllm_ascend/ops/rotary_embedding.py index fad9d879b456..1ae120be74aa 100644 --- a/vllm_ascend/ops/rotary_embedding.py +++ b/vllm_ascend/ops/rotary_embedding.py @@ -148,6 +148,46 @@ def get_cos_and_sin_mla(positions, use_cache=False): return _cos_mla[:num_tokens, ...], _sin_mla[:num_tokens, ...] +class MLARopeMetadataCache: + """Share a persistent RoPE slice across compatible builders in one forward.""" + + def __init__(self) -> None: + self._positions: torch.Tensor | None = None + self._num_tokens = 0 + self._value: tuple[torch.Tensor, torch.Tensor] | None = None + + @staticmethod + def _uses_persistent_buffers(value: tuple[torch.Tensor, torch.Tensor]) -> bool: + cos, sin = value + return ( + _cos_mla is not None + and _sin_mla is not None + and cos.numel() > 0 + and sin.numel() > 0 + and cos.device == _cos_mla.device + and sin.device == _sin_mla.device + and cos.data_ptr() == _cos_mla.data_ptr() + and sin.data_ptr() == _sin_mla.data_ptr() + ) + + def get(self, positions: torch.Tensor, input_positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = input_positions.shape[0] + if ( + self._positions is positions + and self._num_tokens == num_tokens + and self._value is not None + and self._uses_persistent_buffers(self._value) + ): + return self._value + value = get_cos_and_sin_mla(input_positions, use_cache=True) + self._positions = positions + self._num_tokens = num_tokens + # Preserve the original fallback for absent, empty or mismatched graph + # buffers. Only slices of the stable destinations may be reused. + self._value = value if self._uses_persistent_buffers(value) else None + return value + + def _record_cos_sin_cache(cos_sin_cache): global _cos_sin_cache if _cos_sin_cache is not None: diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 282dd646b435..eb11f7e080dc 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -109,6 +109,7 @@ make_qsa_kernel_page_cache, make_slab_view, ) +from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache from vllm_ascend.quantization.methods.kv_cache.turboquant import TURBOQUANT_CACHE_DTYPE from vllm_ascend.quantization.methods.kv_cache.turboquant.cache import uses_turboquant_groups from vllm_ascend.quantization.utils import enable_fa_quant @@ -501,6 +502,7 @@ def build_attn_metadata( copy_sfa_restore_tails: bool = False, draft_layer_names: set[str] | None = None, common_metadata_callback=None, + reuse_hisparse_mla_rope: bool = False, ) -> dict[str, Any]: """Build attention metadata for Ascend NPUs.""" if skip_ring_state_update is None: @@ -547,6 +549,30 @@ def build_attn_metadata( if num_actual_reqs is None: num_actual_reqs = num_reqs + rope_builders: tuple[type, ...] = () + if ( + reuse_hisparse_mla_rope + and full_graph_mode + and not for_cudagraph_capture + and not offload_dummy + and positions is not None + and pcp_context is None + and dcp_local_seq_lens is None + and parallel_config is not None + and parallel_config.prefill_context_parallel_size == 1 + and parallel_config.decode_context_parallel_size == 1 + and model_specific_attn_metadata is None + and common_metadata_callback is None + and not draft_layer_names + and copy_sfa_draft_index is None + and not copy_sfa_restore_tails + ): + from vllm_ascend.attention.indexer import AscendHiSparseIndexerMetadataBuilder + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder + + rope_builders = (AscendHiSparseMetadataBuilder, AscendHiSparseIndexerMetadataBuilder) + shared_mla_rope = None + # positions will not be used directly in graph modes, so it is saft to create it here. if positions is None: positions = torch.zeros(num_input_tokens, dtype=torch.int64, device=query_start_loc_gpu.device) @@ -672,6 +698,14 @@ def build_attn_metadata( if not for_cudagraph_capture and model_specific_attn_metadata is not None else {} ) + if type(attn_metadata_builder) in rope_builders and not getattr(attn_metadata_builder, "nope", False): + if shared_mla_rope is None: + shared_mla_rope = MLARopeMetadataCache() + attn_metadata_extra_kwargs["shared_mla_rope"] = shared_mla_rope + else: + # An unrelated builder may update the same persistent RoPE + # storage. Share only within runs of known compatible builders. + shared_mla_rope = None if is_dsa_builder: # DSA cache groups share request-level metadata during replay. attn_metadata_extra_kwargs.update( diff --git a/vllm_ascend/worker/v2/model_states/default.py b/vllm_ascend/worker/v2/model_states/default.py index d56ae95975ea..e2054834c540 100644 --- a/vllm_ascend/worker/v2/model_states/default.py +++ b/vllm_ascend/worker/v2/model_states/default.py @@ -561,9 +561,16 @@ def prepare_attn( copy_sfa_restore_tails=restore_copy_sfa_tails, draft_layer_names=getattr(self, "_offload_draft_layer_names", None), common_metadata_callback=( - lambda group_id, common: self._decoder_replay_common_metadata.__setitem__(group_id, common) - if self.decoder_replay is not None - else None + self._decoder_replay_common_metadata.__setitem__ if self.decoder_replay is not None else None + ), + reuse_hisparse_mla_rope=( + self.vllm_config.attention_config.hisparse_config is not None + and cudagraph_mode == CUDAGraphMode.FULL + and not for_capture + and not input_batch.is_dummy + # MTP verification can be classified as ChunkedPrefill even + # with no requests actually in prefill. + and not input_batch.is_prefilling_np[:num_actual_reqs].any() ), ) self._prepare_decoder_replay(input_batch, cudagraph_mode, attn_groups, kv_cache_config) From 42ce764cd7b50be1551b9b0b3fdde99741c0bc15 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 19:53:31 +0800 Subject: [PATCH 049/120] perf(hisparse): write Target Prolog C8 rows directly to Resident Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 72 +++++++++++++++++++++++---- vllm_ascend/worker/v2/model_runner.py | 37 ++++++++++++++ 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 07f077d845f8..ae05f54a4916 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -7,6 +7,8 @@ import torch import torch_npu +from vllm.forward_context import get_forward_context +from vllm.logger import logger from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.worker.workspace import current_workspace_manager @@ -138,6 +140,7 @@ def __init__(self, *args, **kwargs): if not isinstance(self.index_group, HiSparseMLAIndexGroup): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") self.hisparse_cache = self.index_group.cache(self.index_group_index) + self._direct_prolog_slots: torch.Tensor | None = None self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: raise ValueError("Ascend HiSparse SFA requires top-k to be divisible by 128") @@ -232,6 +235,20 @@ def allocate(): ("ascend_hisparse_c8_prolog", self._kv_staging_capacity, self.c8_cache_dtype), allocate ) + def _resident_prolog_slots_i64(self, mapping): + context = get_forward_context() + scope = getattr(context, "_ascend_hisparse_resident_i64", None) + if scope is None: + scope = context._ascend_hisparse_resident_i64 = {} + # Exact slot aliases across layers share one cast. The captured cast + # reads current slots again on replay, even when the address is stable. + key = (mapping.data_ptr(), mapping.numel(), mapping.device, torch.npu.current_stream().npu_stream) + cached = scope.get(key) + if cached is None: + cached = (mapping.untyped_storage(), mapping.to(torch.int64)) + scope[key] = cached + return cached[1] + def _sfa_preprocess_prolog_v3(self, hidden_states, kv_cache, cos, sin, slot_mapping, *, attn_metadata=None): from vllm_ascend.hisparse.c8 import C8_ROW_BYTES from vllm_ascend.hisparse.cache_ops import scatter_packed_kv @@ -241,16 +258,53 @@ def _sfa_preprocess_prolog_v3(self, hidden_states, kv_cache, cos, sin, slot_mapp rows = hidden_states.shape[0] if rows > self._kv_staging_capacity: raise ValueError("HiSparse Prolog exceeds its profiled token capacity") - staging, local_slots = self._c8_prolog_workspace() - result = super()._sfa_preprocess_prolog_v3( - hidden_states, (staging,), cos, sin, local_slots[:rows], attn_metadata=attn_metadata - ) - destination, mapping, count = self.hisparse_cache.write_target(rows, slot_mapping.numel()) - packed = staging.view(-1, C8_ROW_BYTES)[:count].view(torch.uint8) - scatter_packed_kv( - packed, packed[:, :0], destination.view(torch.uint8), mapping, attn_metadata.query_start_loc[-1:] + cache = self.hisparse_cache + destination, mapping, count = cache.write_target(rows, slot_mapping.numel()) + resident_slots = self._direct_prolog_slots + direct = ( + resident_slots is not None + and not cache.draft_layer + and attn_metadata.hisparse_for_capture + and attn_metadata.num_prefills == 0 + and 0 < count == rows + and mapping.shape == (rows,) + and mapping.stride(0) == 1 + and mapping.dtype == torch.int32 + and mapping.device == resident_slots.device == destination.device == hidden_states.device + and mapping.data_ptr() == resident_slots.data_ptr() + and destination.ndim == 3 + and destination.shape[1:] == (128, C8_ROW_BYTES) + and destination.dtype == torch.float8_e4m3fn + and destination.storage_offset() % C8_ROW_BYTES == 0 + and destination.stride()[1:] == (C8_ROW_BYTES, 1) + and destination.stride(0) >= 128 * C8_ROW_BYTES + and destination.stride(0) % (128 * C8_ROW_BYTES) == 0 ) - finish_hisparse_kv_update(self.hisparse_cache, attn_metadata) + if direct: + logger.info_once( + "HiSparse Target direct Prolog capture: rows=%d stride_bytes=%d", + rows, + destination.stride(0), + scope="process", + ) + result = super()._sfa_preprocess_prolog_v3( + hidden_states, + (destination.unsqueeze(2),), + cos, + sin, + self._resident_prolog_slots_i64(mapping), + attn_metadata=attn_metadata, + ) + else: + staging, local_slots = self._c8_prolog_workspace() + result = super()._sfa_preprocess_prolog_v3( + hidden_states, (staging,), cos, sin, local_slots[:rows], attn_metadata=attn_metadata + ) + packed = staging.view(-1, C8_ROW_BYTES)[:count].view(torch.uint8) + scatter_packed_kv( + packed, packed[:, :0], destination.view(torch.uint8), mapping, attn_metadata.query_start_loc[-1:] + ) + finish_hisparse_kv_update(cache, attn_metadata) return result def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=None): diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index eecce6a9fc89..a3f479f4c1ef 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -577,6 +577,9 @@ def initialize_kv_cache( self._fill_c8_mxfp_v_scale_caches() + if self.vllm_config.attention_config.hisparse_config is not None: + self._bind_hisparse_direct_prolog() + # Private candidate only: bind after Core seals all cache/stream owners. from vllm_ascend.hisparse.mirror_overlap import overlap_mtp_host_mirror @@ -584,6 +587,40 @@ def initialize_kv_cache( self._sampling_mirror = overlap_connector if overlap_connector is not self.kv_connector else None self.kv_connector = overlap_connector + def _bind_hisparse_direct_prolog(self) -> None: + from vllm.v1.attention.backends.utils import PAD_SLOT_ID + from vllm.v1.hisparse.layout import HISPARSE_RESIDENT_SUFFIX + from vllm.v1.kv_cache_interface import HiSparseResidentSpec + + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + from vllm_ascend.worker.v2.block_table import AscendBlockTables + + block_tables = self.block_tables + # These graphs only replay decode, and the exact slot producer pads + # dummy requests and every unused row with native Prolog's -1 sentinel. + enable = ( + get_ascend_device_type() == AscendDeviceType.A5 + and self.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY + and type(block_tables) is AscendBlockTables + and self.pcp_manager is None + and block_tables.cp_size == 1 + and PAD_SLOT_ID == -1 + ) + layers = self.compilation_config.static_forward_context + for group_id, group in enumerate(self.kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): + continue + for cache_name in group.layer_names: + layer = layers[cache_name.removesuffix(HISPARSE_RESIDENT_SUFFIX)] + impl = layer.impl + if type(impl) is AscendHiSparseImpl: + impl._direct_prolog_slots = ( + block_tables.slot_mappings[group_id] + if enable and block_tables.blocks_per_kv_block[group_id] == 1 + else None + ) + def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" self._c8_mxfp_v_scale_caches = [] From a47cbfd70465270cfef3750f7864824ea95157da Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 22:19:46 +0800 Subject: [PATCH 050/120] perf(hisparse): prefetch shared-indexer followers for C8 MTP Target Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 58 +++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 860257525289..9618f382cb04 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -5,9 +5,14 @@ from __future__ import annotations import importlib +from functools import cache, wraps +from hashlib import sha256 +from pathlib import Path from typing import TYPE_CHECKING, Any import torch +from vllm.config import get_current_vllm_config_or_none +from vllm.config.compilation import CUDAGraphMode from vllm.v1.hisparse.backend import HiSparseBackend if TYPE_CHECKING: @@ -48,6 +53,48 @@ def capture_aware_stream(stream: torch.Stream) -> torch.Stream: return _CaptureAwareStream(stream.stream_id, stream.device_index, stream.device_type) +@cache +def _has_follower_prefetch_contract() -> bool: + from vllm.v1.hisparse import runtime + + return ( + sha256(Path(runtime.__file__).read_bytes()).hexdigest() + == "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759" + ) + + +def _enable_multi_token_prefetch(runtime: HiSparseRuntime) -> None: + """Reuse Core follower staging without changing its KV write classification. + + Remove this adapter when Core exposes independent follower eligibility. + """ + original = runtime._resolve_and_stage_group + if getattr(original, "_ascend_multi_token_prefetch", False): + return + + @wraps(original) + def resolve_and_stage(*, resident, **kwargs): + original(resident=resident, **kwargs) + if ( + resident.decode_batch + or resident.num_decode_tokens <= 0 + or resident.num_decode_tokens != resident.num_actual_tokens + or any(handle.draft_layer for handle in resident.index_group_caches) + ): + return + group = runtime.index_group + with group.copy_stream: + for follower in group.followers: + follower._swap_rows(kwargs["shared_rows"]) + if follower._layer_ready_event is None: + follower._layer_ready_event = torch.Event() + follower._layer_ready_event.record(group.copy_stream) + follower._swap_staged = True + + resolve_and_stage._ascend_multi_token_prefetch = True + runtime._resolve_and_stage_group = resolve_and_stage + + class AscendHiSparseBackend(HiSparseBackend): # Keep private pools until NPU IPC events interoperate with Core's generic streams. stats_dtype = torch.int64 @@ -141,6 +188,17 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: torch.arange(workspace.max_rows + 1, dtype=torch.int32, device=runtime.device), torch.full((1,), attention_view.shape[0], dtype=torch.int32, device=runtime.device), ) + if runtime is group.leader: + config = get_current_vllm_config_or_none() + if ( + config is not None + and config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY + and config.speculative_config is not None + and config.speculative_config.method == "mtp" + and group.followers + and _has_follower_prefetch_contract() + ): + _enable_multi_token_prefetch(runtime) return HiSparseNPUOps( mapping, workspace, From 0a9b0e407c19d6218a02f3ce23292dc1eabec5d2 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 22:31:20 +0800 Subject: [PATCH 051/120] refactor(hisparse): remove unused scalar compact refill branch Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/cache_ops.py | 5 ++--- vllm_ascend/ops/triton/hisparse.py | 8 ++------ 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index aa613dc72290..aab56cc7fff6 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -364,9 +364,9 @@ def _gather( _check_tensor(attention, "attention", device, (torch.int32,), shape=source.shape) if attention_block < block_size: raise ValueError("HiSparse attention block stride must cover a physical block") - if rows == 0: + if rows == 0 or topk == 0: return - parallel_columns = min(topk, 128) if counts is not None and attention is None else 0 + parallel_columns = min(topk, 128) if counts is not None else 0 gather_rows_kernel[(rows, parallel_columns or topk)]( self.host_mapping.address_of(host_cache), hot_cache.view(torch.uint8), @@ -383,7 +383,6 @@ def _gather( HOT_STRIDE_BYTES=hot_stride, TOP_K=topk, ATTENTION_BLOCK=attention_block, - COMPACT=counts is not None, HAS_ATTENTION=attention is not None, HAS_STATE=state is not None, BLOCK=min(1024, triton.next_power_of_2(row_bytes)), diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 88901a903131..30c1a659eaa4 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -25,7 +25,6 @@ def gather_rows_kernel( HOT_STRIDE_BYTES: tl.constexpr, TOP_K: tl.constexpr, ATTENTION_BLOCK: tl.constexpr, - COMPACT: tl.constexpr, HAS_ATTENTION: tl.constexpr, HAS_STATE: tl.constexpr, BLOCK: tl.constexpr, @@ -35,7 +34,7 @@ def gather_rows_kernel( offset = tl.arange(0, BLOCK) padding = tl.load(state + row) < 0 if HAS_STATE else False if PARALLEL_COLUMNS > 0: - tl.static_assert(COMPACT and not HAS_ATTENTION) + tl.static_assert(not HAS_ATTENTION) count = tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) # Empty compact plans do not inspect stale source/destination entries. # Keep enough independent programs for small misses and retain the @@ -63,10 +62,7 @@ def gather_rows_kernel( if HAS_ATTENTION: value = tl.where((dst >= 0) & ~padding, dst // HOT_BLOCK * ATTENTION_BLOCK + dst % HOT_BLOCK, -1) tl.store(attention + index, value) - if COMPACT: - copy = col < tl.minimum(TOP_K, tl.maximum(tl.load(counts + row), 0)) - else: - copy = tl.load(mask_ptr + index) != 0 + copy = tl.load(mask_ptr + index) != 0 if copy & ~padding: src = tl.load(source + index) valid = (src >= 0) & (dst >= 0) & (dst < HOT_ROWS) From 51e61b59e90db904ea41c946b437cd99686d6f11 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 22:37:02 +0800 Subject: [PATCH 052/120] perf(hisparse): prefetch every shared compact swap plan Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 9618f382cb04..4b01bcece08a 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -11,8 +11,6 @@ from typing import TYPE_CHECKING, Any import torch -from vllm.config import get_current_vllm_config_or_none -from vllm.config.compilation import CUDAGraphMode from vllm.v1.hisparse.backend import HiSparseBackend if TYPE_CHECKING: @@ -64,7 +62,7 @@ def _has_follower_prefetch_contract() -> bool: def _enable_multi_token_prefetch(runtime: HiSparseRuntime) -> None: - """Reuse Core follower staging without changing its KV write classification. + """Stage followers after Core publishes their shared compact swap plan. Remove this adapter when Core exposes independent follower eligibility. """ @@ -75,12 +73,9 @@ def _enable_multi_token_prefetch(runtime: HiSparseRuntime) -> None: @wraps(original) def resolve_and_stage(*, resident, **kwargs): original(resident=resident, **kwargs) - if ( - resident.decode_batch - or resident.num_decode_tokens <= 0 - or resident.num_decode_tokens != resident.num_actual_tokens - or any(handle.draft_layer for handle in resident.index_group_caches) - ): + # Core already stages the followers for this classification. Keep + # it unchanged because it also determines the number of KV writes. + if resident.decode_batch: return group = runtime.index_group with group.copy_stream: @@ -188,17 +183,8 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: torch.arange(workspace.max_rows + 1, dtype=torch.int32, device=runtime.device), torch.full((1,), attention_view.shape[0], dtype=torch.int32, device=runtime.device), ) - if runtime is group.leader: - config = get_current_vllm_config_or_none() - if ( - config is not None - and config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY - and config.speculative_config is not None - and config.speculative_config.method == "mtp" - and group.followers - and _has_follower_prefetch_contract() - ): - _enable_multi_token_prefetch(runtime) + if runtime is group.leader and getattr(group, "followers", ()) and _has_follower_prefetch_contract(): + _enable_multi_token_prefetch(runtime) return HiSparseNPUOps( mapping, workspace, From d3096a1d55a5637a1b87d04cb58d0fe2a24f3255 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 22:49:35 +0800 Subject: [PATCH 053/120] perf(hisparse): overlap prefill Host misses across shared-indexer layers Signed-off-by: zhenwenqi2024 (cherry picked from commit 226dcef52499f040a5bcd05a26bee3372f0bc476) Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 44 ++++---- vllm_ascend/hisparse/cache_ops.py | 33 ++++++ vllm_ascend/hisparse/prefill.py | 155 ++++++++++++++++++++++++++ vllm_ascend/ops/triton/hisparse.py | 26 +++++ 4 files changed, 238 insertions(+), 20 deletions(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index ae05f54a4916..742630321d8a 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -15,7 +15,11 @@ from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder from vllm_ascend.attention.utils import PreprocessType -from vllm_ascend.hisparse.prefill import AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan +from vllm_ascend.hisparse.prefill import ( + AscendHiSparsePrefillStagingPlan, + build_hisparse_prefill_staging_plan, + stage_prefill_rows, +) from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache # Bound the temporary D2D packing workspace independently of prefill length. @@ -140,6 +144,7 @@ def __init__(self, *args, **kwargs): if not isinstance(self.index_group, HiSparseMLAIndexGroup): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") self.hisparse_cache = self.index_group.cache(self.index_group_index) + self.hisparse_cache._ascend_prefill_layer_name = self.layer_name self._direct_prolog_slots: torch.Tensor | None = None self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: @@ -570,26 +575,25 @@ def _execute_sparse_flash_attention_process( ) ) if metadata.num_prefills: - host_cache = self.hisparse_cache.runtime.host_cache - staged, staged_table, req_ids = self.index_group.stage_prefill_rows( - self.index_group_index, - host_cache.view(-1, metadata.block_size, host_cache.shape[-1]), - metadata, - ) - selected = logical[num_decode:] - block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size - request_tables = staged_table[req_ids.to(torch.int64)] - physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size - physical = physical + selected.clamp_min(0) % metadata.block_size - physical = physical.masked_fill(selected < 0, -1).to(torch.int32) - outputs.append( - self._attention_physical( - ql_nope[num_decode:], - q_pe[num_decode:] if q_pe is not None else None, - staged, - physical, + staged, staged_table, req_ids, slot = stage_prefill_rows(self.index_group, self.index_group_index, metadata) + try: + selected = logical[num_decode:] + block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size + request_tables = staged_table[req_ids.to(torch.int64)] + physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size + physical = physical + selected.clamp_min(0) % metadata.block_size + physical = physical.masked_fill(selected < 0, -1).to(torch.int32) + outputs.append( + self._attention_physical( + ql_nope[num_decode:], + q_pe[num_decode:] if q_pe is not None else None, + staged, + physical, + ) ) - ) + finally: + if slot is not None: + slot.finish() if not outputs: return torch.zeros_like(ql_nope) result = torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index aab56cc7fff6..7d9c75fcb221 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -17,6 +17,7 @@ from vllm_ascend.ops.triton.hisparse import ( gather_rows_kernel, + gather_valid_resident_rows_kernel, invalidate_kernel, pack_attention_rows_kernel, scatter_packed_kv_kernel, @@ -436,6 +437,38 @@ def hisparse_invalidate_written_slots( ) +def gather_valid_resident_rows(source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor) -> None: + """Fill only Resident rows in a staging cache that already holds Host rows.""" + _check_packed_cache(source) + _check_packed_cache(destination) + rows = destination.size(0) * destination.size(1) + _check_tensor(row_ids, "row_ids", source.device, (torch.int32,), shape=(rows,)) + if ( + source.device != destination.device + or source.dtype != destination.dtype + or source.size(2) != destination.size(2) + or source.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr() + ): + raise ValueError("HiSparse Resident gather requires matching, separate source and staging rows") + row_bytes = source.size(2) * source.element_size() + block_rows = 8 + gather_valid_resident_rows_kernel[(triton.cdiv(rows, block_rows),)]( + source.view(torch.uint8), + destination.view(torch.uint8), + row_ids, + ROWS=rows, + SOURCE_ROWS=source.size(0) * source.size(1), + SOURCE_BLOCK=source.size(1), + SOURCE_STRIDE=source.stride(0) * source.element_size(), + DEST_BLOCK=destination.size(1), + DEST_STRIDE=destination.stride(0) * destination.element_size(), + ROW_BYTES=row_bytes, + BLOCK_ROWS=block_rows, + BLOCK_BYTES=triton.next_power_of_2(row_bytes), + multibuffer=False, + ) + + def pack_attention_rows(packed_cache, physical_topk, selected_latent_out, selected_rope_out, latent_dim) -> None: """Gather packed physical KV rows into the split, dense SFA input layout.""" _check_packed_cache(packed_cache) diff --git a/vllm_ascend/hisparse/prefill.py b/vllm_ascend/hisparse/prefill.py index 0bf43a4b4808..35a997954bd7 100644 --- a/vllm_ascend/hisparse/prefill.py +++ b/vllm_ascend/hisparse/prefill.py @@ -2,11 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Reuse Core's prefill plan with NPU-safe column selection.""" +from __future__ import annotations + +from dataclasses import dataclass, field + import torch from torch.overrides import TorchFunctionMode +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import get_forward_context from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan +from vllm_ascend.hisparse.cache_ops import gather_valid_resident_rows + +_PREFILL_LOOKAHEAD = "ascend_hisparse_prefill_lookahead" + class _PrefillColumnSelection(TorchFunctionMode): """Keep a local workaround for the NPU singleton-column indexing defect. @@ -41,6 +51,151 @@ def ensure_gpu_sources(self, resident_state_rows, state_indices, resident_block_ super().ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size) +@dataclass(eq=False) +class _StagingSlot: + tensor: torch.Tensor + available: torch.npu.Event + in_use: bool = False + + def finish(self) -> None: + # This event follows the last Attention read on the consumer stream. + self.available = torch.npu.Event() + self.available.record() + self.in_use = False + + +@dataclass +class _PrefillLookahead: + slots: list[_StagingSlot] = field(default_factory=list) + pending: tuple | None = None + + +def _host_view(runtime, block_size): + cache = runtime.host_cache + return cache.view(-1, block_size, cache.shape[-1]) + + +def _stage_next(state, cache, plan) -> None: + runtime = cache.runtime + host = _host_view(runtime, plan.block_size) + rows = plan.row_ids.size(1) + if rows == 0: + return + shape = (rows // plan.block_size, plan.block_size, host.size(2)) + slot = next( + ( + slot + for slot in state.slots + if not slot.in_use + and slot.tensor.shape == shape + and slot.tensor.dtype == host.dtype + and slot.tensor.device == plan.block_table.device + ), + None, + ) + if slot is None and len(state.slots) == 2: + return + stream = runtime.index_group.copy_stream + side = torch.npu.Stream( + stream_id=stream.stream_id, device_index=stream.device_index, device_type=stream.device_type + ) + # The scheduler's prior Host publication and this plan's device writes + # precede this boundary on the compute stream. + side.wait_stream(torch.npu.current_stream()) + if slot is not None: + side.wait_event(slot.available) + # A failed enqueue must not make the possibly-written slot reusable. + slot.in_use = True + ready = torch.npu.Event() + with torch.npu.stream(side): + for tensor in (plan.row_ids, plan.dst_rows, plan.miss_mask): + tensor.record_stream(side) + if slot is None: + staged = runtime.gather_prefill_cache(host, plan, resident_cache=None) + else: + slot.tensor.record_stream(side) + runtime.cache_ops.hisparse_gather_plan( + host.view(-1, host.size(2)), + slot.tensor, + plan.row_ids, + plan.dst_rows, + plan.miss_mask, + None, + None, + 0, + ) + ready.record(side) + if slot is None: + slot = _StagingSlot(staged, ready) + state.slots.append(slot) + else: + slot.available = ready + slot.in_use = False + state.pending = (cache, plan, plan.gpu_row_ids, plan.miss_mask, slot) + + +def stage_prefill_rows(group, layer_index, metadata): + """Pipeline historical Host misses for one same-indexer follower.""" + cache = group.cache(layer_index) + plan = metadata.prefill.host_staging_plan + host = _host_view(cache.runtime, plan.block_size) + if cache.all_context_pages_resident or metadata.hisparse_for_capture: + return (*group.stage_prefill_rows(layer_index, host, metadata), None) + context = get_forward_context() + if ( + context.cudagraph_runtime_mode != CUDAGraphMode.NONE + or not isinstance(context.attn_metadata, dict) + or torch.npu.is_current_stream_capturing() + ): + return (*group.stage_prefill_rows(layer_index, host, metadata), None) + + state = context.additional_kwargs.get(_PREFILL_LOOKAHEAD) + pending = None + if state is not None: + pending, state.pending = state.pending, None + slot = None + if pending is not None: + owner, owner_plan, gpu_rows, miss_mask, prefetched = pending + if owner is cache and owner_plan is plan and gpu_rows is plan.gpu_row_ids and miss_mask is plan.miss_mask: + slot = prefetched + if slot is not None: + slot.in_use = True + main = torch.npu.current_stream() + main.wait_event(slot.available) + slot.tensor.record_stream(main) + # The follower has now written its current KV. Patch its Resident + # rows without overwriting the prefetched Host bytes. + gather_valid_resident_rows(cache.view.cache, slot.tensor, plan.gpu_row_ids[0]) + staged, table = slot.tensor, plan.block_table + req_ids = metadata.req_id_per_token[metadata.num_decode_tokens :] + if metadata.num_decodes: + req_ids = req_ids - metadata.num_decodes + else: + staged, table, req_ids = group.stage_prefill_rows(layer_index, host, metadata) + + next_index = layer_index + 1 + if next_index < len(group.caches) and plan.gpu_row_ids is not None and cache.view is not None: + follower = group.cache(next_index) + if ( + context.attn_metadata.get(getattr(follower, "_ascend_prefill_layer_name", None)) is metadata + and follower.view is not None + and follower.block_table is not None + and cache.block_table is not None + and follower.block_table.data_ptr() == cache.block_table.data_ptr() + and follower.view.block_size == cache.view.block_size + and plan.gpu_source_key == (cache.block_table.data_ptr(), cache.view.block_size) + and cache.source_block_table is not None + and follower.source_block_table is cache.source_block_table + and follower.runtime.request_state_indices is cache.runtime.request_state_indices + and follower.runtime.index_group is cache.runtime.index_group + ): + if state is None: + state = _PrefillLookahead() + context.additional_kwargs[_PREFILL_LOOKAHEAD] = state + _stage_next(state, follower, plan) + return staged, table, req_ids, slot + + def build_hisparse_prefill_staging_plan( block_table: torch.Tensor, seq_lens: torch.Tensor, diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 30c1a659eaa4..74179ee6893e 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -158,3 +158,29 @@ def scatter_packed_kv_kernel( b = tl.load(rope + row.to(tl.int64) * ROPE_STRIDE + i - LATENT, mask=valid & (i >= LATENT) & (i < WIDTH), other=0) base = slot.to(tl.int64) // CACHE_BLOCK * CACHE_STRIDE + slot % CACHE_BLOCK * WIDTH tl.store(cache + base + i, tl.where(i < LATENT, a, b), mask=valid & (i < WIDTH)) + + +@triton.jit +def gather_valid_resident_rows_kernel( + source, + destination, + row_ids, + ROWS: tl.constexpr, + SOURCE_ROWS: tl.constexpr, + SOURCE_BLOCK: tl.constexpr, + SOURCE_STRIDE: tl.constexpr, + DEST_BLOCK: tl.constexpr, + DEST_STRIDE: tl.constexpr, + ROW_BYTES: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + BLOCK_BYTES: tl.constexpr, +): + row = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64) + src = tl.load(row_ids + row, mask=row < ROWS, other=-1).to(tl.int64) + src_offset = src // SOURCE_BLOCK * SOURCE_STRIDE + src % SOURCE_BLOCK * ROW_BYTES + dst_offset = row // DEST_BLOCK * DEST_STRIDE + row % DEST_BLOCK * ROW_BYTES + byte = tl.arange(0, BLOCK_BYTES) + valid = (row < ROWS) & (src >= 0) & (src < SOURCE_ROWS) + mask = valid[:, None] & (byte[None, :] < ROW_BYTES) + payload = tl.load(source + src_offset[:, None] + byte[None, :], mask=mask, other=0) + tl.store(destination + dst_offset[:, None] + byte[None, :], payload, mask=mask) From a20218c7f218d97321de0838094dacdf585dcd77 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 23:02:27 +0800 Subject: [PATCH 054/120] fix(hisparse): bound Prefill Resident gather byte tile Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/cache_ops.py | 6 ++++-- vllm_ascend/ops/triton/hisparse.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 7d9c75fcb221..0bf9040d3c0c 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -452,7 +452,9 @@ def gather_valid_resident_rows(source: torch.Tensor, destination: torch.Tensor, raise ValueError("HiSparse Resident gather requires matching, separate source and staging rows") row_bytes = source.size(2) * source.element_size() block_rows = 8 - gather_valid_resident_rows_kernel[(triton.cdiv(rows, block_rows),)]( + block_bytes = min(1024, triton.next_power_of_2(row_bytes)) + grid = (triton.cdiv(rows, block_rows), triton.cdiv(row_bytes, block_bytes)) + gather_valid_resident_rows_kernel[grid]( source.view(torch.uint8), destination.view(torch.uint8), row_ids, @@ -464,7 +466,7 @@ def gather_valid_resident_rows(source: torch.Tensor, destination: torch.Tensor, DEST_STRIDE=destination.stride(0) * destination.element_size(), ROW_BYTES=row_bytes, BLOCK_ROWS=block_rows, - BLOCK_BYTES=triton.next_power_of_2(row_bytes), + BLOCK_BYTES=block_bytes, multibuffer=False, ) diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 74179ee6893e..c2a41e51f3e6 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -179,7 +179,7 @@ def gather_valid_resident_rows_kernel( src = tl.load(row_ids + row, mask=row < ROWS, other=-1).to(tl.int64) src_offset = src // SOURCE_BLOCK * SOURCE_STRIDE + src % SOURCE_BLOCK * ROW_BYTES dst_offset = row // DEST_BLOCK * DEST_STRIDE + row % DEST_BLOCK * ROW_BYTES - byte = tl.arange(0, BLOCK_BYTES) + byte = tl.program_id(1) * BLOCK_BYTES + tl.arange(0, BLOCK_BYTES) valid = (row < ROWS) & (src >= 0) & (src < SOURCE_ROWS) mask = valid[:, None] & (byte[None, :] < ROW_BYTES) payload = tl.load(source + src_offset[:, None] + byte[None, :], mask=mask, other=0) From 4c0ef30c18332fd0dfc417771b44f89b0460d7a0 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 20:46:21 +0800 Subject: [PATCH 055/120] perf(hisparse): group recycled Host block invalidation (cherry picked from commit 10145fa06c9837bac004af4fcfaa6e9278270a62) Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/recycled_invalidate.py | 181 ++++++++++++++++++++ vllm_ascend/hisparse/row_dma.py | 32 +++- 2 files changed, 208 insertions(+), 5 deletions(-) create mode 100644 vllm_ascend/hisparse/recycled_invalidate.py diff --git a/vllm_ascend/hisparse/recycled_invalidate.py b/vllm_ascend/hisparse/recycled_invalidate.py new file mode 100644 index 000000000000..e2221c90db57 --- /dev/null +++ b/vllm_ascend/hisparse/recycled_invalidate.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Grouped invalidation for recycled Host blocks in the pinned HiSparse worker.""" + +import operator + +import torch +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d + +from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + +logger = init_logger(__name__) +_BLOCK_BUCKET = 8 +_TILE = 256 +_INT32_MAX = 2**31 - 1 +_JIT_POINTER_ALIGNMENT = 16 + + +def _bucket_blocks(block_ids, block_size): + try: + blocks = sorted({operator.index(block) for block in block_ids}) + except TypeError: + return None + if ( + not blocks + or len(blocks) > _BLOCK_BUCKET + or blocks[0] < 0 + or blocks[-1] * block_size + block_size - 1 > _INT32_MAX + ): + return None + # Padding with the last valid ID preserves membership at the int32 limit. + return blocks + [blocks[-1]] * (_BLOCK_BUCKET - len(blocks)) + + +@triton.jit +def _invalidate_recycled_groups( + tables, + blocks, + state_indices, + STATES: tl.constexpr, + HOT: tl.constexpr, + BLOCK_COUNT: tl.constexpr, + PAGE_ROWS: tl.constexpr, + ITERATIONS: tl.constexpr, + TILE: tl.constexpr, +): + table = tl.load(tables + tl.program_id(0)).to(tl.pointer_type(tl.int32)) + state = tl.load(state_indices + tl.program_id(1)) + slot = tl.program_id(2) * TILE + tl.arange(0, TILE) + valid = (state >= 0) & (state < STATES) & (slot < HOT) + value = tl.load(table + state * HOT + slot, mask=valid, other=-1) + valid = valid & (value >= 0) + page = tl.maximum(value, 0) // PAGE_ROWS + first = tl.load(blocks) + last = tl.load(blocks + BLOCK_COUNT - 1) + valid = valid & (page >= first) & (page <= last) + low = tl.full((TILE,), 0, tl.int32) + high = tl.full((TILE,), BLOCK_COUNT, tl.int32) + for _ in tl.static_range(ITERATIONS): + active = valid & (low < high) + middle = low + (high - low) // 2 + key = tl.load(blocks + middle, mask=active, other=0) + advance = active & (key < page) + high = tl.where(active & ~advance, middle, high) + low = tl.where(advance, middle + 1, low) + key = tl.load(blocks + low, mask=valid & (low < BLOCK_COUNT), other=0) + tl.store(table + state * HOT + slot, -1, mask=valid & (low < BLOCK_COUNT) & (key == page)) + + +class RecycledGroupInvalidator: + def __init__(self, tables, block_size): + self._rank = get_tensor_model_parallel_rank() + self._tables = tables + self._storages = tuple(table.untyped_storage() for table in tables) + self._addresses = tuple(table.data_ptr() for table in tables) + self._device = tables[0].device + self._states, self._hot = tables[0].shape + self._block_size = block_size + self._stream = torch.accelerator.current_stream() + self._pointer_table = async_tensor_h2d(list(self._addresses), self._device, torch.int64) + self._logged_rows: set[int] = set() + + @classmethod + def bind(cls, worker): + kernel_config = getattr(worker.vllm_config, "kernel_config", None) + if ( + not getattr(kernel_config, "enable_jit_warmup", True) + or get_ascend_device_type() != AscendDeviceType.A5 + or worker.backend.is_current_stream_capturing() + or worker.kernel_block_size <= 0 + or not worker.leader_runtimes + ): + return None + tables = tuple(runtime.index_group.device_global_indices for runtime in worker.leader_runtimes) + first = tables[0] + if first.device.type != "npu" or first.ndim != 2 or not first.numel(): + return None + if any( + table.device != first.device + or table.dtype != torch.int32 + or table.shape != first.shape + or not table.is_contiguous() + for table in tables + ): + return None + ranges = sorted((table.data_ptr(), table.data_ptr() + table.numel() * table.element_size()) for table in tables) + if any(left[1] > right[0] for left, right in zip(ranges, ranges[1:])): + return None + owner = cls(tables, worker.kernel_block_size) + # Compile the live ABI before serving. The device state sentinel masks + # every table read/write while preserving the real dtype and geometry. + dummy_states = async_tensor_h2d([-1], first.device, torch.int64) + if dummy_states.data_ptr() % _JIT_POINTER_ALIGNMENT: + return None + owner._launch([0] * _BLOCK_BUCKET, dummy_states) + logger.info( + "HiSparse recycled invalidation prewarmed: rank=%d groups=%d state_dtype=int64 bucket=8", + owner._rank, + len(tables), + ) + return owner + + def _launch(self, blocks, indices): + device_blocks = async_tensor_h2d(blocks, self._device, torch.int32) + grid = (len(self._tables), indices.numel(), triton.cdiv(self._hot, _TILE)) + _invalidate_recycled_groups[grid]( + self._pointer_table, + device_blocks, + indices, + self._states, + self._hot, + _BLOCK_BUCKET, + self._block_size, + _BLOCK_BUCKET.bit_length(), + _TILE, + ) + + def try_call(self, worker, block_ids, indices) -> bool: + # The pinned caller passes unique, live scheduler indices, not the + # resolver's padded state table. External callers retain Core errors. + if ( + block_ids is not worker._pending_invalid_block_ids + or worker.kernel_block_size != self._block_size + or worker.backend.is_current_stream_capturing() + or torch.accelerator.current_stream() != self._stream + or indices.ndim != 1 + or indices.dtype != torch.int64 + or indices.device != self._device + or not indices.is_contiguous() + or indices.storage_offset() != 0 + or not 0 < indices.numel() <= self._states + or indices.data_ptr() % _JIT_POINTER_ALIGNMENT + or len(worker.leader_runtimes) != len(self._tables) + ): + return False + for runtime, table, address in zip(worker.leader_runtimes, self._tables, self._addresses, strict=True): + current = runtime.index_group.device_global_indices + if ( + current is not table + or current.data_ptr() != address + or current.shape != (self._states, self._hot) + or not current.is_contiguous() + ): + return False + blocks = _bucket_blocks(block_ids, self._block_size) + if blocks is None: + return False + self._launch(blocks, indices) + rows = indices.numel() + if rows not in self._logged_rows: + logger.info( + "HiSparse grouped recycled invalidation: rank=%d groups=%d active_rows=%d bucket=8", + self._rank, + len(self._tables), + rows, + ) + self._logged_rows.add(rows) + return True diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index f7523c41c9c9..153185d1984b 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -10,6 +10,7 @@ import numpy as np _WORKER_SHA256 = "864e2c7ce4e2cff320b4f68f8dccc9691adb829f0e50fdf6a27dc9c7885040e5" +_RUNTIME_SHA256 = "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759" # Large prefill batches keep the scalar path to bound temporary memory. _MAX_DESCRIPTORS = 32768 @@ -108,7 +109,7 @@ def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: return True -def _install_worker(worker_type, backend_type) -> None: +def _install_worker(worker_type, backend_type, bind_recycled=None) -> None: original_init = worker_type._init_dma original_enqueue = worker_type._enqueue_row_dma if getattr(original_enqueue, "_ascend_row_dma", False): @@ -118,8 +119,11 @@ def _install_worker(worker_type, backend_type) -> None: def init(worker): original_init(worker) worker._ascend_row_dma_geometry = None - if isinstance(worker.backend, backend_type) and worker.is_host_writer and worker.shared_host_region is None: - worker._ascend_row_dma_geometry = _Geometry.bind(worker) + if isinstance(worker.backend, backend_type): + if worker.is_host_writer and worker.shared_host_region is None: + worker._ascend_row_dma_geometry = _Geometry.bind(worker) + if bind_recycled is not None: + worker._ascend_recycled_invalidator = bind_recycled(worker) @wraps(original_enqueue) def enqueue(worker, layer_indices, ready_event=None): @@ -130,11 +134,29 @@ def enqueue(worker, layer_indices, ready_event=None): enqueue._ascend_row_dma = True worker_type._init_dma = init worker_type._enqueue_row_dma = enqueue + if bind_recycled is not None: + # Keep the same pinned adapter until Core exposes this backend hook. + original_invalidate = worker_type.invalidate_blocks + + @wraps(original_invalidate) + def invalidate(worker, block_ids, request_state_indices): + owner = getattr(worker, "_ascend_recycled_invalidator", None) + if owner is None or not owner.try_call(worker, block_ids, request_state_indices): + return original_invalidate(worker, block_ids, request_state_indices) + + worker_type.invalidate_blocks = invalidate def install(backend_type) -> None: """Only adapt the exact Core implementation reviewed for this experiment.""" from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker + from vllm.v1.hisparse import runtime + + if sha256(Path(worker.__file__).read_bytes()).hexdigest() != _WORKER_SHA256: + return + bind_recycled = None + if sha256(Path(runtime.__file__).read_bytes()).hexdigest() == _RUNTIME_SHA256: + from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator - if sha256(Path(worker.__file__).read_bytes()).hexdigest() == _WORKER_SHA256: - _install_worker(worker.HiSparseConnectorWorker, backend_type) + bind_recycled = RecycledGroupInvalidator.bind + _install_worker(worker.HiSparseConnectorWorker, backend_type, bind_recycled) From 1e50646357cc073d8a1bc8f6d04ac834636b630e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Fri, 6 Nov 2026 23:20:16 +0800 Subject: [PATCH 056/120] perf(hisparse): group all valid recycled Host blocks Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/recycled_invalidate.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/vllm_ascend/hisparse/recycled_invalidate.py b/vllm_ascend/hisparse/recycled_invalidate.py index e2221c90db57..2d7d1a783def 100644 --- a/vllm_ascend/hisparse/recycled_invalidate.py +++ b/vllm_ascend/hisparse/recycled_invalidate.py @@ -19,20 +19,14 @@ _JIT_POINTER_ALIGNMENT = 16 -def _bucket_blocks(block_ids, block_size): +def _normalize_blocks(block_ids, block_size): try: blocks = sorted({operator.index(block) for block in block_ids}) except TypeError: return None - if ( - not blocks - or len(blocks) > _BLOCK_BUCKET - or blocks[0] < 0 - or blocks[-1] * block_size + block_size - 1 > _INT32_MAX - ): + if not blocks or blocks[0] < 0 or blocks[-1] * block_size + block_size - 1 > _INT32_MAX: return None - # Padding with the last valid ID preserves membership at the int32 limit. - return blocks + [blocks[-1]] * (_BLOCK_BUCKET - len(blocks)) + return blocks @triton.jit @@ -165,10 +159,14 @@ def try_call(self, worker, block_ids, indices) -> bool: or not current.is_contiguous() ): return False - blocks = _bucket_blocks(block_ids, self._block_size) + blocks = _normalize_blocks(block_ids, self._block_size) if blocks is None: return False - self._launch(blocks, indices) + for start in range(0, len(blocks), _BLOCK_BUCKET): + chunk = blocks[start : start + _BLOCK_BUCKET] + # Repeating the last ID preserves membership without a sentinel. + chunk.extend([chunk[-1]] * (_BLOCK_BUCKET - len(chunk))) + self._launch(chunk, indices) rows = indices.numel() if rows not in self._logged_rows: logger.info( From 0fcbdd5457785a13eb6fb23e735191914a229790 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 01:01:22 +0800 Subject: [PATCH 057/120] refactor(hisparse): reuse Core speculative group prefetch Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 44 --------------------------------- vllm_ascend/hisparse/row_dma.py | 7 ++++-- 2 files changed, 5 insertions(+), 46 deletions(-) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 4b01bcece08a..860257525289 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -5,9 +5,6 @@ from __future__ import annotations import importlib -from functools import cache, wraps -from hashlib import sha256 -from pathlib import Path from typing import TYPE_CHECKING, Any import torch @@ -51,45 +48,6 @@ def capture_aware_stream(stream: torch.Stream) -> torch.Stream: return _CaptureAwareStream(stream.stream_id, stream.device_index, stream.device_type) -@cache -def _has_follower_prefetch_contract() -> bool: - from vllm.v1.hisparse import runtime - - return ( - sha256(Path(runtime.__file__).read_bytes()).hexdigest() - == "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759" - ) - - -def _enable_multi_token_prefetch(runtime: HiSparseRuntime) -> None: - """Stage followers after Core publishes their shared compact swap plan. - - Remove this adapter when Core exposes independent follower eligibility. - """ - original = runtime._resolve_and_stage_group - if getattr(original, "_ascend_multi_token_prefetch", False): - return - - @wraps(original) - def resolve_and_stage(*, resident, **kwargs): - original(resident=resident, **kwargs) - # Core already stages the followers for this classification. Keep - # it unchanged because it also determines the number of KV writes. - if resident.decode_batch: - return - group = runtime.index_group - with group.copy_stream: - for follower in group.followers: - follower._swap_rows(kwargs["shared_rows"]) - if follower._layer_ready_event is None: - follower._layer_ready_event = torch.Event() - follower._layer_ready_event.record(group.copy_stream) - follower._swap_staged = True - - resolve_and_stage._ascend_multi_token_prefetch = True - runtime._resolve_and_stage_group = resolve_and_stage - - class AscendHiSparseBackend(HiSparseBackend): # Keep private pools until NPU IPC events interoperate with Core's generic streams. stats_dtype = torch.int64 @@ -183,8 +141,6 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: torch.arange(workspace.max_rows + 1, dtype=torch.int32, device=runtime.device), torch.full((1,), attention_view.shape[0], dtype=torch.int32, device=runtime.device), ) - if runtime is group.leader and getattr(group, "followers", ()) and _has_follower_prefetch_contract(): - _enable_multi_token_prefetch(runtime) return HiSparseNPUOps( mapping, workspace, diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index 153185d1984b..812ed79beb8e 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -10,7 +10,10 @@ import numpy as np _WORKER_SHA256 = "864e2c7ce4e2cff320b4f68f8dccc9691adb829f0e50fdf6a27dc9c7885040e5" -_RUNTIME_SHA256 = "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759" +_RECYCLED_RUNTIME_SHA256S = { + "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759", # Core 24420b7a + "935ad5669ccdcbc0ffa45d66cfd28145cee0039fe979bb1c637f325373133325", # Core 39dbf518 +} # Large prefill batches keep the scalar path to bound temporary memory. _MAX_DESCRIPTORS = 32768 @@ -155,7 +158,7 @@ def install(backend_type) -> None: if sha256(Path(worker.__file__).read_bytes()).hexdigest() != _WORKER_SHA256: return bind_recycled = None - if sha256(Path(runtime.__file__).read_bytes()).hexdigest() == _RUNTIME_SHA256: + if sha256(Path(runtime.__file__).read_bytes()).hexdigest() in _RECYCLED_RUNTIME_SHA256S: from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator bind_recycled = RecycledGroupInvalidator.bind From b494a3ce810368813577445c97816b9bf31498ea Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 01:57:08 +0800 Subject: [PATCH 058/120] fix(mla): accept Core no-argument post-load hook Signed-off-by: zhenwenqi2024 --- vllm_ascend/ops/mla.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index 03707a0d7865..86d9ad15a962 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -251,17 +251,18 @@ def __init__( original_process_weights = self.mla_attn.process_weights_after_loading - def wrapped_process_weights(act_dtype: torch.dtype): + def wrapped_process_weights(act_dtype: torch.dtype | None = None): from vllm_ascend.attention.sfa_v1 import AscendSFAImpl - if self.dcp_q_replicate_enabled: - self.mla_attn.impl.process_weights_after_loading(act_dtype) - elif not isinstance(self.mla_attn.impl, AscendSFAImpl): - # Both supported vLLM versions dispatch to the impl here. - original_process_weights(act_dtype) - else: + if self.dcp_q_replicate_enabled or isinstance(self.mla_attn.impl, AscendSFAImpl): # SFA disposes kv_b_proj, so bypass upstream's dense packing. - self.mla_attn.impl.process_weights_after_loading(act_dtype) + self.mla_attn.impl.process_weights_after_loading( + self.mla_attn.act_dtype if act_dtype is None else act_dtype + ) + elif act_dtype is None: + original_process_weights() + else: + original_process_weights(act_dtype) self.mla_attn.process_weights_after_loading = wrapped_process_weights From f6393dd7de64ab930f9f2e9c4af2a3861495486d Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 02:13:02 +0800 Subject: [PATCH 059/120] fix(worker): accept Core KV allocation callback ABI Signed-off-by: zhenwenqi2024 --- vllm_ascend/worker/v2/attn_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index eb11f7e080dc..38ba536d803a 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -1669,6 +1669,8 @@ def allocate_kv_cache_main( device: torch.device, layout: Any, kernel_block_sizes: list[int], + *, + allocate: Callable[[int], torch.Tensor] | None = None, ) -> dict[str, Any]: """Allocate Ascend KV cache through vLLM main's #51718 entry point. @@ -1677,6 +1679,8 @@ def allocate_kv_cache_main( still consumes separate K/V (and backend-specific state) tensors, so keep the Ascend allocation/reshape contract behind the new entry point. """ + if allocate is not None: + raise NotImplementedError("Extensible KV cache allocation is not supported on Ascend.") del layout vllm_config = get_current_vllm_config() attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase) From c67eb79a83ed0bcddd5709342306ea7cc456c6d9 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 02:34:22 +0800 Subject: [PATCH 060/120] perf(hisparse): read resident pages in eager C8 prefill Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 62 +++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 742630321d8a..e9f39999a452 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -7,6 +7,7 @@ import torch import torch_npu +from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import logger from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder @@ -524,6 +525,61 @@ def _attention_physical( outputs.append(result) return torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] + def _try_c8_resident_prefill(self, query, query_rope, logical_topk, metadata, query_ends, key_lengths): + if ( + not self.enable_sparse_sfa_c8 + or metadata.num_decodes + or metadata.hisparse_for_capture + or query_rope is None + or not (query.shape[0] == logical_topk.shape[0] == metadata.num_actual_tokens) + ): + return None + cache = self.hisparse_cache + view = cache.view + state_indices = cache.runtime.request_state_indices + if ( + not cache.all_context_pages_resident + or cache.dummy_batch + or cache.draft_layer + or view is None + or cache.block_table is None + or state_indices is None + or state_indices.numel() < metadata.num_reqs + or query_ends.shape != (metadata.num_reqs,) + or key_lengths.shape != (metadata.num_reqs,) + ): + return None + if ( + get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.NONE + or torch.npu.is_current_stream_capturing() + ): + return None + + # Read the exact layer's live Resident table after the connector updates + # it. The dense HMA alias includes gaps, so scale only its physical pages. + page_stride, remainder = divmod(view.attention_block_stride, view.block_size) + storage = view.attention_cache + if ( + remainder + or page_stride < 1 + or not storage.is_contiguous() + or view.cache.stride(0) != view.attention_block_stride * view.cache.stride(1) + or storage.size(0) <= (view.cache.size(0) - 1) * page_stride + ): + return None + table = cache.block_table.index_select(0, state_indices[: metadata.num_reqs].to(torch.long)) + table.clamp_min_(0).mul_(page_stride) + return super()._execute_sparse_flash_attention_process( + query, + query_rope, + (storage.unsqueeze(2),), + logical_topk.unsqueeze(1), + metadata, + query_ends, + key_lengths, + block_table=table, + ) + def _execute_sparse_flash_attention_process( self, ql_nope, @@ -545,6 +601,12 @@ def _execute_sparse_flash_attention_process( # Decode's resolver already masks graph padding from this same device # scalar. Keep that validation at the producer of physical indices. self.index_group.set_logical_topk_ready(self.index_group_index) + if metadata.num_prefills and not metadata.num_decodes: + resident = self._try_c8_resident_prefill( + ql_nope, q_pe, logical, metadata, actual_seq_lengths_query, actual_seq_lengths_key + ) + if resident is not None: + return resident outputs = [] num_decode = metadata.num_decode_tokens if num_decode: From bf872fc9e1bed2de465ae5c0ad735e6826d0ea40 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 02:34:22 +0800 Subject: [PATCH 061/120] perf(hisparse): simplify contiguous C8 row mapping Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/c8.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm_ascend/hisparse/c8.py b/vllm_ascend/hisparse/c8.py index d2fe993fe3aa..a93592a77442 100644 --- a/vllm_ascend/hisparse/c8.py +++ b/vllm_ascend/hisparse/c8.py @@ -67,6 +67,8 @@ def map_c8_physical_indices(indices: torch.Tensor, block: int, stride_rows: int, if capacity <= 0 or capacity % block or (capacity // block - 1) * stride_rows + block > _INT32_MAX: raise ValueError("HiSparse C8 physical row capacity is invalid") valid = (indices >= 0) & (indices < capacity) + if stride_rows == block: + return torch.where(valid, indices, -1) safe = indices.clamp_min(0).to(torch.int64) mapped = safe.div(block, rounding_mode="floor") * stride_rows + safe.remainder(block) return torch.where(valid, mapped, -1).to(torch.int32) From b5f75fa3047b858f18a0a5525c45bd13ce271d67 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 02:50:20 +0800 Subject: [PATCH 062/120] perf(hisparse): use resident pages for mixed-batch prefill Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 61 +++++++++++++++------------ 1 file changed, 33 insertions(+), 28 deletions(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index e9f39999a452..eb11ab82ea93 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -526,12 +526,16 @@ def _attention_physical( return torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] def _try_c8_resident_prefill(self, query, query_rope, logical_topk, metadata, query_ends, key_lengths): + num_decodes = metadata.num_decodes + num_decode_tokens = metadata.num_decode_tokens + num_prefill_tokens = metadata.num_actual_tokens - num_decode_tokens if ( not self.enable_sparse_sfa_c8 - or metadata.num_decodes or metadata.hisparse_for_capture or query_rope is None - or not (query.shape[0] == logical_topk.shape[0] == metadata.num_actual_tokens) + or num_prefill_tokens <= 0 + or metadata.num_reqs != num_decodes + metadata.num_prefills + or not (query.shape[0] == logical_topk.shape[0] == num_prefill_tokens) ): return None cache = self.hisparse_cache @@ -567,9 +571,12 @@ def _try_c8_resident_prefill(self, query, query_rope, logical_topk, metadata, qu or storage.size(0) <= (view.cache.size(0) - 1) * page_stride ): return None - table = cache.block_table.index_select(0, state_indices[: metadata.num_reqs].to(torch.long)) + table = cache.block_table.index_select(0, state_indices[num_decodes : metadata.num_reqs].to(torch.long)) table.clamp_min_(0).mul_(page_stride) - return super()._execute_sparse_flash_attention_process( + if num_decodes: + query_ends = query_ends[num_decodes:] - num_decode_tokens + key_lengths = key_lengths[num_decodes:] + output = super()._execute_sparse_flash_attention_process( query, query_rope, (storage.unsqueeze(2),), @@ -579,6 +586,8 @@ def _try_c8_resident_prefill(self, query, query_rope, logical_topk, metadata, qu key_lengths, block_table=table, ) + logger.info_once("[hisparse] Resident paged Prefill active: mixed=%s", bool(num_decodes)) + return output def _execute_sparse_flash_attention_process( self, @@ -601,12 +610,6 @@ def _execute_sparse_flash_attention_process( # Decode's resolver already masks graph padding from this same device # scalar. Keep that validation at the producer of physical indices. self.index_group.set_logical_topk_ready(self.index_group_index) - if metadata.num_prefills and not metadata.num_decodes: - resident = self._try_c8_resident_prefill( - ql_nope, q_pe, logical, metadata, actual_seq_lengths_query, actual_seq_lengths_key - ) - if resident is not None: - return resident outputs = [] num_decode = metadata.num_decode_tokens if num_decode: @@ -637,25 +640,27 @@ def _execute_sparse_flash_attention_process( ) ) if metadata.num_prefills: - staged, staged_table, req_ids, slot = stage_prefill_rows(self.index_group, self.index_group_index, metadata) - try: - selected = logical[num_decode:] - block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size - request_tables = staged_table[req_ids.to(torch.int64)] - physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size - physical = physical + selected.clamp_min(0) % metadata.block_size - physical = physical.masked_fill(selected < 0, -1).to(torch.int32) - outputs.append( - self._attention_physical( - ql_nope[num_decode:], - q_pe[num_decode:] if q_pe is not None else None, - staged, - physical, - ) + query = ql_nope[num_decode:] + query_rope = q_pe[num_decode:] if q_pe is not None else None + prefill_output = self._try_c8_resident_prefill( + query, query_rope, logical[num_decode:], metadata, actual_seq_lengths_query, actual_seq_lengths_key + ) + if prefill_output is None: + staged, staged_table, req_ids, slot = stage_prefill_rows( + self.index_group, self.index_group_index, metadata ) - finally: - if slot is not None: - slot.finish() + try: + selected = logical[num_decode:] + block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size + request_tables = staged_table[req_ids.to(torch.int64)] + physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size + physical = physical + selected.clamp_min(0) % metadata.block_size + physical = physical.masked_fill(selected < 0, -1).to(torch.int32) + prefill_output = self._attention_physical(query, query_rope, staged, physical) + finally: + if slot is not None: + slot.finish() + outputs.append(prefill_output) if not outputs: return torch.zeros_like(ql_nope) result = torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] From f5d96c845e962750c96b8b89be3a9fdec3e17406 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 03:42:18 +0800 Subject: [PATCH 063/120] test(hisparse): report resident prefill activation per worker Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index eb11ab82ea93..e0db162d827a 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -586,7 +586,7 @@ def _try_c8_resident_prefill(self, query, query_rope, logical_topk, metadata, qu key_lengths, block_table=table, ) - logger.info_once("[hisparse] Resident paged Prefill active: mixed=%s", bool(num_decodes)) + logger.info_once("[hisparse] Resident paged Prefill active: mixed=%s", bool(num_decodes), scope="process") return output def _execute_sparse_flash_attention_process( From dbdec2deac6d7558317a5fd8352d7f5f4cc0f58c Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 04:39:54 +0800 Subject: [PATCH 064/120] perf(hisparse): bound large written-slot invalidation Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/cache_ops.py | 42 +++++++++++++++++++++-- vllm_ascend/ops/triton/hisparse.py | 54 ++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 0bf9040d3c0c..ac68b399f89e 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -15,16 +15,20 @@ import torch from vllm.triton_utils import triton +from vllm_ascend.device.device_config import is_950 from vllm_ascend.ops.triton.hisparse import ( gather_rows_kernel, gather_valid_resident_rows_kernel, invalidate_kernel, + invalidate_written_slots_sorted, pack_attention_rows_kernel, + pack_written_slot_keys, scatter_packed_kv_kernel, ) _ACL_HOST_REGISTER_MAPPED = 0 _ACL_HAC_AIV = 3 +_SORTED_WRITTEN_SLOT_COUNT = 8192 class MappedHostTensor: @@ -419,11 +423,45 @@ def hisparse_invalidate_written_slots( _check_tensor(request_state_indices, "request_state_indices", device, (torch.int32,), ndim=1) _check_tensor(written_slots, "written_slots", device, (torch.int32, torch.int64), ndim=1) _check_tensor(req_id_per_token, "req_id_per_token", device, (torch.int32,), shape=written_slots.shape) - if written_slots.numel() == 0: + num_tokens = written_slots.numel() + if num_tokens == 0: return if region == 0 or state_rows == 0 or groups == 0: return - invalidate_kernel[(written_slots.numel(), triton.cdiv(region, 256), groups)]( + # The 8192-row Prefill chunk is the measured large shape. Keep scratch + # bounded and preserve the existing path for capture and other shapes. + if ( + grouped + and num_tokens == _SORTED_WRITTEN_SLOT_COUNT + and is_950() + and not torch.npu.is_current_stream_capturing() + ): + keys = torch.empty((num_tokens,), dtype=torch.int64, device=device) + sorted_keys = torch.empty_like(keys) + order = torch.empty_like(keys) + pack_written_slot_keys[(triton.cdiv(num_tokens, 256),)]( + request_state_indices, + req_id_per_token, + written_slots, + keys, + ROWS=num_tokens, + REQUESTS=request_state_indices.numel(), + STATES=state_rows, + TILE=256, + multibuffer=False, + ) + torch.sort(keys, out=(sorted_keys, order)) + invalidate_written_slots_sorted[(groups, state_rows, triton.cdiv(region, 256))]( + device_global_indices, + sorted_keys, + ROWS=num_tokens, + HOT=region, + TILE=256, + STEPS=num_tokens.bit_length(), + multibuffer=False, + ) + return + invalidate_kernel[(num_tokens, triton.cdiv(region, 256), groups)]( device_global_indices, request_state_indices, req_id_per_token, diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index c2a41e51f3e6..c0f36fecfc92 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -103,6 +103,60 @@ def invalidate_kernel( tl.store(ptr, -1, mask=valid & (cached == value)) +@triton.jit +def pack_written_slot_keys( + states, + requests, + written, + keys, + ROWS: tl.constexpr, + REQUESTS: tl.constexpr, + STATES: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0) * TILE + tl.arange(0, TILE) + request = tl.load(requests + row, mask=row < ROWS, other=-1) + slot = tl.load(written + row, mask=row < ROWS, other=-1) + valid = (request >= 0) & (request < REQUESTS) & (slot >= 0) & (slot <= 2147483647) + state = tl.load(states + request, mask=valid, other=-1) + valid = valid & (state >= 0) & (state < STATES) + # Each key identifies a Host slot within one request state. Slots above + # int32 cannot match a Hot entry and share the padding sentinel. + key = (state.to(tl.int64) << 31) | slot.to(tl.int64) + tl.store(keys + row, tl.where(valid, key, 9223372036854775807), mask=row < ROWS) + + +@triton.jit +def invalidate_written_slots_sorted( + tables, + keys, + ROWS: tl.constexpr, + HOT: tl.constexpr, + TILE: tl.constexpr, + STEPS: tl.constexpr, +): + table = tl.load(tables + tl.program_id(0)).to(tl.pointer_type(tl.int32)) + state = tl.program_id(1) + offset = tl.program_id(2) * TILE + tl.arange(0, TILE) + pointer = table + state.to(tl.int64) * HOT + offset + cached = tl.load(pointer, mask=offset < HOT, other=-1) + valid = (offset < HOT) & (cached >= 0) + if tl.sum(valid.to(tl.int32), 0) > 0: + query = (state.to(tl.int64) << 31) | cached.to(tl.int64) + low = tl.full((TILE,), 0, tl.int32) + high = tl.full((TILE,), ROWS, tl.int32) + for _ in tl.static_range(STEPS): + active = valid & (low < high) + mid = low + (high - low) // 2 + value = tl.load(keys + mid, mask=active, other=9223372036854775807) + advance = active & (value < query) + low = tl.where(advance, mid + 1, low) + high = tl.where(active & ~advance, mid, high) + inside = valid & (low < ROWS) + key = tl.load(keys + low, mask=inside, other=9223372036854775807) + tl.store(pointer, -1, mask=inside & (key == query)) + + @triton.jit def pack_attention_rows_kernel( cache, From 0d07bd6adbab1331652d3edcae5be0f3d05382f5 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 07:01:59 +0800 Subject: [PATCH 065/120] fix(hisparse): bound Ascend device-to-host DMA batches Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 860257525289..153381c4603f 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any import torch +from vllm.logger import logger from vllm.v1.hisparse.backend import HiSparseBackend if TYPE_CHECKING: @@ -51,6 +52,12 @@ def capture_aware_stream(stream: torch.Stream) -> torch.Stream: class AscendHiSparseBackend(HiSparseBackend): # Keep private pools until NPU IPC events interoperate with Core's generic streams. stats_dtype = torch.int64 + # Verified for row and page D2H on the installed A5 runtime. Use this + # conservative bound so the driver never receives an oversized batch. + _D2H_BATCH_LIMIT = 2047 + + def __init__(self) -> None: + self._failed_dma_batches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] def check_available(self) -> None: from vllm_ascend.hisparse.resolver import _load_resolver @@ -178,4 +185,25 @@ def gather_resident_rows(self, source: torch.Tensor, destination: torch.Tensor, super().gather_resident_rows(source, destination, row_ids) def swap_blocks_batch(self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor) -> None: - torch.ops._C_ascend.swap_blocks_batch(src, dst, sizes, 1, prefer_small_async=True) + if self._failed_dma_batches: + raise RuntimeError("HiSparse DMA submission previously failed; restart this worker.") + limit = self._D2H_BATCH_LIMIT + count = src.size(0) + if count > limit and any(t.ndim != 1 or t.size(0) != count or not t.is_contiguous() for t in (src, dst, sizes)): + raise ValueError("HiSparse D2H split requires equal-length contiguous descriptor arrays.") + submit = torch.ops._C_ascend.swap_blocks_batch + try: + if count <= limit: + submit(src, dst, sizes, 1, prefer_small_async=True) + return + for start in range(0, count, limit): + end = min(start + limit, count) + # Keep every chunk on Core's current DMA stream and on the batch API. + submit(src[start:end], dst[start:end], sizes[start:end], 1, prefer_small_async=False) + logger.info_once("[hisparse] Oversized D2H descriptor batch submitted in bounded chunks.", scope="process") + except BaseException: + # Core publishes its completion only after this method returns. Retain + # the original storage in the process-cached backend on an uncertain + # partial submit; Core shutdown drains the stream before freeing payloads. + self._failed_dma_batches.append((src, dst, sizes)) + raise From b86fdc48300b5cbbcfe850817a6ca9f087ebad1e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sat, 7 Nov 2026 12:48:54 +0800 Subject: [PATCH 066/120] perf(hisparse): align cold slot staging with Ascend producer dtype Use the bound Ascend slot producer dtype for empty private Host staging at cold initialization. Preserve Core streams, completion events, consumer int64 address arithmetic and the existing int64 producer fallback. CPU exact-source contract passed 574 checks. Device and model results are recorded separately in the PR qualification report. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/row_dma.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index 812ed79beb8e..b3db68737665 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -8,6 +8,7 @@ from pathlib import Path import numpy as np +import torch _WORKER_SHA256 = "864e2c7ce4e2cff320b4f68f8dccc9691adb829f0e50fdf6a27dc9c7885040e5" _RECYCLED_RUNTIME_SHA256S = { @@ -124,6 +125,15 @@ def init(worker): worker._ascend_row_dma_geometry = None if isinstance(worker.backend, backend_type): if worker.is_host_writer and worker.shared_host_region is None: + # Core allocates int64 staging; the bound Ascend slot producer uses int32. + # Cache binding is complete here and precedes graph capture. + staging = worker._slot_mapping_staging + slots = worker.cache_handles[0].slot_mapping + if staging is not None and not staging.num_tokens and slots.device.type == "npu": + if slots.dtype == torch.int32: + staging.slots = torch.empty( + staging.slots.shape, dtype=slots.dtype, device="cpu", pin_memory=True + ) worker._ascend_row_dma_geometry = _Geometry.bind(worker) if bind_recycled is not None: worker._ascend_recycled_invalidator = bind_recycled(worker) From 8172616dd4fe3fe22b33e2923cbede8351cd86af Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sun, 8 Nov 2026 04:42:06 +0800 Subject: [PATCH 067/120] chore(hisparse): trim stale qualification notes and test diagnostics Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 38 +++++-------------- release_notes/fragments/20926.yaml | 10 +---- .../one_card/test_hisparse_ops.py | 9 ----- 3 files changed, 11 insertions(+), 46 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 3197ecf84c31..98f1e9242ae8 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -99,32 +99,12 @@ or unmerged Core PR is required. ## Validation and performance -The native foundation has actual Ascend 950 eager/Graph component and TP8/MTP3 -full-model evidence. The C8 preprocessing and direct-index changes also passed -focused device checks and a fresh full-model ON/OFF comparison. See the PR for -the exact source pins, commands and qualification boundaries. - -In the current GLM-5.2 C4 comparison (32768 input and 8192 output tokens, -TP8/DP1, MTP3, Model Runner V2, prefix caching off and explicit -`weight_nz_mode=0`), both arms completed two rounds with 65536 actual output -token IDs per arm. ON's settled reported miss/H2D deltas were zero. Mean TPOT was -12.313 ms OFF and 13.249 ms ON: an observed increase of 7.61%. All four rounds -had identical complete input/output-sequence multisets and MTP counters -(59.8936% acceptance). Sampled steady Decode throughput decreased by 4.58%. -TPOT includes the mixed Prefill/Decode portion of the requests, so its increase -must not be attributed entirely to steady Decode or to a single operation. - -In separate ON comparisons with identical inputs, outputs and MTP work, native C8 output and -length reuse reduced observed mean TPOT by 2.19% across two rounds. Bounding -compact-refill programs reduced TPOT by 1.93% in one matching round. These gains -are not additive and do not establish statistical significance or acceptable -performance for other workloads. Compact refill retains current device counts, -page-strided byte copies and the existing completion dependencies. - -The request mapping, page-strided Indexer store and small-copy queue passed -focused eager/Graph NPU checks before the current full-model comparison. These -checks do not establish an isolated end-to-end gain for each individual change. - -Dataset accuracy, the full-model miss path and a TPOT regression within 3% -remain unqualified. Profiling requests are separate from formal performance -timing; reported zero-miss counters do not prove the absence of device DMA. +Ascend 950 eager/Graph component and TP8/MTP3 full-model checks cover the +supported foundation. Performance depends on the workload, residency and model +configuration. See [PR #20926](https://github.com/vllm-project/vllm-ascend/pull/20926) +for exact source pins, commands, measured ON/OFF results and remaining +qualification limits. + +Formal performance runs exclude profiling. Reported miss and H2D counters are +asynchronous software observations, not physical DMA measurements. Qualification +on one workload does not establish dataset accuracy or support for other devices. diff --git a/release_notes/fragments/20926.yaml b/release_notes/fragments/20926.yaml index 2edffa770acd..066e946ad7d3 100644 --- a/release_notes/fragments/20926.yaml +++ b/release_notes/fragments/20926.yaml @@ -4,12 +4,6 @@ entries: text: >- Added experimental Ascend HiSparse BF16 and native C8 sparse MLA/SFA support using a fused AscendC residency resolver, existing Attention - operators and Core lifecycle interfaces. The resolver requires its - compiled extension and does not fall back to Triton. Native C8 reuses - supported SFA Prolog preprocessing and consumes the resolver's physical - Decode indices while preserving the separate refill plan. Decode lengths - are prepared with the cache binding and native empty-query clearing is - reused without redundant output masking. Compact refill bounds the number - of per-query programs while retaining current counts and scalar byte - copies for both native C8 and BF16 payloads. + operators and Core lifecycle interfaces. The native resolver requires + its compiled extension and does not fall back to Triton. hardware: [A5] diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index d0839dc7d2ff..fe972e877365 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -2,12 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Real-NPU HiSparse state and mapped-payload tests against a CPU oracle.""" -import json - import pytest import torch import torch_npu # noqa: F401 - from vllm_ascend.hisparse.backend import AscendHiSparseBackend, capture_aware_stream from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, @@ -308,7 +305,6 @@ def attention_reference(physical): src, dst = int(expected[2][row, col]), int(expected[3][row, col]) torch.testing.assert_close(hot[dst // 4, dst % 4].cpu(), host[src]) torch.testing.assert_close(backing[:, 0].cpu(), torch.full((10, 4, 16), -17.0)) - print(json.dumps({"resolver_mtp_union_lru_compact_gather": "PASS"}), flush=True) latent = torch.empty((4, 8, 1, 12), device="npu") rope = torch.empty((4, 8, 1, 4), device="npu") pack_attention_rows(hot, out, latent, rope, 12) @@ -328,7 +324,6 @@ def attention_reference(physical): torch.npu.synchronize() torch.testing.assert_close(hot[0, 0].cpu(), torch.cat((new_latent[0].flatten().cpu(), new_rope[0].flatten().cpu()))) torch.testing.assert_close(hot[0, 2].cpu(), hot_cpu[0, 2]) - print(json.dumps({"pack_scatter_strides_padding": "PASS"}), flush=True) # Capture warmed kernels, reset state outside capture, then verify replay # after changing the mapped CPU payload. No host selection work runs here. stream = torch.npu.Stream() @@ -356,14 +351,12 @@ def attention_reference(physical): for col in range(int(expected[4][row])): src, dst = int(expected[2][row, col]), int(expected[3][row, col]) torch.testing.assert_close(hot[dst // 4, dst % 4].cpu(), host[src]) - print(json.dumps({"acl_graph_replay_mapped_host_update": "PASS"}), flush=True) for _ in range(100): graph.replay() torch.npu.synchronize() torch.testing.assert_close(out.cpu(), expected[0]) torch.testing.assert_close(dgi_n.cpu(), expected[5]) torch.testing.assert_close(lru_n.cpu(), expected[6]) - print(json.dumps({"repeated_graph_replay": "PASS", "iterations": 100}), flush=True) if device_row_count: for active_rows in [0, 1, 2, 3]: valid_rows.fill_(active_rows) @@ -846,8 +839,6 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): sfa_dcp_replicated_indexer_size=1, ) key, scale = _quantized_indexer_views(packed, spec) - with pytest.raises(RuntimeError, match="view size"): - key.view(-1, width) writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) torch.nn.Module.__init__(writer) writer.k_cache = SimpleNamespace(kv_cache=(key, scale)) From 7a6e3264df3d62789ea888fc7bc324090837151a Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Sun, 8 Nov 2026 08:29:30 +0800 Subject: [PATCH 068/120] test(hisparse): remove duplicated scratch arithmetic assertions Signed-off-by: zhenwenqi2024 --- .../kv_transfer/test_hisparse_foundation_cpu.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index ee404e2e0c17..016baf362e40 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -190,16 +190,6 @@ def group(): assert torch.count_nonzero(second_workspace.first_refs) == 0 -@pytest.mark.parametrize("hot,union", [(16, 8), (8192, 8192), (12288, 64), (32768, 32768)]) -def test_native_resolver_scratch_follows_request_capacity(hot, union): - from vllm_ascend.hisparse.resolver import scratch_layout - - buckets, size = scratch_layout(hot, union) - assert buckets > union and buckets & (buckets - 1) == 0 - words = 2 * buckets + hot + 2 * ((hot + 31) // 32) + 64 - assert size == (words + 15) // 16 * 64 + 512 - - def test_native_resolver_uses_available_ub_without_truncating_the_union(): from vllm_ascend.hisparse.resolver import scratch_layout From 114ab3eca3f6e0da21c43ebc70ac1f39e24c2cd2 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 00:44:49 +0800 Subject: [PATCH 069/120] test(hisparse): bootstrap SFA ops before device initialization Match the platform initialization order so standalone C8 tests can locate custom tiling. A fresh private A5 reference component passed all60 eager/Graph cases after this fixture change, including original and prepared-query controls. This branch carries only the fixture fix; runtime/native sources stay identical to7a6e3264. Exact public-source standalone validation is pending. Signed-off-by: zhenwenqi2024 --- .../e2e/pull_request/one_card/test_hisparse_attention.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index c28fb7997933..cc2f1b5a3ac0 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -13,6 +13,11 @@ @pytest.fixture(autouse=True) def _require_npu_sfa(): + from vllm_ascend.utils import bootstrap_custom_op_env + + # Match platform initialization: publish the OPP path before device queries + # or set_device can initialize the runtime and its tiling registry. + bootstrap_custom_op_env() if not torch.npu.is_available(): pytest.skip("HiSparse SFA tests require an NPU") from vllm_ascend.device.device_config import is_310p @@ -20,10 +25,7 @@ def _require_npu_sfa(): if is_310p(): pytest.skip("The native SFA operation is not registered for Ascend 310P") torch.npu.set_device(0) - from vllm_ascend.utils import bootstrap_custom_op_env - # Load the required SFA extension without changing custom-op selection. - bootstrap_custom_op_env() import vllm_ascend.vllm_ascend_C # noqa: F401 From 47a8276372b67c3a7f99939048edf74a39c0d6a7 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 02:04:33 +0800 Subject: [PATCH 070/120] perf(hisparse): size resolved-row scratch by configured TopK Signed-off-by: zhenwenqi2024 --- .../kv_transfer/test_hisparse_foundation_cpu.py | 13 +++++++++++-- vllm_ascend/hisparse/backend.py | 1 + vllm_ascend/hisparse/cache_ops.py | 17 ++++++++++++++--- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 016baf362e40..515670b2f56c 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -160,7 +160,10 @@ def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups backend = AscendHiSparseBackend() def group(): - owner = NS(copy_stream=NS(device=torch.device("cpu"))) + owner = NS( + copy_stream=NS(device=torch.device("cpu")), + shared_topk=NS(device_topk_rows=torch.empty((max_rows, union_rows), dtype=torch.int32)), + ) leader = NS( index_group=owner, max_swap_rows=max_rows, @@ -181,7 +184,7 @@ def group(): assert follower.host_mapping is first.registered_host_pool._hisparse_registration assert second.cache_ops.workspace is not first.cache_ops.workspace first_workspace, second_workspace = first.cache_ops.workspace, second.cache_ops.workspace - assert first_workspace.first_refs.shape == (max_rows * hot_rows,) + assert first_workspace.first_refs.shape == (max_rows * union_rows,) assert first_workspace.errors.shape == (max_rows, 8) assert first_workspace.max_union_rows == union_rows first_workspace.first_refs.fill_(13) @@ -190,6 +193,12 @@ def group(): assert torch.count_nonzero(second_workspace.first_refs) == 0 +@pytest.mark.parametrize("topk", [0, 17]) +def test_resolver_scratch_rejects_topk_outside_hot_capacity(topk): + with pytest.raises(ValueError, match="TopK capacity"): + hisparse_ops.ResolverWorkspace(4, 16, 8, torch.device("cpu"), max_topk=topk) + + def test_native_resolver_uses_available_ub_without_truncating_the_union(): from vllm_ascend.hisparse.resolver import scratch_layout diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 153381c4603f..88b2e762ea72 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -114,6 +114,7 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: hot_size=runtime.region_stride, max_union_rows=runtime.max_union_rows, device=runtime.device, + max_topk=group.shared_topk.device_topk_rows.shape[1], ) else: leader_ops = group.leader.cache_ops diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index ac68b399f89e..3a7e45f060b6 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -118,17 +118,23 @@ def allocate_mapped_host(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Te class ResolverWorkspace: """Fixed scratch shared by one index group on its ordered copy stream.""" - def __init__(self, max_rows: int, hot_size: int, max_union_rows: int, device: torch.device) -> None: + def __init__( + self, max_rows: int, hot_size: int, max_union_rows: int, device: torch.device, *, max_topk: int | None = None + ) -> None: if not 0 < max_union_rows <= hot_size <= 32768 or max_rows <= 0: raise ValueError("HiSparse requires 0 < union <= hot_size <= 32768 and positive max_rows") self.max_rows = max_rows self.hot_size = hot_size self.max_union_rows = max_union_rows + self.max_topk = hot_size if max_topk is None else max_topk + if not 0 < self.max_topk <= hot_size: + raise ValueError("HiSparse resolver TopK capacity must be positive and fit the Hot cache") from vllm_ascend.hisparse.resolver import _load_resolver, scratch_layout self.native = _load_resolver() scratch_layout(hot_size, max_union_rows, self.native.hisparse_resolver_scratch_capacity()) - self.first_refs = torch.empty(max_rows * hot_size, dtype=torch.int32, device=device) + # One resolved Host row per selected token, not per Hot-cache slot. + self.first_refs = torch.empty(max_rows * self.max_topk, dtype=torch.int32, device=device) self.errors = torch.zeros((max_rows, 8), dtype=torch.int32, device=device) @@ -254,7 +260,12 @@ def hisparse_resolve_residency( hot_size = lru_slots.size(1) if rows == 0: return - if rows > workspace.max_rows or hot_size != workspace.hot_size or max_union_rows != workspace.max_union_rows: + if ( + rows > workspace.max_rows + or topk > workspace.max_topk + or hot_size != workspace.hot_size + or max_union_rows != workspace.max_union_rows + ): raise ValueError("HiSparse resolver input exceeds its prepared workspace") if region_stride != device_global_indices.size(1) or region_stride != hot_size or hot_size < topk: raise ValueError("HiSparse region and LRU widths must match and cover top-k") From 251f0d92fd085760304990440e4f2d5c91dba5d4 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 11:51:16 +0800 Subject: [PATCH 071/120] test(hisparse): route device coverage to A5 and refresh fixtures Signed-off-by: zhenwenqi2024 --- .github/workflows/scripts/test_config.yaml | 6 +++++ .../one_card/test_hisparse_kv_preprocess.py | 3 +++ .../test_hisparse_foundation_cpu.py | 27 ++++++++++++++----- 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/.github/workflows/scripts/test_config.yaml b/.github/workflows/scripts/test_config.yaml index e562e539b6f0..56aa0c225171 100644 --- a/.github/workflows/scripts/test_config.yaml +++ b/.github/workflows/scripts/test_config.yaml @@ -48,6 +48,10 @@ curated_tests: - tests/e2e/pull_request/one_card/test_mla_fp8_decode_layout_a5.py - tests/e2e/pull_request/one_card/test_flash_attn_precision.py - tests/e2e/pull_request/one_card/test_swiglu_mx_quant.py + - tests/e2e/pull_request/one_card/test_hisparse_attention.py + - tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py + - tests/e2e/pull_request/one_card/test_hisparse_ops.py + - tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py # === Skipped Tests === # Removed from any selection after directory scanning, so they are never @@ -121,6 +125,8 @@ runner_mapping: default: a5-4 tests/e2e/pull_request/one_card/test_swiglu_mx_quant\.py$: default: a5-4 + tests/e2e/pull_request/one_card/test_hisparse_(attention|kv_preprocess|ops|prefill_staging)\.py$: + default: a5-4 tests/e2e/pull_request/one_card: default: a2-1 310p: 310p-1 diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index 0c9362336b34..729e5347431a 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -134,6 +134,9 @@ def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False impl.sfa_qsfa_tile_size = block impl._kv_staging_capacity = block + # This case exercises the staged destination, not the separately bound + # direct-Resident capture path installed by the model runner. + impl._direct_prolog_slots = None def mx_layer(inputs, outputs): weight, scale = torch_npu.npu_dynamic_mx_quant( diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 515670b2f56c..c18b4ecf0e81 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -202,10 +202,26 @@ def test_resolver_scratch_rejects_topk_outside_hot_capacity(topk): def test_native_resolver_uses_available_ub_without_truncating_the_union(): from vllm_ascend.hisparse.resolver import scratch_layout - buckets, size = scratch_layout(12288, 8192, 175616) - assert buckets == 15328 and buckets > 8192 and size == 175616 - with pytest.raises(RuntimeError, match="needs at least 118592 UB bytes"): - scratch_layout(12288, 8192, 118591) + hot, union, capacity = 12288, 8192, 175616 + buckets, size = scratch_layout(hot, union, capacity) + assert buckets > union + assert 0 < size <= capacity and size % 64 == 0 + # Find the actual supported boundary without pinning a particular hash + # bucket count or the precise UB layout/error-message arithmetic. + too_small, fits = 0, capacity + while fits - too_small > 1: + candidate = (too_small + fits) // 2 + try: + scratch_layout(hot, union, candidate) + except RuntimeError: + too_small = candidate + else: + fits = candidate + minimum_buckets, minimum_size = scratch_layout(hot, union, fits) + assert minimum_buckets > union and minimum_size <= fits + assert minimum_size % 64 == 0 + with pytest.raises(RuntimeError, match="UB bytes"): + scratch_layout(hot, union, fits - 1) def test_missing_native_resolver_raises_instead_of_falling_back(monkeypatch): @@ -299,8 +315,7 @@ def test_capture_metadata_prevents_warmup_host_mirrors(monkeypatch): real_metadata = NS(hisparse_for_capture=False) submitted = [] cache = NS(finish_kv_update=lambda: submitted.append(True)) - for _ in range(5): - finish_hisparse_kv_update(cache, capture_metadata) + finish_hisparse_kv_update(cache, capture_metadata) assert submitted == [] finish_hisparse_kv_update(cache, real_metadata) assert submitted == [True] From 5a02c50e38e78bf76f509ba2ca99b412a4b2a739 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 08:17:10 +0800 Subject: [PATCH 072/120] perf(hisparse): pass Resident staging extents at runtime Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 48 ++++++++++++++++++- vllm_ascend/ops/triton/hisparse.py | 4 +- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 18c533f62df6..9372c3e3065e 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -10,7 +10,12 @@ from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView from vllm_ascend.hisparse.backend import AscendHiSparseBackend -from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, allocate_mapped_host, scatter_packed_kv +from vllm_ascend.hisparse.cache_ops import ( + HiSparseNPUOps, + allocate_mapped_host, + gather_valid_resident_rows, + scatter_packed_kv, +) from vllm_ascend.hisparse.prefill import build_hisparse_prefill_staging_plan @@ -115,6 +120,47 @@ def test_backend_resident_payload_graph_replay(dtype, width): assert torch.equal(destination.view(torch.uint8).cpu(), expected) +@pytest.mark.parametrize("row_bytes", [37, 656]) +@pytest.mark.parametrize("capture", [False, True]) +def test_prefill_resident_patch_changing_extents(row_bytes, capture): + if not torch.npu.is_available(): + pytest.skip("NPU required") + torch.npu.set_device(0) + block = 4 + for rows, source_blocks in ((8, 4), (24, 6), (40, 8)): + source_cpu = (torch.arange(source_blocks * 5 * block * row_bytes) % 251).to(torch.uint8) + source_cpu = source_cpu.view(source_blocks, 5, block, row_bytes) + source = source_cpu.npu() + destination = torch.full((rows // block, 2, block, row_bytes), 7, dtype=torch.uint8, device="npu:0") + indices = torch.empty(rows, dtype=torch.int32, device="npu:0") + source_view, destination_view = source[:, 0], destination[:, 0] + + indices.fill_(-1) + gather_valid_resident_rows(source_view, destination_view, indices) + torch.npu.synchronize() + graph = None + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + gather_valid_resident_rows(source_view, destination_view, indices) + for iteration in range(2): + values = (torch.arange(rows, dtype=torch.int32) + iteration * 3) % (source_blocks * block + 2) - 1 + current_source = source_cpu + iteration * 13 + source.copy_(current_source) + indices.copy_(values) + destination.fill_(7) + if graph is None: + gather_valid_resident_rows(source_view, destination_view, indices) + else: + graph.replay() + torch.npu.synchronize() + expected = torch.full(destination.shape, 7, dtype=torch.uint8) + for row, selected in enumerate(values.tolist()): + if 0 <= selected < source_blocks * block: + expected[row // block, 0, row % block] = current_source[selected // block, 0, selected % block] + assert torch.equal(destination.cpu(), expected) + + def test_prefill_singleton_column_mapping_tracks_current_residency(): if not torch.npu.is_available(): pytest.skip("HiSparse prefill mapping requires an NPU") diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index c0f36fecfc92..790d0633c57b 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -219,8 +219,8 @@ def gather_valid_resident_rows_kernel( source, destination, row_ids, - ROWS: tl.constexpr, - SOURCE_ROWS: tl.constexpr, + ROWS, + SOURCE_ROWS, SOURCE_BLOCK: tl.constexpr, SOURCE_STRIDE: tl.constexpr, DEST_BLOCK: tl.constexpr, From e661dc8713b2fece6fe8167d9baf074d1ce68f9b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 17:01:23 +0800 Subject: [PATCH 073/120] refactor(hisparse): trim redundant state and tests Signed-off-by: zhenwenqi2024 --- CMakeLists.txt | 2 - .../user_guide/feature_guide/hisparse.md | 64 ++++++++----------- .../one_card/test_hisparse_ops.py | 30 ++------- vllm_ascend/hisparse/mirror_overlap.py | 2 +- vllm_ascend/hisparse/recycled_invalidate.py | 4 +- vllm_ascend/hisparse/row_dma.py | 4 +- vllm_ascend/hisparse/target_snapshot.py | 2 +- vllm_ascend/worker/v2/model_runner.py | 2 +- 8 files changed, 40 insertions(+), 70 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index dbb2fe8328c7..2a9afe8c3e85 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,8 +26,6 @@ if(NOT ${TORCH_VERSION} VERSION_EQUAL "2.10.0") message(FATAL_ERROR "Expected PyTorch version 2.10.0, but found ${TORCH_VERSION}") endif() - - set(RUN_MODE "npu" CACHE STRING "cpu/sim/npu") set(SOC_VERSION ${SOC_VERSION}) message(STATUS "Detected SOC version: ${SOC_VERSION}") diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 98f1e9242ae8..c3ae0969f64a 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -28,12 +28,9 @@ the kernel into `vllm_ascend_C` and registers `torch.ops._C_ascend.hisparse_resolve_residency` with the other Ascend operators; there is no separate resolver build or library-loading step. -The architecture and usable per-AIV UB budget are selected for the build target -in `csrc/hisparse/CMakeLists.txt`; these are hardware/compiler parameters. -The kernel sizes its hash table from the request union and available UB. An -unsupported capacity or missing extension raises an error before execution; -there is no Triton resolver fallback. Builds for other targets retain their -existing operators without registering this resolver. +An unsupported capacity or missing extension raises an error before execution; +there is no Triton resolver fallback. Other build targets retain their existing +operators without registering this resolver. ## Enable HiSparse @@ -69,41 +66,32 @@ victims, and produces the refill plan on device. Hot capacity must fit both the int16 LRU representation (at most 32768 rows) and the compiled target's UB budget. These checks do not change the configured TopK or discard selected tokens. -Native C8 main rows contain 512 E4M3FN bytes, 64 BF16 RoPE values and four FP32 -scales (656 bytes total). Indexer rows contain 128 E4M3FN keys and one FP32 scale. -Views preserve the shared-slab page strides and layer offsets. The existing -quantizer and Attention operations are reused; payloads are copied as bytes. -With MXFP8 projection weights, C8 reuses SFA's existing Prolog V3 preprocessing -and publishes its staged rows through the HiSparse Resident mapping. Other -supported projection schemes retain the native preprocessing path. - -C8 Decode consumes the resolver's mapped, stably compacted indices directly. -The refill plan retains separate virtual page-row IDs. Prefill keeps its existing -layout bridge, including stable ordering and duplicate selections. -Decode's padded query-shape lengths and physical KV extent are prepared before -capture and owned by the current binding. Per-round request mappings and valid -rows remain current. Native C8 Attention already clears empty query outputs; -the adapter returns that output without another index scan and output mask. - -Request-row mapping reuses Core's device implementation, with separate storage -for each captured MTP graph owner and Ascend's padded input extent. Indexer -writes reuse the existing SFA store with independent page strides for keys and -scales, preserving the original payload bits. Small Host-mirror submissions use -PTA's ordered queue with immutable scalar descriptors; full-layer batches keep -the batch-copy API and Core's completion and payload-ownership boundaries. - -Keep completion-before-unregister and graph allocation ownership intact. -Overwriting KV still invalidates any old Hot copy. Equal batch shapes do not -permit reusing stale request mappings or lengths. No upstream vLLM source change -or unmerged Core PR is required. +Native C8 uses the existing Ascend E4M3FN values, BF16 RoPE and FP32 scales. +Weight quantization and KV-cache quantization remain separate choices. HiSparse +preserves the shared-slab strides and uses the existing quantizer and Attention +operations. No upstream vLLM source change or unmerged Core PR is required. + +Zero misses remove cold-row refills; eager Host mirroring and residency +bookkeeping still run. Size Host memory per rank and compare matched workloads +when assessing performance. Prefill, mixed batches and Decode retain current +request mappings, padding and completion-before-reuse rules. ## Validation and performance -Ascend 950 eager/Graph component and TP8/MTP3 full-model checks cover the -supported foundation. Performance depends on the workload, residency and model -configuration. See [PR #20926](https://github.com/vllm-project/vllm-ascend/pull/20926) -for exact source pins, commands, measured ON/OFF results and remaining -qualification limits. +Ascend 950 eager/Graph components and TP8/MTP3 full-model checks cover the +supported paths. A matched C8 test with two rounds per arm, fixed inputs and +8K outputs measured these mean SSE TPOT increases: + +| Input / concurrency | Reported residency | ON versus OFF | +| --- | --- | --- | +| 16K / C8 | Zero misses | +1.03% | +| 32K / C8 | Misses present | +10.01% | + +These measurements used merged Core `39dbf518` (including #63801), Ascend +`5a02c50`, Python 3.12 and CANN 9.2.1 on the same machine. Inputs, outputs and +MTP counters matched. N2 does not establish statistical significance or dataset +accuracy. See [PR #20926](https://github.com/vllm-project/vllm-ascend/pull/20926) +for exact configuration, commands and separate capacity qualifications. Formal performance runs exclude profiling. Reported miss and H2D counters are asynchronous software observations, not physical DMA measurements. Qualification diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index fe972e877365..d23aa74de9c1 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -5,6 +5,7 @@ import pytest import torch import torch_npu # noqa: F401 + from vllm_ascend.hisparse.backend import AscendHiSparseBackend, capture_aware_stream from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, @@ -26,13 +27,12 @@ def _require_mapped_host_device(): AscendHiSparseBackend().check_available() -@pytest.mark.parametrize("descriptor_count", [0, 1, 2, 4, 16, 17, 312]) +@pytest.mark.parametrize("descriptor_count", [0, 1, 16, 17, 312]) @pytest.mark.parametrize("graph_producer", [False, True]) def test_private_host_pool_retains_pinned_storage_for_dma_writeback(descriptor_count, graph_producer): from vllm.v1.hisparse import runtime backend = AscendHiSparseBackend() - backend.check_available() capacity = max(32768, (descriptor_count + 1) * 2048) payload, registered = runtime.allocate_pinned_host_pool(capacity) graph = None @@ -201,7 +201,6 @@ def reference(topk, requests, states, source, resident, hot_table, dgi, lru, blo @pytest.mark.parametrize("device_row_count", [False, True]) @pytest.mark.parametrize("compact_attention", [False, True]) def test_resolver_strides_padding_graph(device_row_count, compact_attention): - torch.npu.set_device(0) raw = torch.empty(8192, dtype=torch.uint8) offset = (-raw.data_ptr()) % 4096 registered = raw[offset : offset + 4096] @@ -351,12 +350,6 @@ def attention_reference(physical): for col in range(int(expected[4][row])): src, dst = int(expected[2][row, col]), int(expected[3][row, col]) torch.testing.assert_close(hot[dst // 4, dst % 4].cpu(), host[src]) - for _ in range(100): - graph.replay() - torch.npu.synchronize() - torch.testing.assert_close(out.cpu(), expected[0]) - torch.testing.assert_close(dgi_n.cpu(), expected[5]) - torch.testing.assert_close(lru_n.cpu(), expected[6]) if device_row_count: for active_rows in [0, 1, 2, 3]: valid_rows.fill_(active_rows) @@ -701,7 +694,7 @@ def forward(): mapping.close() -def test_plan_gather_and_written_slot_invalidation(): +def test_plan_gather_respects_request_state_padding(): host = allocate_mapped_host((32, 16), torch.float32) mapping = host._hisparse_registration host.copy_(torch.arange(host.numel(), dtype=torch.float32).view_as(host)) @@ -718,24 +711,18 @@ def test_plan_gather_and_written_slot_invalidation(): cache = hot.cpu().view(32, 16) torch.testing.assert_close(cache[5:13], host[1:9]) torch.testing.assert_close(attention[1].cpu(), torch.full((8,), -1, dtype=torch.int32)) - dgi = torch.tensor([[4, 7, 9, -1], [11, 13, 15, -1]], dtype=torch.int32, device="npu") - states = torch.tensor([1, 0], dtype=torch.int32, device="npu") - reqs = torch.tensor([0, 1, 1, -1], dtype=torch.int32, device="npu") - written = torch.tensor([13, 4, 9, 11], dtype=torch.int64, device="npu") - ops.hisparse_invalidate_written_slots(dgi, states, reqs, written) - torch.npu.synchronize() - torch.testing.assert_close(dgi.cpu(), torch.tensor([[-1, 7, -1, -1], [11, -1, 15, -1]], dtype=torch.int32)) finally: torch.npu.synchronize() mapping.close() -def test_invalidate_large_hot_capacity_graph(): - dgi = torch.arange(20480, dtype=torch.int32, device="npu").view(2, 10240) +@pytest.mark.parametrize("hot_capacity", [4, 10240]) +def test_invalidate_hot_capacity_graph(hot_capacity): + dgi = torch.arange(2 * hot_capacity, dtype=torch.int32, device="npu").view(2, hot_capacity) initial = dgi.cpu() states = torch.tensor([1, 0], dtype=torch.int32, device="npu") requests = torch.tensor([0, 1, 1, -1], dtype=torch.int32, device="npu") - written = torch.tensor([20479, 0, 10239, 10240], dtype=torch.int64, device="npu") + written = torch.tensor([2 * hot_capacity - 1, 0, hot_capacity - 1, hot_capacity], dtype=torch.int64, device="npu") expected = initial.clone() expected[1, -1] = expected[0, 0] = expected[0, -1] = -1 @@ -817,9 +804,6 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.hisparse.binding import _quantized_indexer_views - if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): - pytest.skip("The native E4M3FN HiSparse format currently targets A5") - torch.npu.set_device(0) pages, block, width, page_stride = 3, 128, 128, 419840 size = (pages - 1) * page_stride + block * width + block * 4 raw = torch.full((size,), 90, dtype=torch.uint8, device="npu") diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/mirror_overlap.py index 3dbcd4f7ee7a..e8db7472124e 100644 --- a/vllm_ascend/hisparse/mirror_overlap.py +++ b/vllm_ascend/hisparse/mirror_overlap.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Private decode experiment for Core 24420b7's sampling and draft boundaries.""" +"""Overlap HiSparse Host mirroring with Target sampling and guarded MTP.""" from __future__ import annotations diff --git a/vllm_ascend/hisparse/recycled_invalidate.py b/vllm_ascend/hisparse/recycled_invalidate.py index 2d7d1a783def..501356e6bd64 100644 --- a/vllm_ascend/hisparse/recycled_invalidate.py +++ b/vllm_ascend/hisparse/recycled_invalidate.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Grouped invalidation for recycled Host blocks in the pinned HiSparse worker.""" +"""Batch recycled Host-slot invalidation within Core's worker lifecycle.""" import operator @@ -67,8 +67,8 @@ def _invalidate_recycled_groups( class RecycledGroupInvalidator: def __init__(self, tables, block_size): self._rank = get_tensor_model_parallel_rank() + # Retain the bound tensors and their storage through queued invalidations. self._tables = tables - self._storages = tuple(table.untyped_storage() for table in tables) self._addresses = tuple(table.data_ptr() for table in tables) self._device = tables[0].device self._states, self._hot = tables[0].shape diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index b3db68737665..e2c76dea8137 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Private experiment for the frozen Core HiSparse row-mirror worker.""" +"""Prepare HiSparse row-mirror DMA through Core's descriptor ownership.""" from dataclasses import dataclass from functools import wraps @@ -161,7 +161,7 @@ def invalidate(worker, block_ids, request_state_indices): def install(backend_type) -> None: - """Only adapt the exact Core implementation reviewed for this experiment.""" + """Only adapt the exact Core worker implementation reviewed for this backend.""" from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker from vllm.v1.hisparse import runtime diff --git a/vllm_ascend/hisparse/target_snapshot.py b/vllm_ascend/hisparse/target_snapshot.py index ece4ed6c5ac1..7281b8c2c878 100644 --- a/vllm_ascend/hisparse/target_snapshot.py +++ b/vllm_ascend/hisparse/target_snapshot.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Private Target snapshot experiment for the frozen HiSparse worker.""" +"""Own Target row snapshots until asynchronous Host mirroring completes.""" from dataclasses import dataclass diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 893a60dd48ce..4e87cd1a550e 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -527,7 +527,7 @@ def initialize_kv_cache( if self.vllm_config.attention_config.hisparse_config is not None: self._bind_hisparse_direct_prolog() - # Private candidate only: bind after Core seals all cache/stream owners. + # Bind after Core seals all cache and stream owners. from vllm_ascend.hisparse.mirror_overlap import overlap_mtp_host_mirror overlap_connector = overlap_mtp_host_mirror(self) From c81a48facce309b52f2ce780d2f5fa3627f911a7 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 15:15:57 +0800 Subject: [PATCH 074/120] perf(hisparse): reuse target logits index output workspace Signed-off-by: zhenwenqi2024 (cherry picked from commit 9f1f86a8995bdc173fba0cd2066c797d359b7ff1) --- .../model_runner_v2/test_input_batch.py | 123 ++++++++++++++++++ tests/ut/worker/test_model_runner_v2.py | 6 +- vllm_ascend/worker/v2/input_batch.py | 52 +++++++- vllm_ascend/worker/v2/model_runner.py | 8 +- 4 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py diff --git a/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py b/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py new file mode 100644 index 000000000000..60fa121fed39 --- /dev/null +++ b/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real-NPU input preparation tests for the V2 model runner.""" + +import pytest +import torch +import torch_npu # noqa: F401 +from vllm.v1.worker.gpu.input_batch import combine_sampled_and_draft_tokens + +from vllm_ascend.worker.v2.input_batch import AscendInputBuffers + + +@pytest.fixture(autouse=True) +def _require_npu(): + if not torch.npu.is_available(): + pytest.skip("NPU input preparation requires an available NPU") + torch.npu.set_device(0) + + +def _combine_inputs_cpu(input_ids, metadata, num_new): + """Independent value oracle, including the kernel's input-ID side effects.""" + result, indices = input_ids.clone(), [] + mapping, sampled, starts, lengths, prefill, drafts, cumulative = metadata.values() + for row, state in enumerate(mapping.tolist()): + count = int(cumulative[row + 1] - cumulative[row]) + end = int(starts[row + 1]) + indices.extend(range(end - count, end)) + if int(lengths[row]) <= int(prefill[state]): + continue + if num_new and int(lengths[row]) - count >= int(prefill[state]): + result[end - count] = sampled[state] + draft_count = count - num_new + if draft_count > 0: + result[end - draft_count : end] = drafts[state, :draft_count] + return result, torch.tensor(indices, dtype=torch.int64) + + +@pytest.mark.parametrize("num_new", [0, 1]) +@pytest.mark.parametrize("capture", [False, True]) +def test_target_logits_workspace_matches_core_with_changing_inputs(num_new, capture): + owners = [AscendInputBuffers(3, 32, torch.device("npu:0")) for _ in range(2)] + for owner in owners: + owner.bind_logits_indices_workspace() + assert owners[0].logits_indices_workspace.data_ptr() != owners[1].logits_indices_workspace.data_ptr() + count_sets = ( + (((1, 4), (4, 1), (2, 3)), ((1, 1),) * 3) + if num_new + else ( + ((0, 3), (3, 0), (1, 2)), + ((0, 1), (1, 0), (0, 1)), + ) + ) + graphs, retained_metadata = [], [] + try: + for owner_index, owner in enumerate(owners): + cpu = { + "mapping": torch.tensor([2, 0], dtype=torch.int32), + "sampled": torch.empty(3, dtype=torch.int32), + "starts": torch.tensor([0, 4, 8], dtype=torch.int32), + "lengths": torch.tensor([16, 4], dtype=torch.int32), + "prefill": torch.tensor([4, 20, 9], dtype=torch.int32), + "drafts": torch.empty((3, 3), dtype=torch.int32), + "cumulative": torch.empty(3, dtype=torch.int32), + } + device = {name: value.to("npu") for name, value in cpu.items()} + retained_metadata.append((cpu, device)) + + def update(change, counts, owner_index, owner, cpu, device): + salt = owner_index * 100 + change * 10 + cpu["mapping"].copy_(torch.tensor([2, 0] if change % 2 == 0 else [0, 2], dtype=torch.int32)) + cpu["sampled"].copy_(torch.arange(3, dtype=torch.int32) + 1000 + salt) + cpu["drafts"].copy_((torch.arange(9, dtype=torch.int32) + 2000 + salt).reshape(3, 3)) + cpu["cumulative"].copy_(torch.tensor([0, counts[0], sum(counts)], dtype=torch.int32)) + for name, value in cpu.items(): + device[name].copy_(value) + initial = torch.arange(32, dtype=torch.int32) + 70 + salt + owner.input_ids.copy_(initial) + owner.logits_indices_workspace.fill_(-777) + return initial + + for counts_sequence in count_sets: + num_logits = sum(counts_sequence[0]) + update(0, counts_sequence[0], owner_index, owner, cpu, device) + output = owner.combine_sampled_and_draft_tokens(owner.input_ids, *device.values(), num_logits, num_new) + torch.npu.synchronize() + graph = None + if capture: + graph = torch.npu.NPUGraph() + graphs.append(graph) + with torch.npu.graph(graph): + output = owner.combine_sampled_and_draft_tokens( + owner.input_ids, *device.values(), num_logits, num_new + ) + for change, counts in enumerate(counts_sequence): + initial = update(change, counts, owner_index, owner, cpu, device) + if graph is None: + output = owner.combine_sampled_and_draft_tokens( + owner.input_ids, *device.values(), num_logits, num_new + ) + else: + graph.replay() + baseline_ids = initial.to("npu") + baseline = combine_sampled_and_draft_tokens(baseline_ids, *device.values(), num_logits, num_new) + expected_ids, expected_indices = _combine_inputs_cpu(initial, cpu, num_new) + torch.testing.assert_close(output.cpu(), expected_indices, rtol=0, atol=0) + torch.testing.assert_close(baseline.cpu(), expected_indices, rtol=0, atol=0) + torch.testing.assert_close(owner.input_ids.cpu(), expected_ids, rtol=0, atol=0) + torch.testing.assert_close(baseline_ids.cpu(), expected_ids, rtol=0, atol=0) + assert output.data_ptr() == owner.logits_indices_workspace.data_ptr() + assert torch.all(owner.logits_indices_workspace[num_logits:].cpu() == -777) + empty = dict(device) + empty["mapping"] = device["mapping"][:0] + empty["starts"] = device["starts"][:1] + empty["lengths"] = device["lengths"][:0] + empty["cumulative"] = device["cumulative"][:1] + assert owner.combine_sampled_and_draft_tokens(owner.input_ids, *empty.values(), 0, num_new).numel() == 0 + assert combine_sampled_and_draft_tokens(owner.input_ids, *empty.values(), 0, num_new).numel() == 0 + with pytest.raises(ValueError, match="target input capacity"): + owner.combine_sampled_and_draft_tokens(owner.input_ids, *device.values(), 33, num_new) + finally: + torch.npu.synchronize() + for graph in graphs: + graph.reset() diff --git a/tests/ut/worker/test_model_runner_v2.py b/tests/ut/worker/test_model_runner_v2.py index 762018e6f152..60045efcfe7d 100644 --- a/tests/ut/worker/test_model_runner_v2.py +++ b/tests/ut/worker/test_model_runner_v2.py @@ -1492,6 +1492,7 @@ def _prepare_inputs_runner(*, draft=False, full_cg=False, use_dcp=False, use_pp= runner.num_computed_tokens_event = MagicMock() runner.num_computed_tokens_cpu = torch.tensor([3, 4, 0, 0], dtype=torch.int32) runner.input_buffers = AscendInputBuffers(4, 16, runner.device) + runner._combine_sampled_and_draft_tokens = vllm_model_runner.combine_sampled_and_draft_tokens runner.input_buffers.dcp_local_seq_lens = torch.zeros(4, dtype=torch.int32) runner.req_states = SimpleNamespace( req_id_to_index={"r0": 0, "r1": 1}, @@ -1564,8 +1565,9 @@ def _partition(_pcp_manager, input_batch, **_kwargs): patch("vllm_ascend.worker.v2.model_runner.prepare_prefill_inputs", side_effect=prefill_inputs), patch("vllm_ascend.worker.v2.model_runner.prepare_pos_seq_lens"), patch("vllm_ascend.worker.v2.model_runner.prepare_dcp_local_seq_lens", create=True), - patch( - "vllm_ascend.worker.v2.model_runner.combine_sampled_and_draft_tokens", + patch.object( + runner, + "_combine_sampled_and_draft_tokens", return_value=torch.tensor([0, 1], dtype=torch.int32), side_effect=combine_tokens, ), diff --git a/vllm_ascend/worker/v2/input_batch.py b/vllm_ascend/worker/v2/input_batch.py index 6945361d2f20..3d026e980286 100644 --- a/vllm_ascend/worker/v2/input_batch.py +++ b/vllm_ascend/worker/v2/input_batch.py @@ -20,7 +20,12 @@ import numpy as np import torch -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.triton_utils import triton +from vllm.v1.worker.gpu.input_batch import ( + InputBatch, + InputBuffers, + _combine_sampled_and_draft_tokens_kernel, +) from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.ops.rotary_embedding import update_cos_sin @@ -67,6 +72,51 @@ def __init__( # CPP latency profiling must model one real prefill request. The generic # MRV2 dummy path otherwise describes every dummy batch as decode-only. self.profile_cpp: bool = False + self.logits_indices_workspace: torch.Tensor | None = None + + def bind_logits_indices_workspace(self) -> None: + """Give the target its own output storage before capture or execution.""" + if self.logits_indices_workspace is not None: + raise RuntimeError("Logits indices workspace is already bound") + self.logits_indices_workspace = torch.empty(self.max_num_tokens, dtype=torch.int64, device=self.device) + + def combine_sampled_and_draft_tokens( + self, + input_ids: torch.Tensor, + idx_mapping: torch.Tensor, + last_sampled_tokens: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + prefill_len: torch.Tensor, + draft_tokens: torch.Tensor, + cu_num_logits: torch.Tensor, + num_logits: int, + num_new_sampled_tokens: int = 1, + ) -> torch.Tensor: + """Write current values with Core's kernel into this target-owned prefix.""" + assert num_new_sampled_tokens in (0, 1), f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" + workspace = self.logits_indices_workspace + assert workspace is not None, "Logits indices workspace must be bound before use" + if input_ids.device != workspace.device: + raise ValueError("Logits indices workspace must share the input device") + if not 0 <= num_logits <= workspace.numel(): + raise ValueError("Logits indices exceed the target input capacity") + logits_indices = workspace[:num_logits] + _combine_sampled_and_draft_tokens_kernel[(idx_mapping.shape[0],)]( + input_ids, + idx_mapping, + last_sampled_tokens, + query_start_loc, + seq_lens, + prefill_len, + draft_tokens, + draft_tokens.stride(0), + cu_num_logits, + logits_indices, + NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens, + BLOCK_SIZE=triton.next_power_of_2(draft_tokens.shape[-1] + num_new_sampled_tokens), + ) + return logits_indices @dataclass diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 4e87cd1a550e..7768beb1b87c 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -247,6 +247,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): max_num_tokens=self.max_num_tokens, device=self.device, ) + self._combine_sampled_and_draft_tokens = combine_sampled_and_draft_tokens + if self.vllm_config.attention_config.hisparse_config is not None: + # Draft/PCP input providers retain their separate ownership. Sampling + # consumes this target prefix before the next input preparation. + self.input_buffers.bind_logits_indices_workspace() + self._combine_sampled_and_draft_tokens = self.input_buffers.combine_sampled_and_draft_tokens # Pinned D2H staging for corrected device state after spec rejection. # The authoritative host state is the shared NumPy/torch RequestState view. @@ -1034,7 +1040,7 @@ def _prepare_inputs_impl( # Some input token ids are directly read from the last sampled tokens # and draft tokens. Also, get the logits indices to sample tokens from. - logits_indices = combine_sampled_and_draft_tokens( + logits_indices = self._combine_sampled_and_draft_tokens( self.input_buffers.input_ids, idx_mapping, self.req_states.last_sampled_tokens, From 5a54eb466db539b0b3e893116258f7d261c8649f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Mon, 9 Nov 2026 18:22:26 +0800 Subject: [PATCH 075/120] refactor(hisparse): trim dead scratch and duplicate test logic Remove the unread BF16 length workspace and its test placeholder. Compare the logits workspace output and input-ID side effects directly with Core's existing implementation instead of maintaining a second CPU algorithm. Drop per-batch diagnostic log sets while retaining cold enable/prewarm logs and all completion, snapshot, ownership and invalidation guards. Focused Ruff, format, AST and diff checks passed across 40 Python files. The full format command remained in actionlint environment installation at its 55-second bound. No new NPU or model qualification is claimed. Signed-off-by: zhenwenqi2024 --- .../model_runner_v2/test_input_batch.py | 25 ++----------------- .../one_card/test_hisparse_attention.py | 1 - vllm_ascend/attention/sfa_hisparse.py | 3 +-- vllm_ascend/hisparse/mirror_overlap.py | 25 ++----------------- vllm_ascend/hisparse/recycled_invalidate.py | 10 -------- 5 files changed, 5 insertions(+), 59 deletions(-) diff --git a/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py b/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py index 60fa121fed39..b35ac4201c5b 100644 --- a/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py +++ b/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py @@ -17,24 +17,6 @@ def _require_npu(): torch.npu.set_device(0) -def _combine_inputs_cpu(input_ids, metadata, num_new): - """Independent value oracle, including the kernel's input-ID side effects.""" - result, indices = input_ids.clone(), [] - mapping, sampled, starts, lengths, prefill, drafts, cumulative = metadata.values() - for row, state in enumerate(mapping.tolist()): - count = int(cumulative[row + 1] - cumulative[row]) - end = int(starts[row + 1]) - indices.extend(range(end - count, end)) - if int(lengths[row]) <= int(prefill[state]): - continue - if num_new and int(lengths[row]) - count >= int(prefill[state]): - result[end - count] = sampled[state] - draft_count = count - num_new - if draft_count > 0: - result[end - draft_count : end] = drafts[state, :draft_count] - return result, torch.tensor(indices, dtype=torch.int64) - - @pytest.mark.parametrize("num_new", [0, 1]) @pytest.mark.parametrize("capture", [False, True]) def test_target_logits_workspace_matches_core_with_changing_inputs(num_new, capture): @@ -101,11 +83,8 @@ def update(change, counts, owner_index, owner, cpu, device): graph.replay() baseline_ids = initial.to("npu") baseline = combine_sampled_and_draft_tokens(baseline_ids, *device.values(), num_logits, num_new) - expected_ids, expected_indices = _combine_inputs_cpu(initial, cpu, num_new) - torch.testing.assert_close(output.cpu(), expected_indices, rtol=0, atol=0) - torch.testing.assert_close(baseline.cpu(), expected_indices, rtol=0, atol=0) - torch.testing.assert_close(owner.input_ids.cpu(), expected_ids, rtol=0, atol=0) - torch.testing.assert_close(baseline_ids.cpu(), expected_ids, rtol=0, atol=0) + torch.testing.assert_close(output.cpu(), baseline.cpu(), rtol=0, atol=0) + torch.testing.assert_close(owner.input_ids.cpu(), baseline_ids.cpu(), rtol=0, atol=0) assert output.data_ptr() == owner.logits_indices_workspace.data_ptr() assert torch.all(owner.logits_indices_workspace[num_logits:].cpu() == -777) empty = dict(device) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index cc2f1b5a3ac0..6dad6291d0f2 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -47,7 +47,6 @@ def test_sparse_attention_reference(num_queries, strided, capture): torch.arange(capacity * topk // block_size, dtype=torch.int32, device=device).view(capacity, -1), torch.arange(topk, dtype=torch.int32, device=device).view(1, 1, topk), torch.arange(1, capacity + 1, dtype=torch.int32, device=device), - torch.full((capacity,), topk, dtype=torch.int32, device=device), ) impl = SimpleNamespace(kv_lora_rank=latent_dim, qk_rope_head_dim=rope_dim, scale=scale, enable_sparse_sfa_c8=False) impl._workspace = lambda dtype: workspace diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index e0db162d827a..48f50978e048 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -188,7 +188,6 @@ def allocate(): ), torch.arange(self.topk, dtype=torch.int32, device=device).view(1, 1, self.topk), torch.arange(1, _ATTENTION_CHUNK_ROWS + 1, dtype=torch.int32, device=device), - torch.full((_ATTENTION_CHUNK_ROWS,), self.topk, dtype=torch.int32, device=device), ) return current_workspace_manager().get_persistent_resource(key, allocate) @@ -484,7 +483,7 @@ def _attention_physical( ) from vllm_ascend.hisparse.cache_ops import pack_attention_rows - latent, rope, block_table, indices, query_ends, seq_lens = self._workspace(query.dtype) + latent, rope, block_table, indices, query_ends = self._workspace(query.dtype) positions = indices.view(1, -1).to(torch.float32) outputs = [] for start in range(0, query.shape[0], _ATTENTION_CHUNK_ROWS): diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/mirror_overlap.py index e8db7472124e..157aee07e8cb 100644 --- a/vllm_ascend/hisparse/mirror_overlap.py +++ b/vllm_ascend/hisparse/mirror_overlap.py @@ -10,21 +10,15 @@ class SamplingMirrorOverlapConnector: """Let Target logits and sampling run before joining the Host mirror.""" - def __init__(self, connector, worker, output_copy_stream, logger, rank, snapshot=None): + def __init__(self, connector, worker, output_copy_stream, snapshot=None): self._connector = connector self._worker = worker self._output_copy_stream = output_copy_stream - self._logger = logger - self._rank = rank self._snapshot = snapshot self._overlap_draft = False self._pending_event = None self._allow_step = False self._eligible_forward = False - self._batch_size = 0 - self._logged_batch_sizes: set[int] = set() - self._logged_snapshot_sizes: set[int] = set() - self._logged_direct_sizes: set[int] = set() def __getattr__(self, name): return getattr(self._connector, name) @@ -50,8 +44,6 @@ def pre_forward(self, scheduler_output, input_batch=None, attn_metadata=None): and not input_batch.has_prefill and 0 < input_batch.num_tokens == input_batch.num_tokens_after_padding ) - if self._eligible_forward: - self._batch_size = input_batch.num_reqs return self._connector.pre_forward(scheduler_output, input_batch, attn_metadata=attn_metadata) def finish_forward(self) -> None: @@ -92,19 +84,6 @@ def finish_forward(self) -> None: self._overlap_draft = ( len(pending) == pending_before + 1 and pending[-1][0] is snapshot.submitted_completion ) - if self._overlap_draft and self._batch_size not in self._logged_snapshot_sizes: - self._logged_snapshot_sizes.add(self._batch_size) - self._logger.info( - "HiSparse Target snapshot used: rank=%d requests=%d", self._rank, self._batch_size - ) - if snapshot is not None and not self._overlap_draft and self._batch_size not in self._logged_direct_sizes: - self._logged_direct_sizes.add(self._batch_size) - self._logger.info( - "HiSparse Target snapshot fallback: rank=%d requests=%d", self._rank, self._batch_size - ) - if self._batch_size not in self._logged_batch_sizes: - self._logged_batch_sizes.add(self._batch_size) - self._logger.info("HiSparse MTP mirror overlap used: rank=%d requests=%d", self._rank, self._batch_size) else: worker.host_write_event.record(compute_stream) worker._release_completed_dma_descriptors() @@ -182,4 +161,4 @@ def overlap_mtp_host_mirror(runner): logger.info("HiSparse MTP mirror overlap enabled: rank=%d", rank) if snapshot is not None: logger.info("HiSparse Target snapshot enabled: rank=%d", rank) - return SamplingMirrorOverlapConnector(connector, worker, output_stream, logger, rank, snapshot) + return SamplingMirrorOverlapConnector(connector, worker, output_stream, snapshot) diff --git a/vllm_ascend/hisparse/recycled_invalidate.py b/vllm_ascend/hisparse/recycled_invalidate.py index 501356e6bd64..69a29532ac37 100644 --- a/vllm_ascend/hisparse/recycled_invalidate.py +++ b/vllm_ascend/hisparse/recycled_invalidate.py @@ -75,7 +75,6 @@ def __init__(self, tables, block_size): self._block_size = block_size self._stream = torch.accelerator.current_stream() self._pointer_table = async_tensor_h2d(list(self._addresses), self._device, torch.int64) - self._logged_rows: set[int] = set() @classmethod def bind(cls, worker): @@ -167,13 +166,4 @@ def try_call(self, worker, block_ids, indices) -> bool: # Repeating the last ID preserves membership without a sentinel. chunk.extend([chunk[-1]] * (_BLOCK_BUCKET - len(chunk))) self._launch(chunk, indices) - rows = indices.numel() - if rows not in self._logged_rows: - logger.info( - "HiSparse grouped recycled invalidation: rank=%d groups=%d active_rows=%d bucket=8", - self._rank, - len(self._tables), - rows, - ) - self._logged_rows.add(rows) return True From ad32f8430c030139b119ae73460bcfd91d6732ab Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 17:13:34 +0800 Subject: [PATCH 076/120] test(hisparse): initialize shared TopK in rebind fixture Signed-off-by: zhenwenqi2024 --- tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index c18b4ecf0e81..4681095197f7 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -368,6 +368,7 @@ def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind(monkeypatc kv_dtype=torch.float8_e4m3fn, registered_host_pool=NS(_hisparse_registration=object()), ) + group.shared_topk = NS(device_topk_rows=torch.empty((max_rows, runtime.max_union_rows), dtype=torch.int32)) group.leader = runtime def bind(num_pages): From 0271d60dbf5390427129d0d8c09ec776c3dc0b19 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 19:30:16 +0800 Subject: [PATCH 077/120] refactor(hisparse): consolidate host storage and mirror ownership Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 3 +- .../one_card/test_hisparse_prefill_staging.py | 2 +- .../test_hisparse_foundation_cpu.py | 7 +- vllm_ascend/hisparse/backend.py | 6 +- vllm_ascend/hisparse/cache_ops.py | 95 +--------- .../{mirror_overlap.py => host_mirror.py} | 172 ++++++++++++++++- vllm_ascend/hisparse/host_pool.py | 127 +++++++++++++ vllm_ascend/hisparse/target_snapshot.py | 174 ------------------ vllm_ascend/patch/__init__.py | 2 +- .../patch/worker/patch_hisparse_host_pool.py | 32 ---- vllm_ascend/worker/v2/model_runner.py | 2 +- 11 files changed, 312 insertions(+), 310 deletions(-) rename vllm_ascend/hisparse/{mirror_overlap.py => host_mirror.py} (50%) create mode 100644 vllm_ascend/hisparse/host_pool.py delete mode 100644 vllm_ascend/hisparse/target_snapshot.py delete mode 100644 vllm_ascend/patch/worker/patch_hisparse_host_pool.py diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index d23aa74de9c1..adab4001bc1b 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -9,12 +9,11 @@ from vllm_ascend.hisparse.backend import AscendHiSparseBackend, capture_aware_stream from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, - MappedHostTensor, ResolverWorkspace, - allocate_mapped_host, pack_attention_rows, scatter_packed_kv, ) +from vllm_ascend.hisparse.host_pool import MappedHostTensor, allocate_mapped_host @pytest.fixture(autouse=True) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 9372c3e3065e..022fd2b7ce19 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -12,10 +12,10 @@ from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, - allocate_mapped_host, gather_valid_resident_rows, scatter_packed_kv, ) +from vllm_ascend.hisparse.host_pool import allocate_mapped_host from vllm_ascend.hisparse.prefill import build_hisparse_prefill_staging_plan diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 4681095197f7..82d8923d2b57 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -19,6 +19,7 @@ from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.hisparse import cache_ops as hisparse_ops +from vllm_ascend.hisparse import host_pool from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.binding import ( _bind_indexer_views, @@ -90,7 +91,7 @@ def close(self): if len(events) == 1: raise RuntimeError("driver did not retire registration") - monkeypatch.setattr(hisparse_ops, "MappedHostTensor", Registration) + monkeypatch.setattr(host_pool, "MappedHostTensor", Registration) backend = AscendHiSparseBackend() backend.pin_tensor(tensor) registration = tensor._hisparse_registration @@ -107,7 +108,7 @@ def test_private_host_allocator_keeps_other_backends_on_core_path(monkeypatch): from vllm.v1.hisparse import backend as core_backend from vllm.v1.hisparse import runtime - from vllm_ascend.patch.worker.patch_hisparse_host_pool import install + from vllm_ascend.hisparse.host_pool import install class OtherBackend(core_backend.HiSparseBackend): def pin_tensor(self, tensor): @@ -131,7 +132,7 @@ def pin_tensor(self, tensor): def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): backing = torch.arange(128, dtype=torch.uint8) - mapping = hisparse_ops.MappedHostTensor.__new__(hisparse_ops.MappedHostTensor) + mapping = host_pool.MappedHostTensor.__new__(host_pool.MappedHostTensor) mapping.tensor = backing mapping.host_address = backing.data_ptr() mapping.nbytes = backing.numel() diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 88b2e762ea72..b11a3eefc689 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -60,9 +60,9 @@ def __init__(self) -> None: self._failed_dma_batches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] def check_available(self) -> None: + from vllm_ascend.hisparse.host_pool import install from vllm_ascend.hisparse.resolver import _load_resolver from vllm_ascend.hisparse.row_dma import install as install_row_dma - from vllm_ascend.patch.worker.patch_hisparse_host_pool import install from vllm_ascend.utils import bootstrap_custom_op_env install() @@ -83,7 +83,7 @@ def check_available(self) -> None: install_row_dma(AscendHiSparseBackend) def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Tensor]: - from vllm_ascend.hisparse.cache_ops import allocate_mapped_host + from vllm_ascend.hisparse.host_pool import allocate_mapped_host payload = allocate_mapped_host((size,), torch.int8) registration = payload._hisparse_registration @@ -92,7 +92,7 @@ def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Ten return registered[:size], registered def pin_tensor(self, tensor: torch.Tensor) -> None: - from vllm_ascend.hisparse.cache_ops import MappedHostTensor + from vllm_ascend.hisparse.host_pool import MappedHostTensor tensor._hisparse_registration = MappedHostTensor(tensor) diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 3a7e45f060b6..80d78c63c172 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -1,21 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared HiSparse cache operations and explicit mapped-host memory ownership. +"""HiSparse NPU cache operations and fixed resolver workspaces. -Mapping and scratch allocation happen at cache initialization. Decode planning -and payload copies stay on the NPU, including during ACL graph replay. +Decode planning and payload copies stay on the NPU, including during ACL graph +replay. Workspaces are owned by the bound cache operations. """ from __future__ import annotations -import ctypes import math -import mmap import torch from vllm.triton_utils import triton from vllm_ascend.device.device_config import is_950 +from vllm_ascend.hisparse.host_pool import MappedHostTensor from vllm_ascend.ops.triton.hisparse import ( gather_rows_kernel, gather_valid_resident_rows_kernel, @@ -26,95 +25,9 @@ scatter_packed_kv_kernel, ) -_ACL_HOST_REGISTER_MAPPED = 0 -_ACL_HAC_AIV = 3 _SORTED_WRITTEN_SLOT_COUNT = 8192 -class MappedHostTensor: - """Own an ACL mapping and retain the CPU allocation until explicit close. - - Close only after all streams and graphs using this mapping have completed - and been destroyed. No destructor unmaps memory behind asynchronous work. - """ - - def __init__(self, tensor: torch.Tensor) -> None: - if tensor.device.type != "cpu" or not tensor.is_contiguous(): - raise ValueError("HiSparse host mapping requires a contiguous CPU tensor") - self.tensor = tensor - self._lib = ctypes.CDLL("libascendcl.so") - self._lib.aclrtHostRegister.argtypes = [ - ctypes.c_void_p, - ctypes.c_uint64, - ctypes.c_int, - ctypes.POINTER(ctypes.c_void_p), - ] - self._lib.aclrtHostRegister.restype = ctypes.c_int - self._lib.aclrtHostUnregister.argtypes = [ctypes.c_void_p] - self._lib.aclrtHostUnregister.restype = ctypes.c_int - self._lib.aclrtHostMemMapCapabilities.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.POINTER(ctypes.c_int)] - self._lib.aclrtHostMemMapCapabilities.restype = ctypes.c_int - device_index = torch.npu.current_device() - self.device = torch.device("npu", device_index) - capability = ctypes.c_int() - status = self._lib.aclrtHostMemMapCapabilities(device_index, _ACL_HAC_AIV, ctypes.byref(capability)) - if status or not capability.value: - raise RuntimeError( - f"HiSparse requires AIV host-memory mapping: status={status}, capability={capability.value}" - ) - # The upstream host-pool allocator supplies an aligned subview of a - # larger backing allocation. Register that range, not its unaligned - # storage base; retained tensor ownership still retains the backing. - self.host_address = tensor.data_ptr() - self.nbytes = tensor.numel() * tensor.element_size() - if not self.nbytes: - raise ValueError("Cannot register an empty HiSparse host allocation") - mapped = ctypes.c_void_p() - status = self._lib.aclrtHostRegister( - self.host_address, self.nbytes, _ACL_HOST_REGISTER_MAPPED, ctypes.byref(mapped) - ) - if status: - raise RuntimeError(f"aclrtHostRegister for HiSparse failed: {status}") - self.device_ptr = int(mapped.value) - self.closed = False - - def address_of(self, tensor: torch.Tensor) -> int: - offset = tensor.data_ptr() - self.host_address - size = tensor.numel() * tensor.element_size() - if self.closed or tensor.device.type != "cpu" or not tensor.is_contiguous(): - raise ValueError("HiSparse source must be a contiguous view of a live mapped allocation") - if offset < 0 or offset + size > self.nbytes: - raise ValueError("HiSparse source is outside its registered host allocation") - return self.device_ptr + offset - - def close(self) -> None: - if not self.closed: - status = self._lib.aclrtHostUnregister(self.host_address) - if status: - raise RuntimeError(f"aclrtHostUnregister for HiSparse failed: {status}") - self.closed = True - - -def allocate_mapped_host(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: - """Allocate page-locked, aligned storage for both DMA and AIV access.""" - count = math.prod(shape) - nbytes = count * torch.empty((), dtype=dtype).element_size() - if nbytes <= 0: - raise ValueError("HiSparse host allocation must be nonempty") - allocated_bytes = triton.cdiv(nbytes, mmap.PAGESIZE) * mmap.PAGESIZE - # Mapping permits AIV access; DMA also requires a page-locked allocation. - # Align a view of PTA's pinned allocation without creating a second pool. - owner = torch.empty(allocated_bytes + mmap.PAGESIZE - 1, dtype=torch.uint8, pin_memory=True) - offset = (-owner.data_ptr()) % mmap.PAGESIZE - raw = owner[offset : offset + allocated_bytes] - mapping = MappedHostTensor(raw) - tensor = raw[:nbytes].view(dtype).view(shape) - # Both aliases retain the pinned allocation. Runtime owns the registration - # through this base tensor until its explicit, drained shutdown. - tensor._hisparse_registration = mapping - return tensor - - class ResolverWorkspace: """Fixed scratch shared by one index group on its ordered copy stream.""" diff --git a/vllm_ascend/hisparse/mirror_overlap.py b/vllm_ascend/hisparse/host_mirror.py similarity index 50% rename from vllm_ascend/hisparse/mirror_overlap.py rename to vllm_ascend/hisparse/host_mirror.py index 157aee07e8cb..abcd8b06b1c9 100644 --- a/vllm_ascend/hisparse/mirror_overlap.py +++ b/vllm_ascend/hisparse/host_mirror.py @@ -1,12 +1,181 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Overlap HiSparse Host mirroring with Target sampling and guarded MTP.""" +"""Own HiSparse Host snapshots and sampling/draft mirror completion fences.""" from __future__ import annotations +from dataclasses import dataclass + +import numpy as np import torch +def _host_disjoint(caches, draft_layers): + if any(cache.ndim != 2 or not cache.is_contiguous() or cache.numel() == 0 for cache in caches): + return False + spans = [(cache.data_ptr(), cache.data_ptr() + cache.numel() * cache.element_size()) for cache in caches] + return all( + left_end <= right_start or right_end <= left_start + for index, (left_start, left_end) in enumerate(spans) + if index in draft_layers + for other, (right_start, right_end) in enumerate(spans) + if other not in draft_layers + ) + + +def _draft_groups_disjoint(handles, draft_layers): + draft_runtimes = {id(handles[index].runtime) for index in draft_layers} + # A draft-led resolver can refill every follower from Host in one call. + for index in draft_layers: + group = handles[index].runtime.index_group + if id(group.leader) not in draft_runtimes or any( + id(follower) not in draft_runtimes for follower in group.followers + ): + return False + return True + + +def _index_plan(addresses, sizes, base, width, source_rows, limit): + if not len(addresses): + return None + first, alignment = np.divmod(addresses - base, width) + counts, size_remainder = np.divmod(sizes, width) + if ( + np.any(alignment) + or np.any(size_remainder) + or np.any(counts <= 0) + or np.any(first < 0) + or np.any(first + counts > source_rows) + ): + return None + ends = np.cumsum(counts) + total = int(ends[-1]) + if total > limit: + return None + starts = ends - counts + rows = np.repeat(first - starts, counts) + np.arange(total, dtype=np.int64) + return starts, rows + + +@dataclass +class _Scratch: + cpu_indices: torch.Tensor + device_indices: torch.Tensor + payload: torch.Tensor + ready: torch.Event + completion: torch.Event | None = None + in_flight: bool = False + + def available(self): + return not self.in_flight or (self.completion is not None and self.completion.query()) + + +class TargetMirrorSnapshot: + """Stage exact Target source bytes before draft can overwrite shared HMA.""" + + def __init__(self, worker, layers, width, capacity): + backing = worker.hot_backing + self._source = backing.view(torch.uint8).reshape(-1, width) + self._base = backing.data_ptr() + self._width = width + self._layers = layers + self._sources = worker.resident_caches + self._destinations = worker.host_caches + self._capacity = capacity + self._slots = tuple( + _Scratch( + torch.empty(capacity, dtype=torch.int64, pin_memory=True), + torch.empty(capacity, dtype=torch.int64, device=backing.device), + torch.empty((capacity, width), dtype=torch.uint8, device=backing.device), + torch.Event(), + ) + for _ in range(2) + ) + self.armed = False + self.submitted_completion = None + + @classmethod + def bind(cls, worker, speculator): + geometry = getattr(worker, "_ascend_row_dma_geometry", None) + draft = set(worker._draft_layers) + names = worker.cache_layer_names + expected = set(getattr(speculator, "draft_attn_layer_names", ())) & set(names) + if geometry is None or not draft or {names[i] for i in draft} != expected: + return None + layers = tuple(index for index in range(len(names)) if index not in draft) + if ( + not layers + or not _draft_groups_disjoint(worker.cache_handles, draft) + or not _host_disjoint(worker.host_caches, draft) + ): + return None + widths = {int(geometry.row_bytes[i]) for i in layers} + if len(widths) != 1: + return None + width = widths.pop() + backing = worker.hot_backing + base = backing.data_ptr() + if ( + width <= 0 + or backing.device.type != "npu" + or backing.dtype not in (torch.int8, torch.uint8) + or backing.ndim != 1 + or not backing.is_contiguous() + or backing.storage_offset() != 0 + or backing.numel() % width + or any( + cache.untyped_storage().data_ptr() != base + or (cache.data_ptr() - base) % width + or geometry.source_page_bytes[index] % width + for index in layers + for cache in (worker.resident_caches[index],) + ) + ): + return None + config = worker.vllm_config + capacity = config.scheduler_config.max_num_seqs * (config.num_lookahead_tokens + 1) * len(layers) + return cls(worker, layers, width, capacity) + + def try_submit(self, worker, descriptors, count, layer_indices, ready_event): + if ( + not self.armed + or self.submitted_completion is not None + or ready_event is not worker._forward_ready_event + or tuple(layer_indices) != self._layers + or worker.resident_caches is not self._sources + or worker.host_caches is not self._destinations + ): + return False + slot = next((item for item in self._slots if item.available()), None) + if slot is None: + return False + plan = _index_plan( + descriptors.src_np[:count], + descriptors.sizes_np[:count], + self._base, + self._width, + self._source.shape[0], + self._capacity, + ) + if plan is None: + return False + starts, rows = plan + size = len(rows) + # Owners remain pinned if any submission fails. Never retry a partial + # enqueue using the original device addresses. + slot.in_flight = True + slot.completion = None + slot.cpu_indices[:size].numpy()[:] = rows + slot.device_indices[:size].copy_(slot.cpu_indices[:size], non_blocking=True) + torch.index_select(self._source, 0, slot.device_indices[:size], out=slot.payload[:size]) + slot.ready.record() + descriptors.src_np[:count] = slot.payload.data_ptr() + starts * self._width + worker._submit_dma_descriptors(descriptors, count, ready_event=slot.ready) + slot.completion = worker._pending_dma_descriptors[-1][0] + self.submitted_completion = slot.completion + return True + + class SamplingMirrorOverlapConnector: """Let Target logits and sampling run before joining the Host mirror.""" @@ -118,7 +287,6 @@ def overlap_mtp_host_mirror(runner): from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector from vllm_ascend.hisparse.backend import AscendHiSparseBackend - from vllm_ascend.hisparse.target_snapshot import TargetMirrorSnapshot connector = runner.kv_connector config = runner.vllm_config.parallel_config diff --git a/vllm_ascend/hisparse/host_pool.py b/vllm_ascend/hisparse/host_pool.py new file mode 100644 index 000000000000..ba318e9ae172 --- /dev/null +++ b/vllm_ascend/hisparse/host_pool.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Own pinned HiSparse Host storage and explicit ACL registration. + +Core registers an aligned pageable allocation through ``pin_tensor``. AIV +mapping alone does not make that storage a valid asynchronous DMA destination +on Ascend. Reuse PTA's pinned allocator and Core's existing release ownership. +Remove the allocation adapter when Core exposes a backend Host-allocation hook. +""" + +from __future__ import annotations + +import ctypes +import math +import mmap +from functools import wraps + +import torch +from vllm.triton_utils import triton + +_ACL_HOST_REGISTER_MAPPED = 0 +_ACL_HAC_AIV = 3 + + +class MappedHostTensor: + """Own an ACL mapping and retain the CPU allocation until explicit close. + + Close only after all streams and graphs using this mapping have completed + and been destroyed. No destructor unmaps memory behind asynchronous work. + """ + + def __init__(self, tensor: torch.Tensor) -> None: + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError("HiSparse host mapping requires a contiguous CPU tensor") + self.tensor = tensor + self._lib = ctypes.CDLL("libascendcl.so") + self._lib.aclrtHostRegister.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_int, + ctypes.POINTER(ctypes.c_void_p), + ] + self._lib.aclrtHostRegister.restype = ctypes.c_int + self._lib.aclrtHostUnregister.argtypes = [ctypes.c_void_p] + self._lib.aclrtHostUnregister.restype = ctypes.c_int + self._lib.aclrtHostMemMapCapabilities.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.POINTER(ctypes.c_int)] + self._lib.aclrtHostMemMapCapabilities.restype = ctypes.c_int + device_index = torch.npu.current_device() + self.device = torch.device("npu", device_index) + capability = ctypes.c_int() + status = self._lib.aclrtHostMemMapCapabilities(device_index, _ACL_HAC_AIV, ctypes.byref(capability)) + if status or not capability.value: + raise RuntimeError( + f"HiSparse requires AIV host-memory mapping: status={status}, capability={capability.value}" + ) + # The upstream host-pool allocator supplies an aligned subview of a + # larger backing allocation. Register that range, not its unaligned + # storage base; retained tensor ownership still retains the backing. + self.host_address = tensor.data_ptr() + self.nbytes = tensor.numel() * tensor.element_size() + if not self.nbytes: + raise ValueError("Cannot register an empty HiSparse host allocation") + mapped = ctypes.c_void_p() + status = self._lib.aclrtHostRegister( + self.host_address, self.nbytes, _ACL_HOST_REGISTER_MAPPED, ctypes.byref(mapped) + ) + if status: + raise RuntimeError(f"aclrtHostRegister for HiSparse failed: {status}") + self.device_ptr = int(mapped.value) + self.closed = False + + def address_of(self, tensor: torch.Tensor) -> int: + offset = tensor.data_ptr() - self.host_address + size = tensor.numel() * tensor.element_size() + if self.closed or tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError("HiSparse source must be a contiguous view of a live mapped allocation") + if offset < 0 or offset + size > self.nbytes: + raise ValueError("HiSparse source is outside its registered host allocation") + return self.device_ptr + offset + + def close(self) -> None: + if not self.closed: + status = self._lib.aclrtHostUnregister(self.host_address) + if status: + raise RuntimeError(f"aclrtHostUnregister for HiSparse failed: {status}") + self.closed = True + + +def allocate_mapped_host(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + """Allocate page-locked, aligned storage for both DMA and AIV access.""" + count = math.prod(shape) + nbytes = count * torch.empty((), dtype=dtype).element_size() + if nbytes <= 0: + raise ValueError("HiSparse host allocation must be nonempty") + allocated_bytes = triton.cdiv(nbytes, mmap.PAGESIZE) * mmap.PAGESIZE + # Mapping permits AIV access; DMA also requires a page-locked allocation. + # Align a view of PTA's pinned allocation without creating a second pool. + owner = torch.empty(allocated_bytes + mmap.PAGESIZE - 1, dtype=torch.uint8, pin_memory=True) + offset = (-owner.data_ptr()) % mmap.PAGESIZE + raw = owner[offset : offset + allocated_bytes] + mapping = MappedHostTensor(raw) + tensor = raw[:nbytes].view(dtype).view(shape) + # Both aliases retain the pinned allocation. Runtime owns the registration + # through this base tensor until its explicit, drained shutdown. + tensor._hisparse_registration = mapping + return tensor + + +def install() -> None: + from vllm.v1.hisparse import runtime + from vllm.v1.hisparse.backend import get_hisparse_backend + + from vllm_ascend.hisparse.backend import AscendHiSparseBackend + + original = runtime.allocate_pinned_host_pool + if getattr(original, "_ascend_pinned_private_pool", False): + return + + @wraps(original) + def allocate(size): + backend = get_hisparse_backend() + if isinstance(backend, AscendHiSparseBackend): + return backend.allocate_private_host_pool(size) + return original(size) + + allocate._ascend_pinned_private_pool = True + runtime.allocate_pinned_host_pool = allocate diff --git a/vllm_ascend/hisparse/target_snapshot.py b/vllm_ascend/hisparse/target_snapshot.py deleted file mode 100644 index 7281b8c2c878..000000000000 --- a/vllm_ascend/hisparse/target_snapshot.py +++ /dev/null @@ -1,174 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Own Target row snapshots until asynchronous Host mirroring completes.""" - -from dataclasses import dataclass - -import numpy as np -import torch - - -def _host_disjoint(caches, draft_layers): - if any(cache.ndim != 2 or not cache.is_contiguous() or cache.numel() == 0 for cache in caches): - return False - spans = [(cache.data_ptr(), cache.data_ptr() + cache.numel() * cache.element_size()) for cache in caches] - return all( - left_end <= right_start or right_end <= left_start - for index, (left_start, left_end) in enumerate(spans) - if index in draft_layers - for other, (right_start, right_end) in enumerate(spans) - if other not in draft_layers - ) - - -def _draft_groups_disjoint(handles, draft_layers): - draft_runtimes = {id(handles[index].runtime) for index in draft_layers} - # A draft-led resolver can refill every follower from Host in one call. - for index in draft_layers: - group = handles[index].runtime.index_group - if id(group.leader) not in draft_runtimes or any( - id(follower) not in draft_runtimes for follower in group.followers - ): - return False - return True - - -def _index_plan(addresses, sizes, base, width, source_rows, limit): - if not len(addresses): - return None - first, alignment = np.divmod(addresses - base, width) - counts, size_remainder = np.divmod(sizes, width) - if ( - np.any(alignment) - or np.any(size_remainder) - or np.any(counts <= 0) - or np.any(first < 0) - or np.any(first + counts > source_rows) - ): - return None - ends = np.cumsum(counts) - total = int(ends[-1]) - if total > limit: - return None - starts = ends - counts - rows = np.repeat(first - starts, counts) + np.arange(total, dtype=np.int64) - return starts, rows - - -@dataclass -class _Scratch: - cpu_indices: torch.Tensor - device_indices: torch.Tensor - payload: torch.Tensor - ready: torch.Event - completion: torch.Event | None = None - in_flight: bool = False - - def available(self): - return not self.in_flight or (self.completion is not None and self.completion.query()) - - -class TargetMirrorSnapshot: - """Stage exact Target source bytes before draft can overwrite shared HMA.""" - - def __init__(self, worker, layers, width, capacity): - backing = worker.hot_backing - self._source = backing.view(torch.uint8).reshape(-1, width) - self._base = backing.data_ptr() - self._width = width - self._layers = layers - self._sources = worker.resident_caches - self._destinations = worker.host_caches - self._capacity = capacity - self._slots = tuple( - _Scratch( - torch.empty(capacity, dtype=torch.int64, pin_memory=True), - torch.empty(capacity, dtype=torch.int64, device=backing.device), - torch.empty((capacity, width), dtype=torch.uint8, device=backing.device), - torch.Event(), - ) - for _ in range(2) - ) - self.armed = False - self.submitted_completion = None - - @classmethod - def bind(cls, worker, speculator): - geometry = getattr(worker, "_ascend_row_dma_geometry", None) - draft = set(worker._draft_layers) - names = worker.cache_layer_names - expected = set(getattr(speculator, "draft_attn_layer_names", ())) & set(names) - if geometry is None or not draft or {names[i] for i in draft} != expected: - return None - layers = tuple(index for index in range(len(names)) if index not in draft) - if ( - not layers - or not _draft_groups_disjoint(worker.cache_handles, draft) - or not _host_disjoint(worker.host_caches, draft) - ): - return None - widths = {int(geometry.row_bytes[i]) for i in layers} - if len(widths) != 1: - return None - width = widths.pop() - backing = worker.hot_backing - base = backing.data_ptr() - if ( - width <= 0 - or backing.device.type != "npu" - or backing.dtype not in (torch.int8, torch.uint8) - or backing.ndim != 1 - or not backing.is_contiguous() - or backing.storage_offset() != 0 - or backing.numel() % width - or any( - cache.untyped_storage().data_ptr() != base - or (cache.data_ptr() - base) % width - or geometry.source_page_bytes[index] % width - for index in layers - for cache in (worker.resident_caches[index],) - ) - ): - return None - config = worker.vllm_config - capacity = config.scheduler_config.max_num_seqs * (config.num_lookahead_tokens + 1) * len(layers) - return cls(worker, layers, width, capacity) - - def try_submit(self, worker, descriptors, count, layer_indices, ready_event): - if ( - not self.armed - or self.submitted_completion is not None - or ready_event is not worker._forward_ready_event - or tuple(layer_indices) != self._layers - or worker.resident_caches is not self._sources - or worker.host_caches is not self._destinations - ): - return False - slot = next((item for item in self._slots if item.available()), None) - if slot is None: - return False - plan = _index_plan( - descriptors.src_np[:count], - descriptors.sizes_np[:count], - self._base, - self._width, - self._source.shape[0], - self._capacity, - ) - if plan is None: - return False - starts, rows = plan - size = len(rows) - # Owners remain pinned if any submission fails. Never retry a partial - # enqueue using the original device addresses. - slot.in_flight = True - slot.completion = None - slot.cpu_indices[:size].numpy()[:] = rows - slot.device_indices[:size].copy_(slot.cpu_indices[:size], non_blocking=True) - torch.index_select(self._source, 0, slot.device_indices[:size], out=slot.payload[:size]) - slot.ready.record() - descriptors.src_np[:count] = slot.payload.data_ptr() + starts * self._width - worker._submit_dma_descriptors(descriptors, count, ready_event=slot.ready) - slot.completion = worker._pending_dma_descriptors[-1][0] - self.submitted_completion = slot.completion - return True diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 29330ff48335..465493b5d496 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -28,7 +28,7 @@ # What's Patched and how it works: # -------------------------------- -# * worker/patch_hisparse_host_pool.py: +# * vllm_ascend/hisparse/host_pool.py (HiSparse backend initialization): # Target: vllm.v1.hisparse.runtime.allocate_pinned_host_pool. # Why: AIV mapping alone does not enable DMA to pageable Host storage. # How: Installed only by the HiSparse backend; allocate private pinned storage diff --git a/vllm_ascend/patch/worker/patch_hisparse_host_pool.py b/vllm_ascend/patch/worker/patch_hisparse_host_pool.py deleted file mode 100644 index b8432769d2a6..000000000000 --- a/vllm_ascend/patch/worker/patch_hisparse_host_pool.py +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Allocate pinned private HiSparse storage before Core creates cache views. - -Core registers an aligned pageable allocation through ``pin_tensor``. AIV -mapping alone does not make that storage a valid asynchronous DMA destination -on Ascend. Reuse PTA's pinned allocator and Core's existing release ownership. -Remove this adapter when Core exposes a backend Host-allocation hook. -""" - -from functools import wraps - - -def install() -> None: - from vllm.v1.hisparse import runtime - from vllm.v1.hisparse.backend import get_hisparse_backend - - from vllm_ascend.hisparse.backend import AscendHiSparseBackend - - original = runtime.allocate_pinned_host_pool - if getattr(original, "_ascend_pinned_private_pool", False): - return - - @wraps(original) - def allocate(size): - backend = get_hisparse_backend() - if isinstance(backend, AscendHiSparseBackend): - return backend.allocate_private_host_pool(size) - return original(size) - - allocate._ascend_pinned_private_pool = True - runtime.allocate_pinned_host_pool = allocate diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 7768beb1b87c..c371ec78b4c8 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -534,7 +534,7 @@ def initialize_kv_cache( self._bind_hisparse_direct_prolog() # Bind after Core seals all cache and stream owners. - from vllm_ascend.hisparse.mirror_overlap import overlap_mtp_host_mirror + from vllm_ascend.hisparse.host_mirror import overlap_mtp_host_mirror overlap_connector = overlap_mtp_host_mirror(self) self._sampling_mirror = overlap_connector if overlap_connector is not self.kv_connector else None From 32ae6955d334004674cbdc2a6738e51185b47d0f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 20:19:28 +0800 Subject: [PATCH 078/120] fix(hisparse): include binding workspaces in Core memory profiling Prepare the existing Ascend resolver storage during profile_run and consume it at KV binding. Reuse Core non-KV memory accounting for HiSparse budgets and startup suggestions while preserving standalone bindings and C8 length owners across rebinds. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 69 ++++++++++++++++++++++----- vllm_ascend/worker/v2/model_runner.py | 8 ++++ vllm_ascend/worker/worker.py | 13 +++-- 3 files changed, 74 insertions(+), 16 deletions(-) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index b11a3eefc689..dcbef6f0ec8c 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -5,6 +5,7 @@ from __future__ import annotations import importlib +from collections.abc import Iterable from typing import TYPE_CHECKING, Any import torch @@ -16,6 +17,51 @@ from vllm.v1.kv_cache_interface import KVCacheConfig +def _allocate_resolver_workspace(runtime: HiSparseRuntime): + from vllm_ascend.hisparse.cache_ops import ResolverWorkspace + + return ResolverWorkspace( + max_rows=runtime.max_swap_rows, + hot_size=runtime.region_stride, + max_union_rows=runtime.max_union_rows, + device=runtime.device, + max_topk=runtime.index_group.shared_topk.device_topk_rows.shape[1], + ) + + +def _allocate_attention_lengths(runtime: HiSparseRuntime) -> tuple[torch.Tensor, torch.Tensor] | None: + if runtime.kv_dtype != torch.float8_e4m3fn: + return None + return ( + torch.arange(runtime.max_swap_rows + 1, dtype=torch.int32, device=runtime.device), + torch.empty((1,), dtype=torch.int32, device=runtime.device), + ) + + +def prepare_hisparse_workspaces(runtimes: Iterable[HiSparseRuntime]) -> None: + """Create real next-binding storage inside the existing profile scope. + + Core fixes group geometry when loading the model. Retain one scratch owner + per actual group, and fresh C8 lengths for each consumed binding generation. + """ + from vllm.v1.worker.workspace import current_workspace_manager + + manager = current_workspace_manager() + for runtime in runtimes: + group = runtime.index_group + workspace = getattr(group, "_ascend_profiled_resolver", None) + if workspace is None: + # Identity, not shape, separates groups. The manager key does not + # retain the runtime or any Host-pool/Graph owner. + key = ("ascend_hisparse_resolver", object()) + workspace = manager.get_persistent_resource( + key, lambda runtime=runtime: _allocate_resolver_workspace(runtime) + ) + group._ascend_profiled_resolver = workspace + if not hasattr(runtime, "_ascend_profiled_cache_ops"): + runtime._ascend_profiled_cache_ops = workspace, _allocate_attention_lengths(runtime) + + class _CaptureAwareStream(torch.Stream): """Propagate a captured dependency before PTA allocates CANN workspace. @@ -102,20 +148,15 @@ def unpin_tensor(self, tensor: torch.Tensor) -> None: del tensor._hisparse_registration def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: - from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps, ResolverWorkspace + from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps group = runtime.index_group # Runtime copied the model group's stream reference at layer creation. # Adapt the live runtime binding before any graph captures use it. group.copy_stream = capture_aware_stream(group.copy_stream) + prepared = getattr(runtime, "_ascend_profiled_cache_ops", None) if runtime is group.leader: - workspace = ResolverWorkspace( - max_rows=runtime.max_swap_rows, - hot_size=runtime.region_stride, - max_union_rows=runtime.max_union_rows, - device=runtime.device, - max_topk=group.shared_topk.device_topk_rows.shape[1], - ) + workspace = prepared[0] if prepared is not None else _allocate_resolver_workspace(runtime) else: leader_ops = group.leader.cache_ops if not isinstance(leader_ops, HiSparseNPUOps): @@ -145,17 +186,19 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: # values describe immutable geometry, not the current valid rows. # Prepare them before capture and retain them with this binding; # a rebind gets its own table and current physical KV span. - attention_lengths = ( - torch.arange(workspace.max_rows + 1, dtype=torch.int32, device=runtime.device), - torch.full((1,), attention_view.shape[0], dtype=torch.int32, device=runtime.device), - ) - return HiSparseNPUOps( + attention_lengths = prepared[1] if prepared is not None else _allocate_attention_lengths(runtime) + assert attention_lengths is not None + attention_lengths[1].fill_(attention_view.shape[0]) + ops = HiSparseNPUOps( mapping, workspace, attention_capacity=attention_capacity, attention_block_stride=attention_block_stride, attention_lengths=attention_lengths, ) + if prepared is not None: + del runtime._ascend_profiled_cache_ops + return ops def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: from .binding import get_hisparse_indexer_specs diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index c371ec78b4c8..512e29362c9f 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -796,6 +796,14 @@ def profile_run(self, randomize_inputs: bool = False) -> None: necessary HCCL buffer for the MC2 operator before standard `profile_run`. Additionally, we set override_mrv2_in_profile_run to True to force moe load to be balanced when executing `profile_run` """ + if self.vllm_config.attention_config.hisparse_config is not None: + from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces + + prepare_hisparse_workspaces( + cache.runtime + for layer in self.compilation_config.static_forward_context.values() + if (cache := getattr(layer, "hisparse_cache", None)) is not None + ) sparse_cfg = self.ascend_config.sparse_kv_offload_config if sparse_cfg.enabled: allocate_kv_offload_topk_profile_buffers(self.get_kv_cache_spec(), self.vllm_config, sparse_cfg) diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 83c91f6806a7..82064f324b19 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -723,9 +723,14 @@ def determine_available_memory(self) -> int: # Override torch_peak_increase with the pre-graph-capture value to # avoid double-counting graph pool memory as activation memory. profile_result.torch_peak_increase = profile_torch_peak - profile_result.before_profile.torch_peak - profile_result.non_kv_cache_memory = ( - profile_result.non_torch_increase + profile_result.torch_peak_increase + profile_result.weights_memory - ) + if self.vllm_config.attention_config.hisparse_config is None: + profile_result.non_kv_cache_memory = ( + profile_result.non_torch_increase + profile_result.torch_peak_increase + profile_result.weights_memory + ) + else: + # Reuse Core's measured total, including owners retained from an + # earlier profile. Preserve this authority for startup suggestions. + self._hisparse_profiled_non_kv_memory = profile_result.non_kv_cache_memory # Save per-category memory for use in compile_or_warm_up_model() (step 5). self.peak_activation_memory = profile_result.torch_peak_increase @@ -1048,6 +1053,8 @@ def compile_or_warm_up_model(self) -> CompilationTimes: + self.non_torch_memory + npugraph_memory_bytes ) + if self.vllm_config.attention_config.hisparse_config is not None: + non_kv_memory = self._hisparse_profiled_non_kv_memory + npugraph_memory_bytes self.npugraph_memory_bytes = npugraph_memory_bytes suggested_to_requested = int(self.requested_memory) - non_kv_memory - redundancy_buffer suggested_to_gpu_limit = int(self.init_snapshot.free_memory) - non_kv_memory - redundancy_buffer From c366b41c33551a3885fe9c4d33af203cae48b5dd Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 20:51:24 +0800 Subject: [PATCH 079/120] fix(hisparse): allocate profiled workspaces outside inference mode Keep persistent workspace and C8 length tensors writable when KV binding runs outside the profiling inference context. Disable inference mode only while preparing the original owners, without cloning storage or changing bind and graph lifetimes. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index dcbef6f0ec8c..b2b224f46de0 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -38,6 +38,7 @@ def _allocate_attention_lengths(runtime: HiSparseRuntime) -> tuple[torch.Tensor, ) +@torch.inference_mode(False) def prepare_hisparse_workspaces(runtimes: Iterable[HiSparseRuntime]) -> None: """Create real next-binding storage inside the existing profile scope. From 6bab22a292d38210ac1609fa2895712998d183a9 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 21:12:47 +0800 Subject: [PATCH 080/120] test(hisparse): cover profiled workspace owners and rebinds Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 97 +++++++++++++++++++ .../test_hisparse_foundation_cpu.py | 48 ++++++++- 2 files changed, 142 insertions(+), 3 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 022fd2b7ce19..a401dc3f98d6 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -189,3 +189,100 @@ def resolve(): ) torch.testing.assert_close(gpu_rows.cpu().flatten(), expected, rtol=0, atol=0) torch.testing.assert_close(mask.cpu().flatten(), (expected < 0).int(), rtol=0, atol=0) + + +@pytest.mark.parametrize("pages,stride_rows", [(2, 128), (3, 136)]) +def test_profiled_workspace_rebind_preserves_prior_length_graphs(pages, stride_rows): + from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + + from vllm_ascend.hisparse.backend import capture_aware_stream + from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces as actual_prepare + + prepare_hisparse_workspaces = torch.inference_mode()(actual_prepare) + + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse mapped-host owners currently target Ascend 950") + torch.npu.set_device(0) + device = torch.device("npu:0") + reset_workspace_manager() + init_workspace_manager(device) + backend = AscendHiSparseBackend() + backend.check_available() + host = allocate_mapped_host((4096,), torch.uint8) + registration = host._hisparse_registration + graphs = [] + + def runtime(): + group = SimpleNamespace( + copy_stream=capture_aware_stream(torch.Stream(device=device)), + shared_topk=SimpleNamespace(device_topk_rows=torch.empty((4, 64), dtype=torch.int32, device=device)), + ) + owner = SimpleNamespace( + index_group=group, + max_swap_rows=4, + region_stride=128, + max_union_rows=64, + device=device, + kv_dtype=torch.float8_e4m3fn, + registered_host_pool=host, + ) + group.leader = owner + return owner + + def bind(owner, count): + raw = torch.empty(count * stride_rows * 656, dtype=torch.float8_e4m3fn, device=device) + cache = raw.as_strided((count, 128, 656), (stride_rows * 656, 656, 1)) + owner.hot = SimpleNamespace(cache=cache, attention_cache=cache, block_size=128) + owner.cache_ops = backend.create_cache_ops(owner) + return owner.cache_ops + + def capture_lengths(ops): + query, key = ops.attention_lengths + q_out, k_out = torch.empty_like(query), torch.empty_like(key) + q_out.copy_(query) + k_out.copy_(key) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + graphs.append(graph) + with torch.npu.graph(graph): + q_out.copy_(query) + k_out.copy_(key) + return graph, q_out, k_out + + try: + first = runtime() + prepare_hisparse_workspaces([first]) + pending = first._ascend_profiled_cache_ops + workspace = pending[0] + prepare_hisparse_workspaces([first]) + assert first._ascend_profiled_cache_ops is pending + for tensor in (workspace.first_refs, workspace.errors, *pending[1]): + assert not torch.is_inference(tensor) + old = bind(first, pages) + assert old.workspace is workspace and not hasattr(first, "_ascend_profiled_cache_ops") + captured = [capture_lengths(old)] + prepare_hisparse_workspaces([first]) + new = bind(first, pages + 1) + assert new.workspace is workspace + for tensor in new.attention_lengths: + assert not torch.is_inference(tensor) + assert new.attention_lengths[1].data_ptr() != old.attention_lengths[1].data_ptr() + captured.append(capture_lengths(new)) + prepare_hisparse_workspaces([first]) + third = bind(first, pages + 2) + assert third.workspace is workspace + for tensor in third.attention_lengths: + assert not torch.is_inference(tensor) + for index, (graph, q_out, k_out) in enumerate(captured): + graph.replay() + torch.npu.synchronize() + torch.testing.assert_close(q_out.cpu(), torch.arange(5, dtype=torch.int32)) + torch.testing.assert_close( + k_out.cpu(), torch.tensor([(pages + index - 1) * stride_rows + 128], dtype=torch.int32) + ) + finally: + torch.npu.synchronize() + for graph in graphs: + graph.reset() + registration.close() + reset_workspace_manager() diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 82d8923d2b57..ec6ec0d14df4 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -35,6 +35,21 @@ from vllm_ascend.ops.triton import sfa_indexer_store +@pytest.fixture(params=[False, True], ids=["standalone", "profiled"]) +def prepare_workspaces(request): + from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + + from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces + + reset_workspace_manager() + try: + if request.param: + init_workspace_manager(torch.device("cpu")) + yield torch.inference_mode()(prepare_hisparse_workspaces) if request.param else None + finally: + reset_workspace_manager() + + @pytest.mark.parametrize("mxfp8", [False, True]) def test_c8_prolog_requires_matching_fp8_weight_scheme(mxfp8): from vllm_ascend.attention.utils import PreprocessType @@ -150,7 +165,7 @@ def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): @pytest.mark.parametrize("max_rows,hot_rows,union_rows", [(1, 16, 8), (9, 48, 24), (37, 128, 72)]) def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( - monkeypatch, max_rows, hot_rows, union_rows + monkeypatch, max_rows, hot_rows, union_rows, prepare_workspaces ): from vllm_ascend.hisparse import resolver @@ -175,7 +190,15 @@ def group(): registered_host_pool=NS(_hisparse_registration=object()), ) owner.leader = leader + if prepare_workspaces is not None: + prepare_workspaces([leader, leader]) + pending = leader._ascend_profiled_cache_ops + prepare_workspaces([leader]) + assert leader._ascend_profiled_cache_ops is pending leader.cache_ops = backend.create_cache_ops(leader) + if prepare_workspaces is not None: + assert leader.cache_ops.workspace is pending[0] + assert not hasattr(leader, "_ascend_profiled_cache_ops") follower = NS(**vars(leader)) return leader, backend.create_cache_ops(follower) @@ -188,6 +211,9 @@ def group(): assert first_workspace.first_refs.shape == (max_rows * union_rows,) assert first_workspace.errors.shape == (max_rows, 8) assert first_workspace.max_union_rows == union_rows + assert not torch.is_inference(first_workspace.first_refs) + assert not torch.is_inference(first_workspace.errors) + first_workspace.errors.fill_(17) first_workspace.first_refs.fill_(13) second_workspace.first_refs.zero_() assert torch.all(follower.workspace.first_refs == 13) @@ -354,7 +380,9 @@ def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): @pytest.mark.parametrize("pages,stride_rows,max_rows", [(2, 128, 9), (3, 136, 37)]) -def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind(monkeypatch, pages, stride_rows, max_rows): +def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind( + monkeypatch, pages, stride_rows, max_rows, prepare_workspaces +): from vllm_ascend.hisparse import resolver monkeypatch.setattr(resolver, "_load_resolver", lambda: NS(hisparse_resolver_scratch_capacity=lambda: 175616)) @@ -376,15 +404,29 @@ def bind(num_pages): backing = torch.empty(num_pages * stride_rows * 656, dtype=torch.float8_e4m3fn) cache = backing.as_strided((num_pages, 128, 656), (stride_rows * 656, 656, 1)) runtime.hot = NS(cache=cache, attention_cache=cache, block_size=128) - return backend.create_cache_ops(runtime) + if prepare_workspaces is not None: + prepare_workspaces([runtime]) + pending = runtime._ascend_profiled_cache_ops + prepare_workspaces([runtime]) + assert runtime._ascend_profiled_cache_ops is pending + ops = backend.create_cache_ops(runtime) + if prepare_workspaces is not None: + assert ops.workspace is pending[0] + assert not hasattr(runtime, "_ascend_profiled_cache_ops") + return ops first = bind(pages) query, key = first.attention_lengths + for tensor in (first.workspace.first_refs, first.workspace.errors, query, key): + assert not torch.is_inference(tensor) + first.workspace.errors.fill_(17) for rows in (1, max_rows, 3, 1): torch.testing.assert_close(query[rows : rows + 1], torch.tensor([rows], dtype=torch.int32)) expected = (pages - 1) * stride_rows + 128 torch.testing.assert_close(key, torch.tensor([expected], dtype=torch.int32)) second = bind(pages + 1) + if prepare_workspaces is not None: + assert second.workspace is first.workspace torch.testing.assert_close(second.attention_lengths[1], torch.tensor([expected + stride_rows], dtype=torch.int32)) torch.testing.assert_close(key, torch.tensor([expected], dtype=torch.int32)) assert key.data_ptr() != second.attention_lengths[1].data_ptr() From 29b9f33d297bbc5c406f9fd58a7f1e26e1106739 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 22:24:47 +0800 Subject: [PATCH 081/120] refactor(worker): reuse Core token combination Remove the HiSparse-gated logits-index workspace and private kernel wrapper. Call the existing public Core function directly, remove workspace-only E2E coverage, and retarget existing input preparation tests. Signed-off-by: zhenwenqi2024 --- .../model_runner_v2/test_input_batch.py | 102 ------------------ tests/ut/worker/test_model_runner_v2.py | 6 +- vllm_ascend/worker/v2/input_batch.py | 52 +-------- vllm_ascend/worker/v2/model_runner.py | 8 +- 4 files changed, 4 insertions(+), 164 deletions(-) delete mode 100644 tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py diff --git a/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py b/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py deleted file mode 100644 index b35ac4201c5b..000000000000 --- a/tests/e2e/pull_request/one_card/model_runner_v2/test_input_batch.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Real-NPU input preparation tests for the V2 model runner.""" - -import pytest -import torch -import torch_npu # noqa: F401 -from vllm.v1.worker.gpu.input_batch import combine_sampled_and_draft_tokens - -from vllm_ascend.worker.v2.input_batch import AscendInputBuffers - - -@pytest.fixture(autouse=True) -def _require_npu(): - if not torch.npu.is_available(): - pytest.skip("NPU input preparation requires an available NPU") - torch.npu.set_device(0) - - -@pytest.mark.parametrize("num_new", [0, 1]) -@pytest.mark.parametrize("capture", [False, True]) -def test_target_logits_workspace_matches_core_with_changing_inputs(num_new, capture): - owners = [AscendInputBuffers(3, 32, torch.device("npu:0")) for _ in range(2)] - for owner in owners: - owner.bind_logits_indices_workspace() - assert owners[0].logits_indices_workspace.data_ptr() != owners[1].logits_indices_workspace.data_ptr() - count_sets = ( - (((1, 4), (4, 1), (2, 3)), ((1, 1),) * 3) - if num_new - else ( - ((0, 3), (3, 0), (1, 2)), - ((0, 1), (1, 0), (0, 1)), - ) - ) - graphs, retained_metadata = [], [] - try: - for owner_index, owner in enumerate(owners): - cpu = { - "mapping": torch.tensor([2, 0], dtype=torch.int32), - "sampled": torch.empty(3, dtype=torch.int32), - "starts": torch.tensor([0, 4, 8], dtype=torch.int32), - "lengths": torch.tensor([16, 4], dtype=torch.int32), - "prefill": torch.tensor([4, 20, 9], dtype=torch.int32), - "drafts": torch.empty((3, 3), dtype=torch.int32), - "cumulative": torch.empty(3, dtype=torch.int32), - } - device = {name: value.to("npu") for name, value in cpu.items()} - retained_metadata.append((cpu, device)) - - def update(change, counts, owner_index, owner, cpu, device): - salt = owner_index * 100 + change * 10 - cpu["mapping"].copy_(torch.tensor([2, 0] if change % 2 == 0 else [0, 2], dtype=torch.int32)) - cpu["sampled"].copy_(torch.arange(3, dtype=torch.int32) + 1000 + salt) - cpu["drafts"].copy_((torch.arange(9, dtype=torch.int32) + 2000 + salt).reshape(3, 3)) - cpu["cumulative"].copy_(torch.tensor([0, counts[0], sum(counts)], dtype=torch.int32)) - for name, value in cpu.items(): - device[name].copy_(value) - initial = torch.arange(32, dtype=torch.int32) + 70 + salt - owner.input_ids.copy_(initial) - owner.logits_indices_workspace.fill_(-777) - return initial - - for counts_sequence in count_sets: - num_logits = sum(counts_sequence[0]) - update(0, counts_sequence[0], owner_index, owner, cpu, device) - output = owner.combine_sampled_and_draft_tokens(owner.input_ids, *device.values(), num_logits, num_new) - torch.npu.synchronize() - graph = None - if capture: - graph = torch.npu.NPUGraph() - graphs.append(graph) - with torch.npu.graph(graph): - output = owner.combine_sampled_and_draft_tokens( - owner.input_ids, *device.values(), num_logits, num_new - ) - for change, counts in enumerate(counts_sequence): - initial = update(change, counts, owner_index, owner, cpu, device) - if graph is None: - output = owner.combine_sampled_and_draft_tokens( - owner.input_ids, *device.values(), num_logits, num_new - ) - else: - graph.replay() - baseline_ids = initial.to("npu") - baseline = combine_sampled_and_draft_tokens(baseline_ids, *device.values(), num_logits, num_new) - torch.testing.assert_close(output.cpu(), baseline.cpu(), rtol=0, atol=0) - torch.testing.assert_close(owner.input_ids.cpu(), baseline_ids.cpu(), rtol=0, atol=0) - assert output.data_ptr() == owner.logits_indices_workspace.data_ptr() - assert torch.all(owner.logits_indices_workspace[num_logits:].cpu() == -777) - empty = dict(device) - empty["mapping"] = device["mapping"][:0] - empty["starts"] = device["starts"][:1] - empty["lengths"] = device["lengths"][:0] - empty["cumulative"] = device["cumulative"][:1] - assert owner.combine_sampled_and_draft_tokens(owner.input_ids, *empty.values(), 0, num_new).numel() == 0 - assert combine_sampled_and_draft_tokens(owner.input_ids, *empty.values(), 0, num_new).numel() == 0 - with pytest.raises(ValueError, match="target input capacity"): - owner.combine_sampled_and_draft_tokens(owner.input_ids, *device.values(), 33, num_new) - finally: - torch.npu.synchronize() - for graph in graphs: - graph.reset() diff --git a/tests/ut/worker/test_model_runner_v2.py b/tests/ut/worker/test_model_runner_v2.py index 60045efcfe7d..762018e6f152 100644 --- a/tests/ut/worker/test_model_runner_v2.py +++ b/tests/ut/worker/test_model_runner_v2.py @@ -1492,7 +1492,6 @@ def _prepare_inputs_runner(*, draft=False, full_cg=False, use_dcp=False, use_pp= runner.num_computed_tokens_event = MagicMock() runner.num_computed_tokens_cpu = torch.tensor([3, 4, 0, 0], dtype=torch.int32) runner.input_buffers = AscendInputBuffers(4, 16, runner.device) - runner._combine_sampled_and_draft_tokens = vllm_model_runner.combine_sampled_and_draft_tokens runner.input_buffers.dcp_local_seq_lens = torch.zeros(4, dtype=torch.int32) runner.req_states = SimpleNamespace( req_id_to_index={"r0": 0, "r1": 1}, @@ -1565,9 +1564,8 @@ def _partition(_pcp_manager, input_batch, **_kwargs): patch("vllm_ascend.worker.v2.model_runner.prepare_prefill_inputs", side_effect=prefill_inputs), patch("vllm_ascend.worker.v2.model_runner.prepare_pos_seq_lens"), patch("vllm_ascend.worker.v2.model_runner.prepare_dcp_local_seq_lens", create=True), - patch.object( - runner, - "_combine_sampled_and_draft_tokens", + patch( + "vllm_ascend.worker.v2.model_runner.combine_sampled_and_draft_tokens", return_value=torch.tensor([0, 1], dtype=torch.int32), side_effect=combine_tokens, ), diff --git a/vllm_ascend/worker/v2/input_batch.py b/vllm_ascend/worker/v2/input_batch.py index 3d026e980286..6945361d2f20 100644 --- a/vllm_ascend/worker/v2/input_batch.py +++ b/vllm_ascend/worker/v2/input_batch.py @@ -20,12 +20,7 @@ import numpy as np import torch -from vllm.triton_utils import triton -from vllm.v1.worker.gpu.input_batch import ( - InputBatch, - InputBuffers, - _combine_sampled_and_draft_tokens_kernel, -) +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.ops.rotary_embedding import update_cos_sin @@ -72,51 +67,6 @@ def __init__( # CPP latency profiling must model one real prefill request. The generic # MRV2 dummy path otherwise describes every dummy batch as decode-only. self.profile_cpp: bool = False - self.logits_indices_workspace: torch.Tensor | None = None - - def bind_logits_indices_workspace(self) -> None: - """Give the target its own output storage before capture or execution.""" - if self.logits_indices_workspace is not None: - raise RuntimeError("Logits indices workspace is already bound") - self.logits_indices_workspace = torch.empty(self.max_num_tokens, dtype=torch.int64, device=self.device) - - def combine_sampled_and_draft_tokens( - self, - input_ids: torch.Tensor, - idx_mapping: torch.Tensor, - last_sampled_tokens: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - prefill_len: torch.Tensor, - draft_tokens: torch.Tensor, - cu_num_logits: torch.Tensor, - num_logits: int, - num_new_sampled_tokens: int = 1, - ) -> torch.Tensor: - """Write current values with Core's kernel into this target-owned prefix.""" - assert num_new_sampled_tokens in (0, 1), f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" - workspace = self.logits_indices_workspace - assert workspace is not None, "Logits indices workspace must be bound before use" - if input_ids.device != workspace.device: - raise ValueError("Logits indices workspace must share the input device") - if not 0 <= num_logits <= workspace.numel(): - raise ValueError("Logits indices exceed the target input capacity") - logits_indices = workspace[:num_logits] - _combine_sampled_and_draft_tokens_kernel[(idx_mapping.shape[0],)]( - input_ids, - idx_mapping, - last_sampled_tokens, - query_start_loc, - seq_lens, - prefill_len, - draft_tokens, - draft_tokens.stride(0), - cu_num_logits, - logits_indices, - NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens, - BLOCK_SIZE=triton.next_power_of_2(draft_tokens.shape[-1] + num_new_sampled_tokens), - ) - return logits_indices @dataclass diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 512e29362c9f..d558b64a5ebf 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -247,12 +247,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): max_num_tokens=self.max_num_tokens, device=self.device, ) - self._combine_sampled_and_draft_tokens = combine_sampled_and_draft_tokens - if self.vllm_config.attention_config.hisparse_config is not None: - # Draft/PCP input providers retain their separate ownership. Sampling - # consumes this target prefix before the next input preparation. - self.input_buffers.bind_logits_indices_workspace() - self._combine_sampled_and_draft_tokens = self.input_buffers.combine_sampled_and_draft_tokens # Pinned D2H staging for corrected device state after spec rejection. # The authoritative host state is the shared NumPy/torch RequestState view. @@ -1048,7 +1042,7 @@ def _prepare_inputs_impl( # Some input token ids are directly read from the last sampled tokens # and draft tokens. Also, get the logits indices to sample tokens from. - logits_indices = self._combine_sampled_and_draft_tokens( + logits_indices = combine_sampled_and_draft_tokens( self.input_buffers.input_ids, idx_mapping, self.req_states.last_sampled_tokens, From 7565876e2cd6d9a84d2d39c0e96a0a7022f9a0ff Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 22:36:28 +0800 Subject: [PATCH 082/120] refactor(hisparse): restore Core mirror completion flow Remove the custom sampling/draft mirror overlap and Target snapshots. Historical component source-release timing does not establish an isolated, repeatable model latency benefit. Keep batched DMA descriptor preparation and let Core finish_forward own the write-completion wait. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/host_mirror.py | 332 -------------------------- vllm_ascend/hisparse/row_dma.py | 4 +- vllm_ascend/worker/v2/model_runner.py | 18 +- 3 files changed, 2 insertions(+), 352 deletions(-) delete mode 100644 vllm_ascend/hisparse/host_mirror.py diff --git a/vllm_ascend/hisparse/host_mirror.py b/vllm_ascend/hisparse/host_mirror.py deleted file mode 100644 index abcd8b06b1c9..000000000000 --- a/vllm_ascend/hisparse/host_mirror.py +++ /dev/null @@ -1,332 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Own HiSparse Host snapshots and sampling/draft mirror completion fences.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np -import torch - - -def _host_disjoint(caches, draft_layers): - if any(cache.ndim != 2 or not cache.is_contiguous() or cache.numel() == 0 for cache in caches): - return False - spans = [(cache.data_ptr(), cache.data_ptr() + cache.numel() * cache.element_size()) for cache in caches] - return all( - left_end <= right_start or right_end <= left_start - for index, (left_start, left_end) in enumerate(spans) - if index in draft_layers - for other, (right_start, right_end) in enumerate(spans) - if other not in draft_layers - ) - - -def _draft_groups_disjoint(handles, draft_layers): - draft_runtimes = {id(handles[index].runtime) for index in draft_layers} - # A draft-led resolver can refill every follower from Host in one call. - for index in draft_layers: - group = handles[index].runtime.index_group - if id(group.leader) not in draft_runtimes or any( - id(follower) not in draft_runtimes for follower in group.followers - ): - return False - return True - - -def _index_plan(addresses, sizes, base, width, source_rows, limit): - if not len(addresses): - return None - first, alignment = np.divmod(addresses - base, width) - counts, size_remainder = np.divmod(sizes, width) - if ( - np.any(alignment) - or np.any(size_remainder) - or np.any(counts <= 0) - or np.any(first < 0) - or np.any(first + counts > source_rows) - ): - return None - ends = np.cumsum(counts) - total = int(ends[-1]) - if total > limit: - return None - starts = ends - counts - rows = np.repeat(first - starts, counts) + np.arange(total, dtype=np.int64) - return starts, rows - - -@dataclass -class _Scratch: - cpu_indices: torch.Tensor - device_indices: torch.Tensor - payload: torch.Tensor - ready: torch.Event - completion: torch.Event | None = None - in_flight: bool = False - - def available(self): - return not self.in_flight or (self.completion is not None and self.completion.query()) - - -class TargetMirrorSnapshot: - """Stage exact Target source bytes before draft can overwrite shared HMA.""" - - def __init__(self, worker, layers, width, capacity): - backing = worker.hot_backing - self._source = backing.view(torch.uint8).reshape(-1, width) - self._base = backing.data_ptr() - self._width = width - self._layers = layers - self._sources = worker.resident_caches - self._destinations = worker.host_caches - self._capacity = capacity - self._slots = tuple( - _Scratch( - torch.empty(capacity, dtype=torch.int64, pin_memory=True), - torch.empty(capacity, dtype=torch.int64, device=backing.device), - torch.empty((capacity, width), dtype=torch.uint8, device=backing.device), - torch.Event(), - ) - for _ in range(2) - ) - self.armed = False - self.submitted_completion = None - - @classmethod - def bind(cls, worker, speculator): - geometry = getattr(worker, "_ascend_row_dma_geometry", None) - draft = set(worker._draft_layers) - names = worker.cache_layer_names - expected = set(getattr(speculator, "draft_attn_layer_names", ())) & set(names) - if geometry is None or not draft or {names[i] for i in draft} != expected: - return None - layers = tuple(index for index in range(len(names)) if index not in draft) - if ( - not layers - or not _draft_groups_disjoint(worker.cache_handles, draft) - or not _host_disjoint(worker.host_caches, draft) - ): - return None - widths = {int(geometry.row_bytes[i]) for i in layers} - if len(widths) != 1: - return None - width = widths.pop() - backing = worker.hot_backing - base = backing.data_ptr() - if ( - width <= 0 - or backing.device.type != "npu" - or backing.dtype not in (torch.int8, torch.uint8) - or backing.ndim != 1 - or not backing.is_contiguous() - or backing.storage_offset() != 0 - or backing.numel() % width - or any( - cache.untyped_storage().data_ptr() != base - or (cache.data_ptr() - base) % width - or geometry.source_page_bytes[index] % width - for index in layers - for cache in (worker.resident_caches[index],) - ) - ): - return None - config = worker.vllm_config - capacity = config.scheduler_config.max_num_seqs * (config.num_lookahead_tokens + 1) * len(layers) - return cls(worker, layers, width, capacity) - - def try_submit(self, worker, descriptors, count, layer_indices, ready_event): - if ( - not self.armed - or self.submitted_completion is not None - or ready_event is not worker._forward_ready_event - or tuple(layer_indices) != self._layers - or worker.resident_caches is not self._sources - or worker.host_caches is not self._destinations - ): - return False - slot = next((item for item in self._slots if item.available()), None) - if slot is None: - return False - plan = _index_plan( - descriptors.src_np[:count], - descriptors.sizes_np[:count], - self._base, - self._width, - self._source.shape[0], - self._capacity, - ) - if plan is None: - return False - starts, rows = plan - size = len(rows) - # Owners remain pinned if any submission fails. Never retry a partial - # enqueue using the original device addresses. - slot.in_flight = True - slot.completion = None - slot.cpu_indices[:size].numpy()[:] = rows - slot.device_indices[:size].copy_(slot.cpu_indices[:size], non_blocking=True) - torch.index_select(self._source, 0, slot.device_indices[:size], out=slot.payload[:size]) - slot.ready.record() - descriptors.src_np[:count] = slot.payload.data_ptr() + starts * self._width - worker._submit_dma_descriptors(descriptors, count, ready_event=slot.ready) - slot.completion = worker._pending_dma_descriptors[-1][0] - self.submitted_completion = slot.completion - return True - - -class SamplingMirrorOverlapConnector: - """Let Target logits and sampling run before joining the Host mirror.""" - - def __init__(self, connector, worker, output_copy_stream, snapshot=None): - self._connector = connector - self._worker = worker - self._output_copy_stream = output_copy_stream - self._snapshot = snapshot - self._overlap_draft = False - self._pending_event = None - self._allow_step = False - self._eligible_forward = False - - def __getattr__(self, name): - return getattr(self._connector, name) - - def join_pending(self) -> None: - if self._pending_event is not None: - torch.accelerator.current_stream().wait_event(self._pending_event) - self._pending_event = None - self._overlap_draft = False - - def begin_execution(self, allow_step: bool) -> None: - # The runner calls this before preemption or any scheduler state writes. - self.join_pending() - self._allow_step = allow_step - self._eligible_forward = False - - def pre_forward(self, scheduler_output, input_batch=None, attn_metadata=None): - self._eligible_forward = ( - self._allow_step - and scheduler_output.has_sync_kv_loads - and input_batch is not None - and attn_metadata is not None - and not input_batch.has_prefill - and 0 < input_batch.num_tokens == input_batch.num_tokens_after_padding - ) - return self._connector.pre_forward(scheduler_output, input_batch, attn_metadata=attn_metadata) - - def finish_forward(self) -> None: - eligible = self._eligible_forward - self._eligible_forward = False - if not eligible or self._connector._disabled: - return self._connector.finish_forward() - - worker = self._worker - compute_stream = torch.accelerator.current_stream() - worker._forward_ready_event.record(compute_stream) - snapshot = self._snapshot - pending_before = 0 - if snapshot is not None: - pending_before = len(worker._pending_dma_descriptors) - snapshot.submitted_completion = None - snapshot.armed = ( - not worker._dma_submitted and not worker._per_layer_mirrored and not worker._submitted_mirror_layers - ) - try: - worker._finish_mirror_phase(worker._forward_ready_event) - except BaseException: - # The runner's failure boundary joins an already-recorded copy. - if worker._dma_submitted: - self._pending_event = worker.host_write_event - raise - finally: - if snapshot is not None: - snapshot.armed = False - if worker._dma_submitted: - self._pending_event = worker.host_write_event - # Preserve the Host durability fence on AsyncOutput's copy stream. - self._output_copy_stream.wait_event(self._pending_event) - worker._dma_submitted = False - if snapshot is not None and snapshot.submitted_completion is not None: - pending = worker._pending_dma_descriptors - # Direct fallback or an extra transfer keeps the sampling-only fence. - self._overlap_draft = ( - len(pending) == pending_before + 1 and pending[-1][0] is snapshot.submitted_completion - ) - else: - worker.host_write_event.record(compute_stream) - worker._release_completed_dma_descriptors() - - def post_forward(self, finished_req_ids, defer_finalize=False): - if not defer_finalize or not self._overlap_draft: - self.join_pending() - return self._connector.post_forward(finished_req_ids, defer_finalize=defer_finalize) - - def finalize_step(self): - self.join_pending() - return self._connector.finalize_step() - - def no_forward(self, scheduler_output): - self.begin_execution(False) - return self._connector.no_forward(scheduler_output) - - def set_disabled(self, disabled: bool) -> None: - self.begin_execution(False) - self._connector.set_disabled(disabled) - - def reset_capture_state(self) -> None: - self.begin_execution(False) - self._connector.reset_capture_state() - - -def overlap_mtp_host_mirror(runner): - """Bind only the audited single-stage GLM MTP sampling path.""" - from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector - from vllm.distributed.parallel_state import get_tensor_model_parallel_rank - from vllm.logger import init_logger - from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector - - from vllm_ascend.hisparse.backend import AscendHiSparseBackend - - connector = runner.kv_connector - config = runner.vllm_config.parallel_config - if ( - type(connector) is not ActiveKVConnector - or type(connector.kv_connector) is not HiSparseConnector - or runner.speculator is None - or runner.speculative_config.method != "mtp" - or runner.model_config.architecture != "GlmMoeDsaForCausalLM" - or config.pipeline_parallel_size != 1 - or config.data_parallel_size != 1 - or config.decode_context_parallel_size != 1 - or runner.pcp_manager is not None - or runner.is_pooling_model - or runner.supports_mm_inputs - or runner.aux_output_connector is not None - or runner.kvpp.scheduler is not None - or getattr(runner, "pd_dspark_aux_layer_ids", ()) - ): - return connector - worker = connector.kv_connector.connector_worker - if ( - worker is None - or not worker._initialized - or not isinstance(worker.backend, AscendHiSparseBackend) - or worker.shared_host_region is not None - or not worker.is_host_writer - ): - return connector - output = runner.output_copy_stream - output_stream = torch.Stream( - stream_id=output.stream_id, device_index=output.device_index, device_type=output.device_type - ) - if output_stream == torch.accelerator.current_stream(): - return connector - logger = init_logger(__name__) - rank = get_tensor_model_parallel_rank() - snapshot = TargetMirrorSnapshot.bind(worker, runner.speculator) - worker._ascend_target_snapshot = snapshot - logger.info("HiSparse MTP mirror overlap enabled: rank=%d", rank) - if snapshot is not None: - logger.info("HiSparse Target snapshot enabled: rank=%d", rank) - return SamplingMirrorOverlapConnector(connector, worker, output_stream, snapshot) diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index e2c76dea8137..265de6c4bbd8 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -107,9 +107,7 @@ def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: np.multiply(destinations, row_bytes, out=dst) dst += geometry.destination_addresses[layers] descriptors.sizes_np[:count].reshape(shape)[:] = row_counts * row_bytes - snapshot = getattr(worker, "_ascend_target_snapshot", None) - if snapshot is None or not snapshot.try_submit(worker, descriptors, count, layer_indices, ready_event): - worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) + worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) return True diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index d558b64a5ebf..79b738c6a2c2 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -338,11 +338,7 @@ def sample_tokens(self, grammar_output): "prompt logprobs and fixed-token prompt scoring are not supported with lmhead TP." ) - try: - output = super().sample_tokens(grammar_output) - finally: - if (mirror := getattr(self, "_sampling_mirror", None)) is not None: - mirror.join_pending() + output = super().sample_tokens(grammar_output) # Close the per-step msprobe dump cycle on the normal last-PP-rank # path (v1 parity: v1 finalizes at the end of sample_tokens). self._finalize_dump_data() @@ -527,13 +523,6 @@ def initialize_kv_cache( if self.vllm_config.attention_config.hisparse_config is not None: self._bind_hisparse_direct_prolog() - # Bind after Core seals all cache and stream owners. - from vllm_ascend.hisparse.host_mirror import overlap_mtp_host_mirror - - overlap_connector = overlap_mtp_host_mirror(self) - self._sampling_mirror = overlap_connector if overlap_connector is not self.kv_connector else None - self.kv_connector = overlap_connector - def _bind_hisparse_direct_prolog(self) -> None: from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.hisparse.layout import HISPARSE_RESIDENT_SUFFIX @@ -658,9 +647,6 @@ def execute_model( valid_dummy_state_slots: bool = False, randomize_inputs: bool = False, ): - mirror = getattr(self, "_sampling_mirror", None) - if mirror is not None: - mirror.begin_execution(not dummy_run and not is_profile) self._cpp_execution_time_ms = None profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config execution_start_time = _start_profiling_chunk_timing( @@ -702,8 +688,6 @@ def execute_model( ) forward_failed = False finally: - if forward_failed and mirror is not None: - mirror.join_pending() finish_execution = getattr(self.model_state, "finish_execution", None) if finish_execution is not None: finish_execution(failed=forward_failed) From fd388b21539fb6cd358c7f5c21b17f44970864ef Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 23:08:10 +0800 Subject: [PATCH 083/120] refactor(hisparse): write C8 Prolog output to the bound cache Use Core write_target for native Prolog destinations by default, pad a shortened valid slot prefix with -1, and remove runner opt-in binding plus staging/scatter fallback for this path. Extend the existing eager/Graph comparison with real Core bindings and changing padded slots. Device and whole-model qualification remain pending. Signed-off-by: zhenwenqi2024 --- .../user_guide/feature_guide/hisparse.md | 8 +- .../one_card/test_hisparse_kv_preprocess.py | 48 +++++--- vllm_ascend/attention/sfa_hisparse.py | 105 +++++------------- vllm_ascend/worker/v2/model_runner.py | 37 ------ 4 files changed, 64 insertions(+), 134 deletions(-) diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index c3ae0969f64a..0e1433d6bcc7 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -78,9 +78,11 @@ request mappings, padding and completion-before-reuse rules. ## Validation and performance -Ascend 950 eager/Graph components and TP8/MTP3 full-model checks cover the -supported paths. A matched C8 test with two rounds per arm, fixed inputs and -8K outputs measured these mean SSE TPOT increases: +Ascend 950 eager/Graph components and TP8/MTP3 full-model checks were recorded +at the source revisions below. The default-direct C8 Prolog write path still +requires fresh device and model qualification. A historical matched C8 test +with two rounds per arm, fixed inputs and 8K outputs measured these mean SSE +TPOT increases: | Input / concurrency | Reported residency | ON versus OFF | | --- | --- | --- | diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index 729e5347431a..1e6213b1a9a3 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -117,8 +117,12 @@ def forward(): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("rows", [4, 16]) -def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): +@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) +def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture, slot_dtype): _setup() + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.v1.hisparse.runtime import HiSparseCacheHandle + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl from vllm_ascend.attention.sfa_v1 import AscendSFAImpl from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod @@ -133,10 +137,6 @@ def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture): impl.local_num_heads, impl.has_indexer = 8, True impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False impl.sfa_qsfa_tile_size = block - impl._kv_staging_capacity = block - # This case exercises the staged destination, not the separately bound - # direct-Resident capture path installed by the model runner. - impl._direct_prolog_slots = None def mx_layer(inputs, outputs): weight, scale = torch_npu.npu_dynamic_mx_quant( @@ -159,24 +159,35 @@ def mx_layer(inputs, outputs): # the oracle. This case tests ownership/addressing, not Prolog arithmetic. impl._process_weights_for_fused_prolog_v3() impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16, device=device) - staging = torch.empty((1, block, 1, width), dtype=impl.c8_cache_dtype, device=device) + reference = torch.empty((1, block, 1, width), dtype=impl.c8_cache_dtype, device=device) local_slots = torch.arange(block, dtype=torch.int64, device=device) - impl._c8_prolog_workspace = lambda: (staging, local_slots) raw = torch.full((3, 2, block, width), 90, dtype=torch.uint8, device=device) - destination = raw[:, 1].view(impl.c8_cache_dtype) logical = [-1, 129, 0, 257, 128][:rows] + list(range(16, 16 + max(0, rows - 5))) - slots = torch.tensor(logical, dtype=torch.int64, device=device) - impl.hisparse_cache = SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)) + slots = torch.tensor(logical, dtype=slot_dtype, device=device) + cache = HiSparseCacheHandle(SimpleNamespace(max_num_reqs=rows, kv_dtype=impl.c8_cache_dtype, row_width=width)) + cache.bind_cache( + raw.view(-1), + byte_offset=block * width, + block_stride=2 * block * width, + num_blocks=3, + block_size=block, + block_table=torch.empty(0, dtype=torch.int32, device=device), + slot_mapping=slots, + ) + cache.dummy_batch = True # This component test has no Host-mirror connector. + cache.num_actual_tokens = rows + impl.hisparse_cache = cache metadata = SimpleNamespace( - query_start_loc=torch.tensor([0, rows], dtype=torch.int32, device=device), hisparse_for_capture=True + query_start_loc=torch.tensor([0, rows], dtype=torch.int32, device=device), hisparse_for_capture=capture ) x = torch.randn(rows, 6144, dtype=torch.bfloat16, device=device) cos = torch.ones(rows, 1, 1, 64, dtype=x.dtype, device=device) sin = torch.zeros_like(cos) - reference = torch.empty_like(staging) def forward(): - return impl._sfa_preprocess_prolog_v3(x, (), cos, sin, slots, attn_metadata=metadata) + context = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + with override_forward_context(context): + return impl._sfa_preprocess_prolog_v3(x, (), cos, sin, slots, attn_metadata=metadata) def compare(actual, expected): if isinstance(actual, tuple): @@ -198,9 +209,18 @@ def compare(actual, expected): graph = torch.npu.NPUGraph() with torch.npu.graph(graph): captured = forward() - for valid in (rows - 1, 2): + for valid in (rows - 1, 2, 0, rows): x.neg_() metadata.query_start_loc[-1:].fill_(valid) + # The real slot producer marks dummy and padded rows with -1. + # Replay consumes these changed device values at fixed addresses. + slots.copy_( + torch.tensor( + [slot if row < valid else -1 for row, slot in enumerate(logical)], dtype=slot_dtype, device=device + ) + ) + if not capture: + cache.num_actual_tokens = valid raw.fill_(90) if graph is None: actual = forward() diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 48f50978e048..f875a05dde47 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -15,7 +15,6 @@ from vllm.v1.worker.workspace import current_workspace_manager from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder -from vllm_ascend.attention.utils import PreprocessType from vllm_ascend.hisparse.prefill import ( AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan, @@ -146,7 +145,6 @@ def __init__(self, *args, **kwargs): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") self.hisparse_cache = self.index_group.cache(self.index_group_index) self.hisparse_cache._ascend_prefill_layer_name = self.layer_name - self._direct_prolog_slots: torch.Tensor | None = None self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: raise ValueError("Ascend HiSparse SFA requires top-k to be divisible by 128") @@ -157,9 +155,8 @@ def __init__(self, *args, **kwargs): self._kv_staging_capacity = (max_rows + 127) // 128 * 128 def _fused_preprocess_type(self): - # Reuse SFA's supported quantization/weight preparation. Its C8 Prolog - # writes owned staging below; the HiSparse handle retains the final - # Resident destination, row mapping and Host-mirror completion policy. + # Reuse SFA's supported quantization/weight preparation. The HiSparse + # handle supplies the final Resident cache and current write mapping. if self.enable_sparse_sfa_c8: from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod @@ -219,104 +216,52 @@ def allocate(): return current_workspace_manager().get_persistent_resource(("ascend_hisparse_c8_lengths",), allocate) - def _c8_prolog_workspace(self): - from vllm_ascend.hisparse.c8 import C8_ROW_BYTES - - def allocate(): - device = self.index_group.logical_topk_indices.device - return ( - torch.empty( - (self._kv_staging_capacity // 128, 128, 1, C8_ROW_BYTES), - dtype=self.c8_cache_dtype, - device=device, - ), - torch.arange(self._kv_staging_capacity, dtype=torch.int64, device=device), - ) - - # Like the existing preprocessing workspace, these rows are consumed - # on the compute stream before the next layer/graph can overwrite them. - # Host mirroring reads the final Resident cache, never this scratch. - return current_workspace_manager().get_persistent_resource( - ("ascend_hisparse_c8_prolog", self._kv_staging_capacity, self.c8_cache_dtype), allocate - ) - - def _resident_prolog_slots_i64(self, mapping): + def _prolog_slot_mapping(self, mapping, num_input_rows): context = get_forward_context() scope = getattr(context, "_ascend_hisparse_resident_i64", None) if scope is None: scope = context._ascend_hisparse_resident_i64 = {} # Exact slot aliases across layers share one cast. The captured cast # reads current slots again on replay, even when the address is stable. - key = (mapping.data_ptr(), mapping.numel(), mapping.device, torch.npu.current_stream().npu_stream) + key = ( + mapping.data_ptr(), + mapping.numel(), + num_input_rows, + mapping.device, + torch.npu.current_stream().npu_stream, + ) cached = scope.get(key) if cached is None: - cached = (mapping.untyped_storage(), mapping.to(torch.int64)) + slots = mapping.to(torch.int64) + # Core returns the effective write prefix. Prolog computes queries + # for every input row and skips cache writes at -1 slots. Graph + # replay's changing padding is already set by the slot producer. + if mapping.numel() < num_input_rows: + slots = torch.nn.functional.pad(slots, (0, num_input_rows - mapping.numel()), value=-1) + cached = (mapping.untyped_storage(), slots) scope[key] = cached return cached[1] def _sfa_preprocess_prolog_v3(self, hidden_states, kv_cache, cos, sin, slot_mapping, *, attn_metadata=None): - from vllm_ascend.hisparse.c8 import C8_ROW_BYTES - from vllm_ascend.hisparse.cache_ops import scatter_packed_kv - if attn_metadata is None: raise ValueError("HiSparse Prolog requires current row-validity metadata") rows = hidden_states.shape[0] - if rows > self._kv_staging_capacity: - raise ValueError("HiSparse Prolog exceeds its profiled token capacity") cache = self.hisparse_cache - destination, mapping, count = cache.write_target(rows, slot_mapping.numel()) - resident_slots = self._direct_prolog_slots - direct = ( - resident_slots is not None - and not cache.draft_layer - and attn_metadata.hisparse_for_capture - and attn_metadata.num_prefills == 0 - and 0 < count == rows - and mapping.shape == (rows,) - and mapping.stride(0) == 1 - and mapping.dtype == torch.int32 - and mapping.device == resident_slots.device == destination.device == hidden_states.device - and mapping.data_ptr() == resident_slots.data_ptr() - and destination.ndim == 3 - and destination.shape[1:] == (128, C8_ROW_BYTES) - and destination.dtype == torch.float8_e4m3fn - and destination.storage_offset() % C8_ROW_BYTES == 0 - and destination.stride()[1:] == (C8_ROW_BYTES, 1) - and destination.stride(0) >= 128 * C8_ROW_BYTES - and destination.stride(0) % (128 * C8_ROW_BYTES) == 0 + destination, mapping, _ = cache.write_target(rows, slot_mapping.numel()) + result = super()._sfa_preprocess_prolog_v3( + hidden_states, + (destination.unsqueeze(2),), + cos, + sin, + self._prolog_slot_mapping(mapping, rows), + attn_metadata=attn_metadata, ) - if direct: - logger.info_once( - "HiSparse Target direct Prolog capture: rows=%d stride_bytes=%d", - rows, - destination.stride(0), - scope="process", - ) - result = super()._sfa_preprocess_prolog_v3( - hidden_states, - (destination.unsqueeze(2),), - cos, - sin, - self._resident_prolog_slots_i64(mapping), - attn_metadata=attn_metadata, - ) - else: - staging, local_slots = self._c8_prolog_workspace() - result = super()._sfa_preprocess_prolog_v3( - hidden_states, (staging,), cos, sin, local_slots[:rows], attn_metadata=attn_metadata - ) - packed = staging.view(-1, C8_ROW_BYTES)[:count].view(torch.uint8) - scatter_packed_kv( - packed, packed[:, :0], destination.view(torch.uint8), mapping, attn_metadata.query_start_loc[-1:] - ) finish_hisparse_kv_update(cache, attn_metadata) return result def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=None): if self.enable_sparse_sfa_c8: self._c8_length_workspace() - if self.preprocess_type == PreprocessType.PROLOG_V3: - self._c8_prolog_workspace() else: self._workspace(hidden_states.dtype) if self.qk_rope_head_dim and not self.enable_sparse_sfa_c8: diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 79b738c6a2c2..a8a95a8641aa 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -520,43 +520,6 @@ def initialize_kv_cache( self._fill_c8_mxfp_v_scale_caches() - if self.vllm_config.attention_config.hisparse_config is not None: - self._bind_hisparse_direct_prolog() - - def _bind_hisparse_direct_prolog(self) -> None: - from vllm.v1.attention.backends.utils import PAD_SLOT_ID - from vllm.v1.hisparse.layout import HISPARSE_RESIDENT_SUFFIX - from vllm.v1.kv_cache_interface import HiSparseResidentSpec - - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type - from vllm_ascend.worker.v2.block_table import AscendBlockTables - - block_tables = self.block_tables - # These graphs only replay decode, and the exact slot producer pads - # dummy requests and every unused row with native Prolog's -1 sentinel. - enable = ( - get_ascend_device_type() == AscendDeviceType.A5 - and self.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY - and type(block_tables) is AscendBlockTables - and self.pcp_manager is None - and block_tables.cp_size == 1 - and PAD_SLOT_ID == -1 - ) - layers = self.compilation_config.static_forward_context - for group_id, group in enumerate(self.kv_cache_config.kv_cache_groups): - if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): - continue - for cache_name in group.layer_names: - layer = layers[cache_name.removesuffix(HISPARSE_RESIDENT_SUFFIX)] - impl = layer.impl - if type(impl) is AscendHiSparseImpl: - impl._direct_prolog_slots = ( - block_tables.slot_mappings[group_id] - if enable and block_tables.blocks_per_kv_block[group_id] == 1 - else None - ) - def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" self._c8_mxfp_v_scale_caches = [] From a11b1b9d1d1f2415e7e5953137bcae2202c240cb Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 23:28:29 +0800 Subject: [PATCH 084/120] refactor(hisparse): prepare workspaces at worker startup Move layer discovery into the HiSparse backend and prepare retained workspace before both worker memory-profile branches. Reuse Core total-consumption accounting and existing binding owners without additional budget formulas or allocations. Remove the remaining HiSparse-specific model-runner changes and adapt existing workspace tests. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 9 +++++---- .../kv_transfer/test_hisparse_foundation_cpu.py | 9 ++++++++- vllm_ascend/hisparse/backend.py | 15 +++++++++------ vllm_ascend/worker/v2/model_runner.py | 8 -------- vllm_ascend/worker/worker.py | 5 +++++ 5 files changed, 27 insertions(+), 19 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index a401dc3f98d6..569732fb6101 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -251,24 +251,25 @@ def capture_lengths(ops): try: first = runtime() - prepare_hisparse_workspaces([first]) + layers = {"target": SimpleNamespace(hisparse_cache=SimpleNamespace(runtime=first))} + prepare_hisparse_workspaces(layers) pending = first._ascend_profiled_cache_ops workspace = pending[0] - prepare_hisparse_workspaces([first]) + prepare_hisparse_workspaces(layers) assert first._ascend_profiled_cache_ops is pending for tensor in (workspace.first_refs, workspace.errors, *pending[1]): assert not torch.is_inference(tensor) old = bind(first, pages) assert old.workspace is workspace and not hasattr(first, "_ascend_profiled_cache_ops") captured = [capture_lengths(old)] - prepare_hisparse_workspaces([first]) + prepare_hisparse_workspaces(layers) new = bind(first, pages + 1) assert new.workspace is workspace for tensor in new.attention_lengths: assert not torch.is_inference(tensor) assert new.attention_lengths[1].data_ptr() != old.attention_lengths[1].data_ptr() captured.append(capture_lengths(new)) - prepare_hisparse_workspaces([first]) + prepare_hisparse_workspaces(layers) third = bind(first, pages + 2) assert third.workspace is workspace for tensor in third.attention_lengths: diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index ec6ec0d14df4..16014f58e1e2 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -45,7 +45,14 @@ def prepare_workspaces(request): try: if request.param: init_workspace_manager(torch.device("cpu")) - yield torch.inference_mode()(prepare_hisparse_workspaces) if request.param else None + + @torch.inference_mode() + def prepare(runtimes): + layers = {str(i): NS(hisparse_cache=NS(runtime=runtime)) for i, runtime in enumerate(runtimes)} + layers.update(ordinary_layer=NS(), unbound_layer=NS(hisparse_cache=None)) + prepare_hisparse_workspaces(layers) + + yield prepare if request.param else None finally: reset_workspace_manager() diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index b2b224f46de0..15463aa5fb67 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -5,7 +5,6 @@ from __future__ import annotations import importlib -from collections.abc import Iterable from typing import TYPE_CHECKING, Any import torch @@ -39,23 +38,27 @@ def _allocate_attention_lengths(runtime: HiSparseRuntime) -> tuple[torch.Tensor, @torch.inference_mode(False) -def prepare_hisparse_workspaces(runtimes: Iterable[HiSparseRuntime]) -> None: - """Create real next-binding storage inside the existing profile scope. +def prepare_hisparse_workspaces(forward_context: dict[str, Any]) -> None: + """Create real next-binding storage before Worker memory profiling. Core fixes group geometry when loading the model. Retain one scratch owner per actual group, and fresh C8 lengths for each consumed binding generation. + Core's total-consumption measurement includes these retained allocations. """ from vllm.v1.worker.workspace import current_workspace_manager - manager = current_workspace_manager() - for runtime in runtimes: + for layer in forward_context.values(): + cache = getattr(layer, "hisparse_cache", None) + if cache is None: + continue + runtime = cache.runtime group = runtime.index_group workspace = getattr(group, "_ascend_profiled_resolver", None) if workspace is None: # Identity, not shape, separates groups. The manager key does not # retain the runtime or any Host-pool/Graph owner. key = ("ascend_hisparse_resolver", object()) - workspace = manager.get_persistent_resource( + workspace = current_workspace_manager().get_persistent_resource( key, lambda runtime=runtime: _allocate_resolver_workspace(runtime) ) group._ascend_profiled_resolver = workspace diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index a8a95a8641aa..e1cb6b8072af 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -737,14 +737,6 @@ def profile_run(self, randomize_inputs: bool = False) -> None: necessary HCCL buffer for the MC2 operator before standard `profile_run`. Additionally, we set override_mrv2_in_profile_run to True to force moe load to be balanced when executing `profile_run` """ - if self.vllm_config.attention_config.hisparse_config is not None: - from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces - - prepare_hisparse_workspaces( - cache.runtime - for layer in self.compilation_config.static_forward_context.values() - if (cache := getattr(layer, "hisparse_cache", None)) is not None - ) sparse_cfg = self.ascend_config.sparse_kv_offload_config if sparse_cfg.enabled: allocate_kv_offload_topk_profile_buffers(self.get_kv_cache_spec(), self.vllm_config, sparse_cfg) diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 82064f324b19..a01ab0de112f 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -676,6 +676,11 @@ def determine_available_memory(self) -> int: maybe_apply_startup_plan(self) + if self.vllm_config.attention_config.hisparse_config is not None: + from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces + + prepare_hisparse_workspaces(self.vllm_config.compilation_config.static_forward_context) + randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and self.parallel_config.data_parallel_size > 1 # Fast path: user has explicitly specified KV cache size via From 62d66e65e772f98f749fe5550e577a23a9a9757d Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 23:59:41 +0800 Subject: [PATCH 085/120] perf(attention): generalize SFA RoPE metadata reuse Share RoPE between compatible SFA and Indexer builders without a HiSparse or execution-phase opt-in. Invalidate at custom writer boundaries and check table, destination and position ownership and tracked versions. Preserve independently owned drafting storage. Source-isolated CPU contracts: 28 passed. Focused Ruff, format, syntax and diff checks passed. Full pytest lacks vllm; format CI timed out installing actionlint. NPU eager/Graph and model qualification remain pending. Signed-off-by: zhenwenqi2024 --- tests/ut/worker/test_attn_utils_v2.py | 132 ++++++++++++++++++ .../attention/context_parallel/sfa_cp.py | 7 +- vllm_ascend/ops/rotary_embedding.py | 23 ++- vllm_ascend/worker/v2/attn_utils.py | 32 ++--- vllm_ascend/worker/v2/model_states/default.py | 9 -- 5 files changed, 164 insertions(+), 39 deletions(-) diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index 3e234cdde5db..4e26c27847b5 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -68,6 +68,138 @@ from vllm_ascend.worker.v2.utils import AscendV2KVBlockZeroer +@pytest.mark.parametrize("full_graph,for_capture", [(False, False), (True, False), (True, True)]) +@pytest.mark.parametrize("boundary", ["none", "builder", "callback", "provider"]) +@pytest.mark.parametrize("inference", [False, True]) +def test_plain_sfa_and_indexer_share_rope_only_within_one_compatible_build( + monkeypatch, full_graph, for_capture, boundary, inference +): + with torch.inference_mode(inference): + from vllm_ascend.attention.indexer import AscendSFAIndexerMetadataBuilder + from vllm_ascend.ops import rotary_embedding as rotary + + table = torch.arange(128, dtype=torch.float32).reshape(32, 4) + monkeypatch.setattr(rotary, "_cos_cache", table) + monkeypatch.setattr(rotary, "_sin_cache", -table) + monkeypatch.setattr(rotary, "_cos_sin_flat", None) + monkeypatch.setattr(rotary, "_cos_mla", torch.empty(3, 1, 1, 4)) + monkeypatch.setattr(rotary, "_sin_mla", torch.empty(3, 1, 1, 4)) + lookup = rotary.get_cos_and_sin_mla + lookups = [] + + def counted_lookup(positions, use_cache=False): + lookups.append(positions.clone()) + return lookup(positions, use_cache=use_cache) + + monkeypatch.setattr(rotary, "get_cos_and_sin_mla", counted_lookup) + + class Recorder: + def __init__(self, offset=0): + self.offset = offset + self.nope = False + + def build(self, common_prefix_len, common_attn_metadata, shared_mla_rope=None, **kwargs): + common = common_attn_metadata + positions = common.positions[: common.num_input_tokens] + if self.offset: + positions = positions + self.offset + if shared_mla_rope is None: + cos, sin = counted_lookup(positions, use_cache=True) + else: + cos, sin = shared_mla_rope.get(common.positions, positions) + torch.testing.assert_close(cos, table[positions].unsqueeze(1).unsqueeze(2)) + torch.testing.assert_close(sin, -table[positions].unsqueeze(1).unsqueeze(2)) + # The unrelated consumer owns its copies, like a separate builder. + return SimpleNamespace(cos=cos.clone() if self.offset else cos, sin=sin.clone() if self.offset else sin) + + def build_for_cudagraph_capture(self, common_attn_metadata, **kwargs): + return self.build(0, common_attn_metadata, **kwargs) + + class PlainSFA(Recorder, AscendSFAMetadataBuilder): + pass + + class PlainIndexer(Recorder, AscendSFAIndexerMetadataBuilder): + pass + + builders = [("sfa", PlainSFA())] + if boundary == "builder": + builders.append(("other", Recorder(5))) + builders.append(("indexer", PlainIndexer())) + groups = [ + [SimpleNamespace(layer_names=[name], get_metadata_builder=lambda _, builder=builder: builder)] + for name, builder in builders + ] + positions = torch.tensor([1, 2, 3], dtype=torch.int64) + config = SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=None) for _ in builders]) + + def update_positions(group, common=None): + if group: + positions.add_(1) + return {} + + provider = SimpleNamespace( + get_extra_common_attn_kwargs=update_positions, + get_extra_attn_kwargs=lambda *args: {}, + ) + for _ in range(2): + lookups.clear() + metadata = attn_utils.build_attn_metadata( + attn_groups=groups, + num_reqs=1, + num_tokens=3, + query_start_loc_gpu=torch.tensor([0, 3], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 3], dtype=torch.int32), + max_query_len=3, + seq_lens=torch.tensor([10], dtype=torch.int32), + max_seq_len=10, + block_tables=tuple(torch.zeros(1, 1, dtype=torch.int32) for _ in builders), + slot_mappings=tuple(torch.zeros(3, dtype=torch.int32) for _ in builders), + kv_cache_config=config, + positions=positions, + full_graph_mode=full_graph, + for_cudagraph_capture=for_capture, + offload_dummy=for_capture, + is_prefilling=np.array([True]), + common_metadata_callback=update_positions if boundary == "callback" else None, + model_specific_attn_metadata=provider if boundary == "provider" else None, + parallel_config=SimpleNamespace(prefill_context_parallel_size=1, decode_context_parallel_size=1), + ) + expected = table[positions].unsqueeze(1).unsqueeze(2) + for name in ("sfa", "indexer"): + torch.testing.assert_close(metadata[name].cos, expected) + torch.testing.assert_close(metadata[name].sin, -expected) + expected_lookups = 1 if boundary == "none" else len(builders) + assert len(lookups) == expected_lookups + # The same position tensor is updated between calls; no prior-build hit. + positions.add_(1) + + +@pytest.mark.parametrize("change", ["positions", "table", "destination", "rebind"]) +def test_mla_rope_cache_refreshes_current_values_after_mutation(monkeypatch, change): + from vllm_ascend.ops import rotary_embedding as rotary + + table = torch.arange(128, dtype=torch.float32).reshape(32, 4) + monkeypatch.setattr(rotary, "_cos_cache", table) + monkeypatch.setattr(rotary, "_sin_cache", -table) + monkeypatch.setattr(rotary, "_cos_sin_flat", None) + monkeypatch.setattr(rotary, "_cos_mla", torch.empty(3, 1, 1, 4)) + monkeypatch.setattr(rotary, "_sin_mla", torch.empty(3, 1, 1, 4)) + positions = torch.tensor([1, 2, 3]) + cache = rotary.MLARopeMetadataCache() + cache.get(positions, positions) + if change == "positions": + positions.add_(1) + elif change == "table": + table.add_(1) + elif change == "destination": + rotary._cos_mla.zero_() + else: + monkeypatch.setattr(rotary, "_cos_cache", table + 1) + cos, sin = cache.get(positions, positions) + torch.testing.assert_close(cos, rotary._cos_cache[positions].unsqueeze(1).unsqueeze(2)) + torch.testing.assert_close(sin, rotary._sin_cache[positions].unsqueeze(1).unsqueeze(2)) + + class SparseAttentionBackend(AscendAttentionBackend): @classmethod def is_sparse(cls) -> bool: diff --git a/vllm_ascend/attention/context_parallel/sfa_cp.py b/vllm_ascend/attention/context_parallel/sfa_cp.py index f41078d65aed..b0813d10f711 100644 --- a/vllm_ascend/attention/context_parallel/sfa_cp.py +++ b/vllm_ascend/attention/context_parallel/sfa_cp.py @@ -781,7 +781,12 @@ def build( global_dcp_num_blocks = pcp_context.global_block_table_num_blocks[pcp_cache_group_idx] metadata = self._build_with_metadata_view( common_attn_metadata, - lambda: self._build(common_attn_metadata, draft_index=None, pcp_context=pcp_context), + lambda: self._build( + common_attn_metadata, + draft_index=None, + pcp_context=pcp_context, + shared_mla_rope=kwargs.get("shared_mla_rope"), + ), global_dcp_block_table=global_dcp_block_table, global_dcp_num_blocks=global_dcp_num_blocks, ) diff --git a/vllm_ascend/ops/rotary_embedding.py b/vllm_ascend/ops/rotary_embedding.py index 39887275be9a..cda198c35c31 100644 --- a/vllm_ascend/ops/rotary_embedding.py +++ b/vllm_ascend/ops/rotary_embedding.py @@ -149,12 +149,22 @@ def get_cos_and_sin_mla(positions, use_cache=False): class MLARopeMetadataCache: - """Share a persistent RoPE slice across compatible builders in one forward.""" + """Share RoPE within one metadata build, between position/storage writers. + + The caller invalidates sharing across custom providers and other builders. + Inference tensors have no version counter, so these write boundaries also + protect their contents. A new cache is used for every metadata build. + """ def __init__(self) -> None: - self._positions: torch.Tensor | None = None self._num_tokens = 0 self._value: tuple[torch.Tensor, torch.Tensor] | None = None + self._owners: tuple = () + self._versions: tuple = () + + @staticmethod + def _versions_of(tensors: tuple) -> tuple: + return tuple(None if tensor is None or tensor.is_inference() else tensor._version for tensor in tensors) @staticmethod def _uses_persistent_buffers(value: tuple[torch.Tensor, torch.Tensor]) -> bool: @@ -172,16 +182,19 @@ def _uses_persistent_buffers(value: tuple[torch.Tensor, torch.Tensor]) -> bool: def get(self, positions: torch.Tensor, input_positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = input_positions.shape[0] + owners = (positions, _cos_cache, _sin_cache, _cos_sin_flat, _cos_mla, _sin_mla) if ( - self._positions is positions - and self._num_tokens == num_tokens + self._num_tokens == num_tokens and self._value is not None + and all(previous is current for previous, current in zip(self._owners, owners)) + and self._versions == self._versions_of(owners) and self._uses_persistent_buffers(self._value) ): return self._value value = get_cos_and_sin_mla(input_positions, use_cache=True) - self._positions = positions self._num_tokens = num_tokens + self._owners = owners + self._versions = self._versions_of(owners) # Preserve the original fallback for absent, empty or mismatched graph # buffers. Only slices of the stable destinations may be reused. self._value = value if self._uses_persistent_buffers(value) else None diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 76da284ddf92..ef638c816189 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -69,6 +69,7 @@ from vllm_ascend.attention.dsa_v1 import AscendDSAMetadataBuilder from vllm_ascend.attention.dsa_v41 import AscendDSAV41MetadataBuilder from vllm_ascend.attention.fa_v1 import AscendFlashAttentionBackend +from vllm_ascend.attention.indexer import AscendSFAIndexerMetadataBuilder from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.attention.utils import ( AscendCommonAttentionMetadata, @@ -502,7 +503,6 @@ def build_attn_metadata( copy_sfa_restore_tails: bool = False, draft_layer_names: set[str] | None = None, common_metadata_callback=None, - reuse_hisparse_mla_rope: bool = False, ) -> dict[str, Any]: """Build attention metadata for Ascend NPUs.""" if skip_ring_state_update is None: @@ -549,28 +549,6 @@ def build_attn_metadata( if num_actual_reqs is None: num_actual_reqs = num_reqs - rope_builders: tuple[type, ...] = () - if ( - reuse_hisparse_mla_rope - and full_graph_mode - and not for_cudagraph_capture - and not offload_dummy - and positions is not None - and pcp_context is None - and dcp_local_seq_lens is None - and parallel_config is not None - and parallel_config.prefill_context_parallel_size == 1 - and parallel_config.decode_context_parallel_size == 1 - and model_specific_attn_metadata is None - and common_metadata_callback is None - and not draft_layer_names - and copy_sfa_draft_index is None - and not copy_sfa_restore_tails - ): - from vllm_ascend.attention.indexer import AscendHiSparseIndexerMetadataBuilder - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseMetadataBuilder - - rope_builders = (AscendHiSparseMetadataBuilder, AscendHiSparseIndexerMetadataBuilder) shared_mla_rope = None # positions will not be used directly in graph modes, so it is saft to create it here. @@ -623,6 +601,9 @@ def build_attn_metadata( group_causal = causal if isinstance(causal, bool) else causal.get(i, True) common_v41_metadata: dict[str, Any] = {} + if model_specific_attn_metadata is not None: + # Custom providers may update positions or persistent RoPE storage. + shared_mla_rope = None common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) if model_specific_attn_metadata is not None @@ -666,6 +647,7 @@ def build_attn_metadata( ) if common_metadata_callback is not None: common_metadata_callback(i, common_attn_metadata) + shared_mla_rope = None for attn_group in attn_groups[i]: attn_metadata_builder = attn_group.get_metadata_builder(0) @@ -691,6 +673,8 @@ def build_attn_metadata( metadata_variants = ([(target_names, common_attn_metadata)] if target_names else []) + [ (draft_names, draft_common) ] + if not for_cudagraph_capture and model_specific_attn_metadata is not None: + shared_mla_rope = None attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_attn_kwargs( attn_metadata_builder, @@ -699,7 +683,7 @@ def build_attn_metadata( if not for_cudagraph_capture and model_specific_attn_metadata is not None else {} ) - if type(attn_metadata_builder) in rope_builders and not getattr(attn_metadata_builder, "nope", False): + if isinstance(attn_metadata_builder, (AscendSFAMetadataBuilder, AscendSFAIndexerMetadataBuilder)): if shared_mla_rope is None: shared_mla_rope = MLARopeMetadataCache() attn_metadata_extra_kwargs["shared_mla_rope"] = shared_mla_rope diff --git a/vllm_ascend/worker/v2/model_states/default.py b/vllm_ascend/worker/v2/model_states/default.py index 811de8c032a8..a2964fd780a3 100644 --- a/vllm_ascend/worker/v2/model_states/default.py +++ b/vllm_ascend/worker/v2/model_states/default.py @@ -563,15 +563,6 @@ def prepare_attn( common_metadata_callback=( self._decoder_replay_common_metadata.__setitem__ if self.decoder_replay is not None else None ), - reuse_hisparse_mla_rope=( - self.vllm_config.attention_config.hisparse_config is not None - and cudagraph_mode == CUDAGraphMode.FULL - and not for_capture - and not input_batch.is_dummy - # MTP verification can be classified as ChunkedPrefill even - # with no requests actually in prefill. - and not input_batch.is_prefilling_np[:num_actual_reqs].any() - ), ) self._prepare_decoder_replay(input_batch, cudagraph_mode, attn_groups, kv_cache_config) if restore_copy_sfa_tails: From 52e6d78794bdeda656cef775bfb8c00c61e4cfd5 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 00:33:37 +0800 Subject: [PATCH 086/120] perf(hisparse): reuse validity and clamped index intermediates Reuse the BF16 adapter validity mask and existing valid lengths instead of repeated comparisons and an any reduction. Reuse the staged Prefill nonnegative indices for page and row offsets. Preserve native Attention arguments and arithmetic. Ten source-isolated CPU comparisons pass with unchanged native input values, physical indices and results. Focused Ruff, format and diff checks pass. Full format CI timed out installing actionlint. NPU Graph/model performance qualification remains pending. Signed-off-by: zhenwenqi2024 --- vllm_ascend/attention/sfa_hisparse.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index f875a05dde47..0a899cb65245 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -442,10 +442,11 @@ def _attention_physical( physical = physical_topk[start:end] keys = positions + (physical < 0).to(torch.float32) * physical.shape[-1] physical = physical.gather(1, keys.sort(dim=-1).indices) - valid_lengths = (physical >= 0).sum(dim=-1, dtype=torch.int32) + valid = physical >= 0 + valid_lengths = valid.sum(dim=-1, dtype=torch.int32) selected_latent, selected_rope = latent[:count], rope[:count] pack_attention_rows(packed_cache, physical, selected_latent, selected_rope, self.kv_lora_rank) - selected_indices = torch.where(physical[:, None, :] >= 0, indices, -1) + selected_indices = torch.where(valid[:, None, :], indices, -1) result = torch.ops._C_ascend.npu_sparse_flash_attention( query=query[start:end], key=selected_latent.view(-1, 128, 1, self.kv_lora_rank), @@ -465,7 +466,7 @@ def _attention_physical( ) if isinstance(result, tuple): result = result[0] - result = result.masked_fill(~(physical >= 0).any(dim=-1)[:, None, None], 0) + result = result.masked_fill((valid_lengths == 0)[:, None, None], 0) outputs.append(result) return torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] @@ -595,10 +596,11 @@ def _execute_sparse_flash_attention_process( ) try: selected = logical[num_decode:] - block_ids = selected.clamp_min(0).to(torch.int64) // metadata.block_size + nonnegative = selected.clamp_min(0) + block_ids = nonnegative.to(torch.int64) // metadata.block_size request_tables = staged_table[req_ids.to(torch.int64)] physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size - physical = physical + selected.clamp_min(0) % metadata.block_size + physical = physical + nonnegative % metadata.block_size physical = physical.masked_fill(selected < 0, -1).to(torch.int32) prefill_output = self._attention_physical(query, query_rope, staged, physical) finally: From f414e7a41a64ac79d8f3510c8cfa720d2f086a4f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 00:50:54 +0800 Subject: [PATCH 087/120] perf(hisparse): share the resident row-copy kernel Use the existing byte-copy kernel for both initial Resident staging and the prefetched Host-row patch. Preserve Core's negative-ID-to-row0 convention for initial staging and retain invalid destinations for the patch. Support flat and paged staging, int32/int64 row IDs and byte-compatible payload views without an intermediate gather allocation. The actual Triton interpreter passed 96 byte-oracle cases with NPU initialization forbidden. Focused Ruff, formatting, syntax and diff checks pass. Existing NPU eager/Graph tests are extended but not executed. No NPU or whole-model performance benefit is claimed; retain this commit as a private candidate pending qualification. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 45 ++++++++++++------- vllm_ascend/hisparse/backend.py | 11 ++--- vllm_ascend/hisparse/cache_ops.py | 30 +++++++++---- vllm_ascend/hisparse/prefill.py | 4 +- vllm_ascend/ops/triton/hisparse.py | 5 ++- 5 files changed, 59 insertions(+), 36 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 569732fb6101..ac0bcc809743 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -12,7 +12,7 @@ from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, - gather_valid_resident_rows, + gather_resident_rows, scatter_packed_kv, ) from vllm_ascend.hisparse.host_pool import allocate_mapped_host @@ -95,29 +95,42 @@ def test_resident_prefill_plan_shared_layers(resident_block_size, host_pages): @pytest.mark.parametrize("dtype,width", [(torch.bfloat16, 576), (torch.float8_e4m3fn, 656)]) -def test_backend_resident_payload_graph_replay(dtype, width): +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("num_rows", [0, 1, 4, 9]) +def test_backend_resident_payload_graph_replay(dtype, width, index_dtype, num_rows): if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("HiSparse resident gather currently targets Ascend 950") torch.npu.set_device(0) stride = 3 * width raw = torch.arange(7 * width * dtype.itemsize, dtype=torch.int64).to(torch.uint8).to("npu:0") source = raw.view(dtype).as_strided((2, 2, width), (stride, width, 1), storage_offset=width) - rows = torch.tensor([3, 0, 2, -1], dtype=torch.int32, device="npu:0") - destination = torch.empty((4, width), dtype=dtype, device="npu:0") + rows = torch.empty(num_rows, dtype=index_dtype, device="npu:0") + destination = torch.empty((num_rows, width), dtype=dtype, device="npu:0") + rows.copy_((torch.arange(num_rows, dtype=index_dtype, device="npu:0") % 5) - 1) backend = AscendHiSparseBackend() backend.gather_resident_rows(source, destination, rows) torch.npu.synchronize() + if not num_rows: + assert destination.numel() == 0 + return graph = torch.npu.NPUGraph() - with torch.npu.graph(graph): - backend.gather_resident_rows(source, destination, rows) - for values in ([3, 0, 2, -1], [0, 3, 1, 2]): - rows.copy_(torch.tensor(values, dtype=torch.int32, device="npu:0")) - raw.bitwise_xor_(1) - graph.replay() + try: + with torch.npu.graph(graph): + backend.gather_resident_rows(source, destination, rows) + for offset in (0, 2): + values = (torch.arange(num_rows, dtype=index_dtype) + offset) % 5 - 1 + if index_dtype == torch.int64: + values[0] = -(1 << 34) + rows.copy_(values.to("npu:0")) + raw.bitwise_xor_(1) + graph.replay() + torch.npu.synchronize() + safe = values.long().clamp_min(0) + expected = source.view(torch.uint8).cpu()[safe // 2, safe % 2] + assert torch.equal(destination.view(torch.uint8).cpu(), expected) + finally: torch.npu.synchronize() - safe = torch.tensor(values).clamp_min(0) - expected = source.view(torch.uint8).cpu()[safe // 2, safe % 2] - assert torch.equal(destination.view(torch.uint8).cpu(), expected) + graph.reset() @pytest.mark.parametrize("row_bytes", [37, 656]) @@ -136,13 +149,13 @@ def test_prefill_resident_patch_changing_extents(row_bytes, capture): source_view, destination_view = source[:, 0], destination[:, 0] indices.fill_(-1) - gather_valid_resident_rows(source_view, destination_view, indices) + gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) torch.npu.synchronize() graph = None if capture: graph = torch.npu.NPUGraph() with torch.npu.graph(graph): - gather_valid_resident_rows(source_view, destination_view, indices) + gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) for iteration in range(2): values = (torch.arange(rows, dtype=torch.int32) + iteration * 3) % (source_blocks * block + 2) - 1 current_source = source_cpu + iteration * 13 @@ -150,7 +163,7 @@ def test_prefill_resident_patch_changing_extents(row_bytes, capture): indices.copy_(values) destination.fill_(7) if graph is None: - gather_valid_resident_rows(source_view, destination_view, indices) + gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) else: graph.replay() torch.npu.synchronize() diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 15463aa5fb67..f48e783fb37e 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -223,14 +223,9 @@ def is_current_stream_capturing(self) -> bool: return torch.npu.is_current_stream_capturing() def gather_resident_rows(self, source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor) -> None: - # NPU advanced indexing does not support FP8. Copy all packed components - # as bytes, retaining Core's addressing and completion ownership. - fp8 = (torch.float8_e4m3fn, torch.float8_e5m2) - if source.dtype in fp8: - source = source.view(torch.uint8) - if destination.dtype in fp8: - destination = destination.view(torch.uint8) - super().gather_resident_rows(source, destination, row_ids) + from .cache_ops import gather_resident_rows + + gather_resident_rows(source, destination, row_ids) def swap_blocks_batch(self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor) -> None: if self._failed_dma_batches: diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 80d78c63c172..cda0e6f4ff37 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -16,8 +16,8 @@ from vllm_ascend.device.device_config import is_950 from vllm_ascend.hisparse.host_pool import MappedHostTensor from vllm_ascend.ops.triton.hisparse import ( + gather_resident_rows_kernel, gather_rows_kernel, - gather_valid_resident_rows_kernel, invalidate_kernel, invalidate_written_slots_sorted, pack_attention_rows_kernel, @@ -399,24 +399,35 @@ def hisparse_invalidate_written_slots( ) -def gather_valid_resident_rows(source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor) -> None: - """Fill only Resident rows in a staging cache that already holds Host rows.""" +def gather_resident_rows( + source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor, *, skip_invalid: bool = False +) -> None: + """Copy packed Resident rows using Core's negative-ID-to-row0 convention. + + With ``skip_invalid``, retain staged Host bytes at invalid IDs instead. + Source and destination may use different dtypes with the same byte width. + """ _check_packed_cache(source) - _check_packed_cache(destination) + if destination.ndim == 2: + destination = destination.unsqueeze(1) + if destination.ndim != 3: + raise ValueError("HiSparse Resident staging must contain flat or paged token rows") rows = destination.size(0) * destination.size(1) - _check_tensor(row_ids, "row_ids", source.device, (torch.int32,), shape=(rows,)) + _check_tensor(row_ids, "row_ids", source.device, (torch.int32, torch.int64), shape=(rows,)) + row_bytes = source.size(2) * source.element_size() if ( source.device != destination.device - or source.dtype != destination.dtype - or source.size(2) != destination.size(2) + or row_bytes != destination.size(2) * destination.element_size() or source.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr() ): raise ValueError("HiSparse Resident gather requires matching, separate source and staging rows") - row_bytes = source.size(2) * source.element_size() + if rows == 0: + return + _check_packed_cache(destination) block_rows = 8 block_bytes = min(1024, triton.next_power_of_2(row_bytes)) grid = (triton.cdiv(rows, block_rows), triton.cdiv(row_bytes, block_bytes)) - gather_valid_resident_rows_kernel[grid]( + gather_resident_rows_kernel[grid]( source.view(torch.uint8), destination.view(torch.uint8), row_ids, @@ -429,6 +440,7 @@ def gather_valid_resident_rows(source: torch.Tensor, destination: torch.Tensor, ROW_BYTES=row_bytes, BLOCK_ROWS=block_rows, BLOCK_BYTES=block_bytes, + SKIP_INVALID=skip_invalid, multibuffer=False, ) diff --git a/vllm_ascend/hisparse/prefill.py b/vllm_ascend/hisparse/prefill.py index 35a997954bd7..2c169fdedb9d 100644 --- a/vllm_ascend/hisparse/prefill.py +++ b/vllm_ascend/hisparse/prefill.py @@ -13,7 +13,7 @@ from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan -from vllm_ascend.hisparse.cache_ops import gather_valid_resident_rows +from vllm_ascend.hisparse.cache_ops import gather_resident_rows _PREFILL_LOOKAHEAD = "ascend_hisparse_prefill_lookahead" @@ -165,7 +165,7 @@ def stage_prefill_rows(group, layer_index, metadata): slot.tensor.record_stream(main) # The follower has now written its current KV. Patch its Resident # rows without overwriting the prefetched Host bytes. - gather_valid_resident_rows(cache.view.cache, slot.tensor, plan.gpu_row_ids[0]) + gather_resident_rows(cache.view.cache, slot.tensor, plan.gpu_row_ids[0], skip_invalid=True) staged, table = slot.tensor, plan.block_table req_ids = metadata.req_id_per_token[metadata.num_decode_tokens :] if metadata.num_decodes: diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 790d0633c57b..7418c8388cc9 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -215,7 +215,7 @@ def scatter_packed_kv_kernel( @triton.jit -def gather_valid_resident_rows_kernel( +def gather_resident_rows_kernel( source, destination, row_ids, @@ -228,9 +228,12 @@ def gather_valid_resident_rows_kernel( ROW_BYTES: tl.constexpr, BLOCK_ROWS: tl.constexpr, BLOCK_BYTES: tl.constexpr, + SKIP_INVALID: tl.constexpr, ): row = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64) src = tl.load(row_ids + row, mask=row < ROWS, other=-1).to(tl.int64) + if not SKIP_INVALID: + src = tl.maximum(src, 0) src_offset = src // SOURCE_BLOCK * SOURCE_STRIDE + src % SOURCE_BLOCK * ROW_BYTES dst_offset = row // DEST_BLOCK * DEST_STRIDE + row % DEST_BLOCK * ROW_BYTES byte = tl.program_id(1) * BLOCK_BYTES + tl.arange(0, BLOCK_BYTES) From 8d044bc1d0a7d5466f2e04af317a7dfcdfdda950 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Tue, 10 Nov 2026 16:43:40 +0800 Subject: [PATCH 088/120] perf(hisparse): use column-first compact gather launches Schedule compact transfer columns before request rows without changing copied values, stream dependencies, source ownership or native kernels. The frozen a11/A383 A5 model control used Core39, CANN9.2.1 and fixed 32K input/8K output at concurrency8. ON N2 mean TPOT fell from18.1353ms to17.5150ms (3.42%), with complete input/output and MTP matches. Same-source candidate ON/OFF N2 still showed9.75% miss-path regression. These measurements retain the A383 source/native scope; later cleanup-head model qualification is pending. Both compact launch/function bodies match that tested source. (cherry picked from commit a3832cb716298ec62eed8b8c338b0973e2da6783) Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/cache_ops.py | 2 +- vllm_ascend/ops/triton/hisparse.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index cda0e6f4ff37..3211812ab0cf 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -296,7 +296,7 @@ def _gather( if rows == 0 or topk == 0: return parallel_columns = min(topk, 128) if counts is not None else 0 - gather_rows_kernel[(rows, parallel_columns or topk)]( + gather_rows_kernel[(parallel_columns, rows) if parallel_columns else (rows, topk)]( self.host_mapping.address_of(host_cache), hot_cache.view(torch.uint8), source, diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 7418c8388cc9..4752d12b33cd 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -30,7 +30,7 @@ def gather_rows_kernel( BLOCK: tl.constexpr, PARALLEL_COLUMNS: tl.constexpr = 0, ): - row = tl.program_id(0) + row = tl.program_id(1) if PARALLEL_COLUMNS > 0 else tl.program_id(0) offset = tl.arange(0, BLOCK) padding = tl.load(state + row) < 0 if HAS_STATE else False if PARALLEL_COLUMNS > 0: @@ -40,7 +40,7 @@ def gather_rows_kernel( # Keep enough independent programs for small misses and retain the # original contiguous byte-copy shape instead of a wide 2-D payload. if ~padding: - for col in range(tl.program_id(1), count, PARALLEL_COLUMNS): + for col in range(tl.program_id(0), count, PARALLEL_COLUMNS): index = row * TOP_K + col src = tl.load(source + index) dst = tl.load(destination + index) From aa7f5dd79aef1c4c047a2934956ff663c4fa7c83 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 03:33:21 +0800 Subject: [PATCH 089/120] fix(worker): preserve the metadata graph capture API Pass shared MLA RoPE context only through the runtime build interface. Exercise the inherited SFA and Core indexer capture entry points in the existing regression test. Signed-off-by: zhenwenqi2024 --- tests/ut/worker/test_attn_utils_v2.py | 11 ++++++----- vllm_ascend/worker/v2/attn_utils.py | 6 +++++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index 4e26c27847b5..debd5326ee31 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -9,6 +9,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache +from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, @@ -112,18 +113,18 @@ def build(self, common_prefix_len, common_attn_metadata, shared_mla_rope=None, * # The unrelated consumer owns its copies, like a separate builder. return SimpleNamespace(cos=cos.clone() if self.offset else cos, sin=sin.clone() if self.offset else sin) - def build_for_cudagraph_capture(self, common_attn_metadata, **kwargs): - return self.build(0, common_attn_metadata, **kwargs) - class PlainSFA(Recorder, AscendSFAMetadataBuilder): pass class PlainIndexer(Recorder, AscendSFAIndexerMetadataBuilder): pass + class UnrelatedBuilder(Recorder, AttentionMetadataBuilder): + pass + builders = [("sfa", PlainSFA())] if boundary == "builder": - builders.append(("other", Recorder(5))) + builders.append(("other", UnrelatedBuilder(5))) builders.append(("indexer", PlainIndexer())) groups = [ [SimpleNamespace(layer_names=[name], get_metadata_builder=lambda _, builder=builder: builder)] @@ -168,7 +169,7 @@ def update_positions(group, common=None): for name in ("sfa", "indexer"): torch.testing.assert_close(metadata[name].cos, expected) torch.testing.assert_close(metadata[name].sin, -expected) - expected_lookups = 1 if boundary == "none" else len(builders) + expected_lookups = 1 if boundary == "none" and not for_capture else len(builders) assert len(lookups) == expected_lookups # The same position tensor is updated between calls; no prior-build hit. positions.add_(1) diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index ef638c816189..562de2499e1a 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -683,7 +683,11 @@ def build_attn_metadata( if not for_cudagraph_capture and model_specific_attn_metadata is not None else {} ) - if isinstance(attn_metadata_builder, (AscendSFAMetadataBuilder, AscendSFAIndexerMetadataBuilder)): + # Core's capture entry point does not accept the shared RoPE context. + # Keep sharing on the runtime build interface that declares it. + if not for_cudagraph_capture and isinstance( + attn_metadata_builder, (AscendSFAMetadataBuilder, AscendSFAIndexerMetadataBuilder) + ): if shared_mla_rope is None: shared_mla_rope = MLARopeMetadataCache() attn_metadata_extra_kwargs["shared_mla_rope"] = shared_mla_rope From 1f3b0d42487aa7984d59513a55f21efd78885f75 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 05:57:26 +0800 Subject: [PATCH 090/120] perf(hisparse): defer direct target mirror wait until after drafting Reuse Core live HMA layer/page ownership to keep Target row DMA sources stable during the same draft. Gate on the bound source layout, draft group and Host ranges, then retain Core finish ordering for incomplete plans. Preserve the output-copy durability fence and Core finalize/begin/failure joins. No payload snapshot, extra workspace, row DMA rewrite, Core change or new environment option. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/host_mirror.py | 275 ++++++++++++++++++++++++++ vllm_ascend/worker/v2/model_runner.py | 234 ++++++++++++---------- 2 files changed, 402 insertions(+), 107 deletions(-) create mode 100644 vllm_ascend/hisparse/host_mirror.py diff --git a/vllm_ascend/hisparse/host_mirror.py b/vllm_ascend/hisparse/host_mirror.py new file mode 100644 index 000000000000..aa78177bdb56 --- /dev/null +++ b/vllm_ascend/hisparse/host_mirror.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Order direct HiSparse Host writeback against drafting and output completion.""" + +from __future__ import annotations + +import torch + + +def _host_disjoint(caches, draft_layers): + if any(cache.ndim != 2 or not cache.is_contiguous() or cache.numel() == 0 for cache in caches): + return False + spans = [(cache.data_ptr(), cache.data_ptr() + cache.numel() * cache.element_size()) for cache in caches] + return all( + left_end <= right_start or right_end <= left_start + for index, (left_start, left_end) in enumerate(spans) + if index in draft_layers + for other, (right_start, right_end) in enumerate(spans) + if other not in draft_layers + ) + + +def _draft_groups_disjoint(handles, draft_layers): + draft_runtimes = {id(handles[index].runtime) for index in draft_layers} + # A draft-led resolver can refill every follower from Host in one call. + for index in draft_layers: + group = handles[index].runtime.index_group + if id(group.leader) not in draft_runtimes or any( + id(follower) not in draft_runtimes for follower in group.followers + ): + return False + return True + + +def _source_layout_valid(worker): + """Check per-page layer ranges; HMA owns live physical block IDs across groups.""" + backing = worker.hot_backing + if backing.ndim != 1 or not backing.is_contiguous() or backing.storage_offset() != 0: + return False + base = backing.data_ptr() + total = backing.numel() * backing.element_size() + stride = None + ranges = [] + for handle, cache in zip(worker.cache_handles, worker.resident_caches, strict=True): + if cache.ndim != 3: + return False + hot = handle.runtime.hot.cache + item = cache.element_size() + row = cache.shape[-1] * item + page = worker.kernel_block_size * row + step = cache.stride(0) * item + offset = cache.data_ptr() - base + group = handle.runtime.resident_source_index + if ( + cache.shape[1] != worker.kernel_block_size + or cache.stride(1) * item != row + or cache.stride(2) != 1 + or cache.device != backing.device + or cache.untyped_storage().data_ptr() != base + or step <= 0 + or offset < 0 + or offset + page > step + or cache.shape[0] * step > total + or group < 0 + or (stride is not None and stride != step) + or cache.data_ptr() != hot.data_ptr() + or cache.dtype != hot.dtype + or cache.shape != hot.shape + or cache.stride() != hot.stride() + ): + return False + stride = step + ranges.append((group, offset, offset + page)) + # Within a group the same physical page can serve Target and draft, but + # their layer slices must not overlap. Different groups' live writable + # pages are distinct allocations in Core's common BlockPool. + return all( + left[0] != right[0] or left[2] <= right[1] or right[2] <= left[1] + for index, left in enumerate(ranges) + if index in worker._draft_layers + for other, right in enumerate(ranges) + if other not in worker._draft_layers + ) + + +class _DirectMirrorSources: + """Retain one canonical worker binding and recognize its complete Target plan.""" + + def __init__(self, worker, targets): + self._sources = worker.resident_caches + self._destinations = worker.host_caches + self._targets = targets + + @classmethod + def bind(cls, worker, speculator): + names = worker.cache_layer_names + draft = set(worker._draft_layers) + expected = set(getattr(speculator, "draft_attn_layer_names", ())) & set(names) + targets = tuple(index for index in range(len(names)) if index not in draft) + if ( + not draft + or not targets + or {names[index] for index in draft} != expected + or not _draft_groups_disjoint(worker.cache_handles, draft) + or not _host_disjoint(worker.host_caches, draft) + or not _source_layout_valid(worker) + ): + return None + return cls(worker, targets) + + def can_defer(self, worker): + return ( + worker.resident_caches is self._sources + and worker.host_caches is self._destinations + and not worker._dma_submitted + and not worker._per_layer_mirrored + and not worker._submitted_mirror_layers + and tuple( + index + for index, cache in enumerate(worker.cache_handles) + if cache.num_actual_tokens and index not in worker._draft_layers + ) + == self._targets + ) + + +class HostMirrorConnector: + """Join one worker's Host mirror at the existing Core draft finalizer.""" + + def __init__(self, connector, worker, output_copy_stream, sources, logger=None, rank=None): + self._connector = connector + self._worker = worker + self._output_copy_stream = output_copy_stream + self._sources = sources + self._logger = logger + self._rank = rank + self._logged_draft_defer = False + self._logged_finalize = False + self._overlap_draft = False + self._pending_event = None + self._allow_step = False + self._eligible_forward = False + + def __getattr__(self, name): + return getattr(self._connector, name) + + def join_pending(self) -> None: + if self._pending_event is not None: + torch.accelerator.current_stream().wait_event(self._pending_event) + self._pending_event = None + self._overlap_draft = False + + def begin_execution(self, allow_step: bool) -> None: + # The runner calls this before preemption or any scheduler state writes. + self.join_pending() + self._allow_step = allow_step + self._eligible_forward = False + + def pre_forward(self, scheduler_output, input_batch=None, attn_metadata=None): + self._eligible_forward = ( + self._allow_step + and scheduler_output.has_sync_kv_loads + and input_batch is not None + and attn_metadata is not None + and not input_batch.has_prefill + and 0 < input_batch.num_tokens == input_batch.num_tokens_after_padding + ) + return self._connector.pre_forward(scheduler_output, input_batch, attn_metadata=attn_metadata) + + def finish_forward(self) -> None: + eligible = self._eligible_forward + self._eligible_forward = False + worker = self._worker + if not eligible or self._connector._disabled or not self._sources.can_defer(worker): + return self._connector.finish_forward() + + compute_stream = torch.accelerator.current_stream() + worker._forward_ready_event.record(compute_stream) + pending_before = len(worker._pending_dma_descriptors) + try: + worker._finish_mirror_phase(worker._forward_ready_event) + except BaseException: + # The runner's failure boundary joins an already-recorded copy. + if worker._dma_submitted: + self._pending_event = worker.host_write_event + raise + if worker._dma_submitted: + self._pending_event = worker.host_write_event + worker._dma_submitted = False + self._overlap_draft = len(worker._pending_dma_descriptors) == pending_before + 1 + if self._overlap_draft: + # AsyncOutput must not report completion before Host KV is durable. + self._output_copy_stream.wait_event(self._pending_event) + else: + self.join_pending() + else: + worker.host_write_event.record(compute_stream) + worker._release_completed_dma_descriptors() + + def post_forward(self, finished_req_ids, defer_finalize=False): + if not defer_finalize or not self._overlap_draft: + self.join_pending() + elif not self._logged_draft_defer and self._logger is not None: + self._logger.info( + "HiSparse direct Target mirror draft defer used: rank=%d; source-safe=HMA layer/page ownership; " + "Host-durable=DMA host_write_event; join=after draft", + self._rank, + ) + self._logged_draft_defer = True + return self._connector.post_forward(finished_req_ids, defer_finalize=defer_finalize) + + def finalize_step(self): + deferred = self._pending_event is not None and self._overlap_draft + self.join_pending() + if deferred and not self._logged_finalize and self._logger is not None: + self._logger.info( + "HiSparse direct Target mirror Host durability wait enqueued after draft: rank=%d", self._rank + ) + self._logged_finalize = True + return self._connector.finalize_step() + + def no_forward(self, scheduler_output): + self.begin_execution(False) + return self._connector.no_forward(scheduler_output) + + def set_disabled(self, disabled: bool) -> None: + self.begin_execution(False) + self._connector.set_disabled(disabled) + + def reset_capture_state(self) -> None: + self.begin_execution(False) + self._connector.reset_capture_state() + + +def overlap_target_host_mirror(runner): + """Use Core's private Host writer and live HMA binding, or keep its original fence.""" + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import HiSparseConnectorWorker + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + from vllm.logger import init_logger + from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + from vllm_ascend.hisparse.backend import AscendHiSparseBackend + + connector = runner.kv_connector + if type(connector) is not ActiveKVConnector or type(connector.kv_connector) is not HiSparseConnector: + return connector + worker = connector.kv_connector.connector_worker + if ( + type(worker) is not HiSparseConnectorWorker + or not worker._initialized + or not isinstance(worker.backend, AscendHiSparseBackend) + or worker.shared_host_region is not None + or not worker.is_host_writer + or worker.hot_backing.device.type != "npu" + ): + return connector + sources = _DirectMirrorSources.bind(worker, runner.speculator) + if sources is None: + return connector + output = runner.output_copy_stream + output_stream = torch.Stream( + stream_id=output.stream_id, device_index=output.device_index, device_type=output.device_type + ) + compute_stream = torch.accelerator.current_stream() + if ( + output_stream.device != compute_stream.device + or compute_stream.device != worker.hot_backing.device + or output_stream == compute_stream + ): + return connector + logger = init_logger(__name__) + rank = get_tensor_model_parallel_rank() + logger.info("HiSparse direct Target Host mirror overlap enabled: rank=%d", rank) + return HostMirrorConnector(connector, worker, output_stream, sources, logger, rank) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index e1cb6b8072af..09c490b70070 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -322,27 +322,32 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: @runtime_guard_sample_tokens def sample_tokens(self, grammar_output): - if ( - lmhead_tp_enable() - and self.prompt_logprobs_worker is not None - and ( - self.prompt_logprobs_worker.uses_prompt_logprobs.any() - or getattr(self.prompt_logprobs_worker, "token_id_scores", {}) - ) - ): - # The prompt-logprobs worker issues a second compute_logits with - # unpadded rows that desyncs the LM-head collectives and hangs. - # vLLM #54335 adds fixed-token scoring through the same unpadded - # compute_logits path, so it has the same collective restriction. - raise NotImplementedError( - "prompt logprobs and fixed-token prompt scoring are not supported with lmhead TP." - ) + mirror = getattr(self, "_host_mirror", None) + try: + if ( + lmhead_tp_enable() + and self.prompt_logprobs_worker is not None + and ( + self.prompt_logprobs_worker.uses_prompt_logprobs.any() + or getattr(self.prompt_logprobs_worker, "token_id_scores", {}) + ) + ): + # The prompt-logprobs worker issues a second compute_logits with + # unpadded rows that desyncs the LM-head collectives and hangs. + # vLLM #54335 adds fixed-token scoring through the same unpadded + # compute_logits path, so it has the same collective restriction. + raise NotImplementedError( + "prompt logprobs and fixed-token prompt scoring are not supported with lmhead TP." + ) - output = super().sample_tokens(grammar_output) - # Close the per-step msprobe dump cycle on the normal last-PP-rank - # path (v1 parity: v1 finalizes at the end of sample_tokens). - self._finalize_dump_data() - return output + output = super().sample_tokens(grammar_output) + # Close the per-step msprobe dump cycle on the normal last-PP-rank + # path (v1 parity: v1 finalizes at the end of sample_tokens). + self._finalize_dump_data() + return output + finally: + if mirror is not None: + mirror.join_pending() def pool(self): output = super().pool() @@ -519,6 +524,12 @@ def initialize_kv_cache( self.model_state.kvpp_runtime = self.kvpp self._fill_c8_mxfp_v_scale_caches() + if self.vllm_config.attention_config.hisparse_config is not None: + from vllm_ascend.hisparse.host_mirror import overlap_target_host_mirror + + connector = overlap_target_host_mirror(self) + self._host_mirror = connector if connector is not self.kv_connector else None + self.kv_connector = connector def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" @@ -610,97 +621,106 @@ def execute_model( valid_dummy_state_slots: bool = False, randomize_inputs: bool = False, ): - self._cpp_execution_time_ms = None - profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config - execution_start_time = _start_profiling_chunk_timing( - profiling_config, - scheduler_output, - ) + mirror = getattr(self, "_host_mirror", None) + if mirror is not None: + mirror.begin_execution(not dummy_run and not is_profile) + try: + self._cpp_execution_time_ms = None + profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config + execution_start_time = _start_profiling_chunk_timing( + profiling_config, + scheduler_output, + ) - # Preemption stores must complete before the parent updates states and - # reuses or zeroes the preempted requests' physical KV cache blocks. - if has_kv_transfer_group(): - kv_connector_metadata = scheduler_output.kv_connector_metadata - assert kv_connector_metadata is not None - get_kv_transfer_group().handle_preemptions(kv_connector_metadata) - - dump_forward = dummy_run or scheduler_output.total_num_scheduled_tokens > 0 - # Eager PrecisionDebugger is started only for real scheduler steps. - # Graph dumping is already active from load_model(), so dummy forwards - # only need to flush the graph dumper below without writing a step. - if dump_forward and not dummy_run: - self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) - - self.model_state.kvpp_is_dummy_run = dummy_run or is_profile - dp_coordination_context = ( - skip_dp_coordination() if should_skip_allreduce_across_dp_group(self.vllm_config) else nullcontext() - ) - profile_kwargs = {"randomize_inputs": randomize_inputs} - forward_failed = True - with dp_coordination_context: - try: - output = super().execute_model( - scheduler_output, - intermediate_tensors=intermediate_tensors, - dummy_run=dummy_run, - skip_attn_for_dummy_run=skip_attn_for_dummy_run, - is_profile=is_profile, - context_len=context_len, - valid_dummy_state_slots=valid_dummy_state_slots, - **profile_kwargs, - ) - forward_failed = False - finally: - finish_execution = getattr(self.model_state, "finish_execution", None) - if finish_execution is not None: - finish_execution(failed=forward_failed) - if not dummy_run and not is_profile: - self._maybe_send_draft_kv(scheduler_output) - self.model_state.kvpp_is_dummy_run = False - if dummy_run and lmhead_tp_enable() and not is_profile and self.is_last_pp_rank: - # lmhead TP: idle ranks never call sample(); join the target head - # here at capacity, before _dummy_run replays the dummy propose. - if self.execute_model_state is None: - raise RuntimeError( - "lmhead TP dummy join expects execute_model_state published by the upstream dummy execute_model." + # Preemption stores must complete before the parent updates states and + # reuses or zeroes the preempted requests' physical KV cache blocks. + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + assert kv_connector_metadata is not None + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) + + dump_forward = dummy_run or scheduler_output.total_num_scheduled_tokens > 0 + # Eager PrecisionDebugger is started only for real scheduler steps. + # Graph dumping is already active from load_model(), so dummy forwards + # only need to flush the graph dumper below without writing a step. + if dump_forward and not dummy_run: + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) + + self.model_state.kvpp_is_dummy_run = dummy_run or is_profile + dp_coordination_context = ( + skip_dp_coordination() if should_skip_allreduce_across_dp_group(self.vllm_config) else nullcontext() + ) + profile_kwargs = {"randomize_inputs": randomize_inputs} + forward_failed = True + with dp_coordination_context: + try: + output = super().execute_model( + scheduler_output, + intermediate_tensors=intermediate_tensors, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + is_profile=is_profile, + context_len=context_len, + valid_dummy_state_slots=valid_dummy_state_slots, + **profile_kwargs, + ) + forward_failed = False + finally: + finish_execution = getattr(self.model_state, "finish_execution", None) + if finish_execution is not None: + finish_execution(failed=forward_failed) + if not dummy_run and not is_profile: + self._maybe_send_draft_kv(scheduler_output) + self.model_state.kvpp_is_dummy_run = False + if dummy_run and lmhead_tp_enable() and not is_profile and self.is_last_pp_rank: + # lmhead TP: idle ranks never call sample(); join the target head + # here at capacity, before _dummy_run replays the dummy propose. + if self.execute_model_state is None: + raise RuntimeError( + "lmhead TP dummy join expects execute_model_state published by " + "the upstream dummy execute_model." + ) + dummy_indices = torch.zeros( + self._lmhead_tp_max_num_logits(), + dtype=torch.int64, + device=self.device, ) - dummy_indices = torch.zeros( - self._lmhead_tp_max_num_logits(), - dtype=torch.int64, - device=self.device, + self.model.compute_logits(self.execute_model_state.hidden_states[dummy_indices]) + self.kvpp.complete_forward() + + if dump_forward: + if dummy_run: + self._finalize_dump_data(dump=False) + elif isinstance(output, IntermediateTensors) and self.is_pooling_model: + # Non-last pooling ranks do not reach pool(). Decoding ranks + # finalize after sample_tokens(), including non-last PP ranks. + self._finalize_dump_data() + + self._cpp_execution_time_ms = _finish_profiling_chunk_timing( + profiling_config, + execution_start_time, ) - self.model.compute_logits(self.execute_model_state.hidden_states[dummy_indices]) - self.kvpp.complete_forward() - - if dump_forward: - if dummy_run: - self._finalize_dump_data(dump=False) - elif isinstance(output, IntermediateTensors) and self.is_pooling_model: - # Non-last pooling ranks do not reach pool(). Decoding ranks - # finalize after sample_tokens(), including non-last PP ranks. - self._finalize_dump_data() - - self._cpp_execution_time_ms = _finish_profiling_chunk_timing( - profiling_config, - execution_start_time, - ) - # TODO: Remove this workaround once upstream vLLM restores idle DP ranks' - # logical request counts before draft DP synchronization. - # In adaptive verification, idle DP ranks dummy-run at target graph capacity. - # Restore their request counts before draft DP sync to avoid inflating all ranks. - if ( - dummy_run - and not is_profile - and context_len == 0 - and self.adaptive_verification is not None - and self.execute_model_state is not None - ): - state = self.execute_model_state - self.execute_model_state = state._replace( - input_batch=replace(state.input_batch, num_reqs=len(scheduler_output.num_scheduled_tokens)) - ) - return output + # TODO: Remove this workaround once upstream vLLM restores idle DP ranks' + # logical request counts before draft DP synchronization. + # In adaptive verification, idle DP ranks dummy-run at target graph capacity. + # Restore their request counts before draft DP sync to avoid inflating all ranks. + if ( + dummy_run + and not is_profile + and context_len == 0 + and self.adaptive_verification is not None + and self.execute_model_state is not None + ): + state = self.execute_model_state + self.execute_model_state = state._replace( + input_batch=replace(state.input_batch, num_reqs=len(scheduler_output.num_scheduled_tokens)) + ) + return output + except BaseException: + if mirror is not None: + mirror.join_pending() + raise def _maybe_send_draft_kv(self, scheduler_output: SchedulerOutput) -> None: """Write draft KV on P, then let D pull the exact allocated cache pages.""" From f4a1dbb46bc6d224edd74b747f85ed14e5c13d6d Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 06:10:11 +0800 Subject: [PATCH 091/120] test(hisparse): cover direct Target mirror during Draft writes Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 175 ++++++++++++++++++ 1 file changed, 175 insertions(+) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index adab4001bc1b..adeb88c99370 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -873,3 +873,178 @@ def write(): finally: if graph is not None: graph.reset() + + +@pytest.mark.parametrize("graph_producer,shared_group", [(False, True), (True, False)]) +def test_direct_target_mirror_survives_legal_draft_writes(graph_producer, shared_group, record_property): + from types import SimpleNamespace + + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import HiSparseConnectorWorker + from vllm.v1.core.block_pool import BlockPool + from vllm.v1.hisparse import runtime + from vllm.v1.hisparse.runtime import HiSparseCacheHandle, PagedCacheView + from vllm.v1.hisparse.types import SparseKVRowMirror + from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + from vllm_ascend.hisparse.host_mirror import HostMirrorConnector, _DirectMirrorSources + + width, block, gpu_blocks, host_blocks = 656, 128, 5, 3 + page = block * width + stride = 3 * page + names = ("target0", "target1", "draft") + layer_offsets = (0, page, 2 * page if shared_group else 0) + source_groups = (0, 0, 0 if shared_group else 1) + pool = BlockPool(gpu_blocks, enable_caching=False, hash_block_size=block) + target_pages = pool.get_new_blocks(2) + draft_pages = target_pages if shared_group else pool.get_new_blocks(2) + target_ids = [page.block_id for page in target_pages] + draft_ids = [page.block_id for page in draft_pages] + assert (set(target_ids).isdisjoint(draft_ids)) is not shared_group + spans = ((126, 2), (2, 1)) + mirrors = tuple( + SparseKVRowMirror( + (target_ids[i] * block + offset,) + if shared_group + else (target_ids[i] * block + offset, draft_ids[i] * block + offset), + (i + 1) * block + offset, + count, + ) + for i, (offset, count) in enumerate(spans) + ) + backend = AscendHiSparseBackend() + host_rows = host_blocks * block + payload, registered = runtime.allocate_pinned_host_pool(len(names) * host_rows * width) + graph = None + worker = None + try: + assert payload.is_pinned() and registered.is_pinned() + mapping = registered._hisparse_registration + # Include both page and layer in the CPU oracle. Plain uint8 arange + # repeats across these aligned pages and could miss a wrong page ID. + byte = torch.arange(gpu_blocks * stride, dtype=torch.int32) + cpu_source = (byte + (byte // stride) * 19 + ((byte % stride) // page) * 31).to(torch.uint8) + source = cpu_source.to("npu") + staging = source.clone() + states = torch.zeros(1, dtype=torch.int32, device=source.device) + written = torch.tensor( + [m.destination_start + j for m in mirrors for j in range(m.num_rows)], + dtype=torch.int32, + device=source.device, + ) + requests = torch.zeros_like(written) + owners, handles = [], [] + for index, (group, offset) in enumerate(zip(source_groups, layer_offsets)): + span = slice(index * host_rows * width, (index + 1) * host_rows * width) + view = PagedCacheView.bind( + source, + dtype=torch.uint8, + row_width=width, + byte_offset=offset, + block_stride=stride, + num_blocks=gpu_blocks, + block_size=block, + ) + owner = SimpleNamespace( + resident_source_index=group, + request_state_indices=states, + hot=view, + host_cache=payload[span].view(host_rows, width), + eager_host_mirror=True, + is_group_leader=index != 1, + cache_ops=HiSparseNPUOps(mapping), + ) + handle = HiSparseCacheHandle(owner) + handle.view = view + handle.slot_mapping = torch.tensor( + [m.source_starts[group] + j for m in mirrors for j in range(m.num_rows)], + dtype=torch.int32, + device=source.device, + ) + handle.draft_layer = index == 2 + owners.append(owner) + handles.append(handle) + for leader, followers in ((owners[0], [owners[1]]), (owners[2], [])): + group = SimpleNamespace( + leader=leader, + followers=followers, + device_global_indices=torch.full((1, 16), -1, dtype=torch.int32, device=source.device), + ) + for owner in [leader, *followers]: + owner.index_group = group + worker = HiSparseConnectorWorker.__new__(HiSparseConnectorWorker) + worker.backend, worker._initialized = backend, False + worker.vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + num_lookahead_tokens=1, + kernel_config=SimpleNamespace(enable_jit_warmup=False), + ) + worker.initialize(handles, list(names), source, 1, host_blocks, source.device, [registered]) + sources = _DirectMirrorSources.bind(worker, SimpleNamespace(draft_attn_layer_names={"draft"})) + assert sources is not None + connector = HiSparseConnector.__new__(HiSparseConnector) + connector.connector_worker = worker + inner = ActiveKVConnector.__new__(ActiveKVConnector) + inner.kv_connector, inner._disabled = connector, False + output_stream = torch.Stream(device=source.device) + proxy = HostMirrorConnector(inner, worker, output_stream, sources) + values = torch.full((written.numel(), width), 237, dtype=torch.uint8, device=source.device) + valid = torch.tensor([written.numel()], dtype=torch.int32, device=source.device) + if graph_producer: + source.copy_(staging) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + source.copy_(staging) + pending_observed = 0 + for change in (1, 3, 7): + proxy.begin_execution(True) + proxy._eligible_forward = True # Fixed decode metadata for this transport-only test. + worker._set_row_mirrors(mirrors) + for handle in handles: + handle.num_actual_tokens = handle.num_decode_tokens = written.numel() + handle.decode_batch = handle.host_mirror_required = True + handle.mirror_slot_mapping, handle.req_id_per_token = written, requests + payload.fill_(-1) + cpu_source.bitwise_xor_(change) + staging.copy_(cpu_source, non_blocking=True) + source.copy_(staging) if graph is None else graph.replay() + expected = torch.full_like(payload, -1).view(torch.uint8).view(len(names), host_rows, width) + for i, (offset, count) in enumerate(spans): + for layer in (0, 1): + start = target_ids[i] * stride + layer_offsets[layer] + offset * width + end = start + count * width + dest = mirrors[i].destination_start + expected[layer, dest : dest + count] = cpu_source[start:end].view(count, width) + proxy.finish_forward() # Real Target-only descriptor DMA from Resident pages. + assert proxy._overlap_draft and proxy._pending_event is worker.host_write_event + draft = handles[2] + draft.num_actual_tokens = written.numel() + destination, slots, count = draft.write_target(written.numel(), written.numel()) + assert count == written.numel() + # Submit only the actual draft layer's writable rows; do not join + # or overwrite the full slab before this same-round write. + scatter_packed_kv(values, values[:, :0], destination, slots, valid) + pending_observed += int(not worker.host_write_event.query()) + proxy.finalize_step() + output_done = torch.Event() + output_done.record(output_stream) + output_done.synchronize() + torch.testing.assert_close(payload.view(torch.uint8).view_as(expected), expected, rtol=0, atol=0) + for draft_block, (offset, size) in zip(draft_ids, spans): + actual = destination[draft_block, offset : offset + size].cpu() + torch.testing.assert_close(actual, torch.full_like(actual, 237), rtol=0, atol=0) + assert torch.npu.is_initialized() + record_property("actual_NPU_initialized", True) + record_property("direct_mirror_submissions", 3) + record_property("shared_resident_group", shared_group) + record_property("host_event_pending_after_draft_submit", pending_observed) + finally: + torch.npu.synchronize() + dma_stream = getattr(worker, "dma_stream", None) + if dma_stream is not None: + dma_stream.synchronize() + if graph is not None: + graph.reset() + backend.unpin_tensor(registered) + pool.free_blocks(target_pages if shared_group else [*target_pages, *draft_pages]) From fecc2168ecdf05fd0c8c9ead6657cba40de7fcf3 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 09:48:44 +0800 Subject: [PATCH 092/120] refactor(hisparse): keep the mirror join at the sample return Keep the deferred event owned by the connector when sampling raises. The next execution joins before block reuse, and connector shutdown drains DMA before freeing pinned storage. Signed-off-by: zhenwenqi2024 --- vllm_ascend/worker/v2/model_runner.py | 47 +++++++++++++-------------- 1 file changed, 22 insertions(+), 25 deletions(-) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 09c490b70070..2f06e0b82cc2 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -322,32 +322,29 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: @runtime_guard_sample_tokens def sample_tokens(self, grammar_output): - mirror = getattr(self, "_host_mirror", None) - try: - if ( - lmhead_tp_enable() - and self.prompt_logprobs_worker is not None - and ( - self.prompt_logprobs_worker.uses_prompt_logprobs.any() - or getattr(self.prompt_logprobs_worker, "token_id_scores", {}) - ) - ): - # The prompt-logprobs worker issues a second compute_logits with - # unpadded rows that desyncs the LM-head collectives and hangs. - # vLLM #54335 adds fixed-token scoring through the same unpadded - # compute_logits path, so it has the same collective restriction. - raise NotImplementedError( - "prompt logprobs and fixed-token prompt scoring are not supported with lmhead TP." - ) + if ( + lmhead_tp_enable() + and self.prompt_logprobs_worker is not None + and ( + self.prompt_logprobs_worker.uses_prompt_logprobs.any() + or getattr(self.prompt_logprobs_worker, "token_id_scores", {}) + ) + ): + # The prompt-logprobs worker issues a second compute_logits with + # unpadded rows that desyncs the LM-head collectives and hangs. + # vLLM #54335 adds fixed-token scoring through the same unpadded + # compute_logits path, so it has the same collective restriction. + raise NotImplementedError( + "prompt logprobs and fixed-token prompt scoring are not supported with lmhead TP." + ) - output = super().sample_tokens(grammar_output) - # Close the per-step msprobe dump cycle on the normal last-PP-rank - # path (v1 parity: v1 finalizes at the end of sample_tokens). - self._finalize_dump_data() - return output - finally: - if mirror is not None: - mirror.join_pending() + output = super().sample_tokens(grammar_output) + # Close the per-step msprobe dump cycle on the normal last-PP-rank + # path (v1 parity: v1 finalizes at the end of sample_tokens). + self._finalize_dump_data() + if (mirror := getattr(self, "_host_mirror", None)) is not None: + mirror.join_pending() + return output def pool(self): output = super().pool() From 5e64181d496bf56800dad00a6588a32da7c0129b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 10:03:17 +0800 Subject: [PATCH 093/120] refactor(hisparse): narrow runner host mirror hooks Use the current KV connector as the sole mirror binding and handle writeback submission failures in that connector. Keep the pending event owned by the connector after a later runner failure; the next step waits before preemption and shutdown drains DMA before unpinning. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/host_mirror.py | 27 ++-- vllm_ascend/worker/v2/model_runner.py | 192 ++++++++++++-------------- 2 files changed, 105 insertions(+), 114 deletions(-) diff --git a/vllm_ascend/hisparse/host_mirror.py b/vllm_ascend/hisparse/host_mirror.py index aa78177bdb56..4ea45ca18a09 100644 --- a/vllm_ascend/hisparse/host_mirror.py +++ b/vllm_ascend/hisparse/host_mirror.py @@ -179,23 +179,24 @@ def finish_forward(self) -> None: pending_before = len(worker._pending_dma_descriptors) try: worker._finish_mirror_phase(worker._forward_ready_event) + if worker._dma_submitted: + self._pending_event = worker.host_write_event + worker._dma_submitted = False + self._overlap_draft = len(worker._pending_dma_descriptors) == pending_before + 1 + if self._overlap_draft: + # AsyncOutput must not report completion before Host KV is durable. + self._output_copy_stream.wait_event(self._pending_event) + else: + self.join_pending() + else: + worker.host_write_event.record(compute_stream) + worker._release_completed_dma_descriptors() except BaseException: - # The runner's failure boundary joins an already-recorded copy. + # The mirror phase may fail after recording a submitted copy. if worker._dma_submitted: self._pending_event = worker.host_write_event + self.join_pending() raise - if worker._dma_submitted: - self._pending_event = worker.host_write_event - worker._dma_submitted = False - self._overlap_draft = len(worker._pending_dma_descriptors) == pending_before + 1 - if self._overlap_draft: - # AsyncOutput must not report completion before Host KV is durable. - self._output_copy_stream.wait_event(self._pending_event) - else: - self.join_pending() - else: - worker.host_write_event.record(compute_stream) - worker._release_completed_dma_descriptors() def post_forward(self, finished_req_ids, defer_finalize=False): if not defer_finalize or not self._overlap_draft: diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 2f06e0b82cc2..f19ba67ea5b7 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -95,6 +95,7 @@ allocate_kv_offload_topk_profile_buffers, init_sparse_kv_offload_manager, ) +from vllm_ascend.hisparse.host_mirror import HostMirrorConnector, overlap_target_host_mirror from vllm_ascend.models.deepseek_v41.cache_config import uses_a5_packed_cache from vllm_ascend.observability.runtime_guard.hooks import ( runtime_guard_sample_tokens, @@ -342,7 +343,7 @@ def sample_tokens(self, grammar_output): # Close the per-step msprobe dump cycle on the normal last-PP-rank # path (v1 parity: v1 finalizes at the end of sample_tokens). self._finalize_dump_data() - if (mirror := getattr(self, "_host_mirror", None)) is not None: + if isinstance(mirror := getattr(self, "kv_connector", None), HostMirrorConnector): mirror.join_pending() return output @@ -522,11 +523,7 @@ def initialize_kv_cache( self._fill_c8_mxfp_v_scale_caches() if self.vllm_config.attention_config.hisparse_config is not None: - from vllm_ascend.hisparse.host_mirror import overlap_target_host_mirror - - connector = overlap_target_host_mirror(self) - self._host_mirror = connector if connector is not self.kv_connector else None - self.kv_connector = connector + self.kv_connector = overlap_target_host_mirror(self) def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" @@ -618,106 +615,99 @@ def execute_model( valid_dummy_state_slots: bool = False, randomize_inputs: bool = False, ): - mirror = getattr(self, "_host_mirror", None) - if mirror is not None: + if isinstance(mirror := getattr(self, "kv_connector", None), HostMirrorConnector): mirror.begin_execution(not dummy_run and not is_profile) - try: - self._cpp_execution_time_ms = None - profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config - execution_start_time = _start_profiling_chunk_timing( - profiling_config, - scheduler_output, - ) + self._cpp_execution_time_ms = None + profiling_config = self.ascend_config.scheduler_config.profiling_chunk_config + execution_start_time = _start_profiling_chunk_timing( + profiling_config, + scheduler_output, + ) - # Preemption stores must complete before the parent updates states and - # reuses or zeroes the preempted requests' physical KV cache blocks. - if has_kv_transfer_group(): - kv_connector_metadata = scheduler_output.kv_connector_metadata - assert kv_connector_metadata is not None - get_kv_transfer_group().handle_preemptions(kv_connector_metadata) - - dump_forward = dummy_run or scheduler_output.total_num_scheduled_tokens > 0 - # Eager PrecisionDebugger is started only for real scheduler steps. - # Graph dumping is already active from load_model(), so dummy forwards - # only need to flush the graph dumper below without writing a step. - if dump_forward and not dummy_run: - self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) - - self.model_state.kvpp_is_dummy_run = dummy_run or is_profile - dp_coordination_context = ( - skip_dp_coordination() if should_skip_allreduce_across_dp_group(self.vllm_config) else nullcontext() - ) - profile_kwargs = {"randomize_inputs": randomize_inputs} - forward_failed = True - with dp_coordination_context: - try: - output = super().execute_model( - scheduler_output, - intermediate_tensors=intermediate_tensors, - dummy_run=dummy_run, - skip_attn_for_dummy_run=skip_attn_for_dummy_run, - is_profile=is_profile, - context_len=context_len, - valid_dummy_state_slots=valid_dummy_state_slots, - **profile_kwargs, - ) - forward_failed = False - finally: - finish_execution = getattr(self.model_state, "finish_execution", None) - if finish_execution is not None: - finish_execution(failed=forward_failed) - if not dummy_run and not is_profile: - self._maybe_send_draft_kv(scheduler_output) - self.model_state.kvpp_is_dummy_run = False - if dummy_run and lmhead_tp_enable() and not is_profile and self.is_last_pp_rank: - # lmhead TP: idle ranks never call sample(); join the target head - # here at capacity, before _dummy_run replays the dummy propose. - if self.execute_model_state is None: - raise RuntimeError( - "lmhead TP dummy join expects execute_model_state published by " - "the upstream dummy execute_model." - ) - dummy_indices = torch.zeros( - self._lmhead_tp_max_num_logits(), - dtype=torch.int64, - device=self.device, + # Preemption stores must complete before the parent updates states and + # reuses or zeroes the preempted requests' physical KV cache blocks. + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + assert kv_connector_metadata is not None + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) + + dump_forward = dummy_run or scheduler_output.total_num_scheduled_tokens > 0 + # Eager PrecisionDebugger is started only for real scheduler steps. + # Graph dumping is already active from load_model(), so dummy forwards + # only need to flush the graph dumper below without writing a step. + if dump_forward and not dummy_run: + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) + + self.model_state.kvpp_is_dummy_run = dummy_run or is_profile + dp_coordination_context = ( + skip_dp_coordination() if should_skip_allreduce_across_dp_group(self.vllm_config) else nullcontext() + ) + profile_kwargs = {"randomize_inputs": randomize_inputs} + forward_failed = True + with dp_coordination_context: + try: + output = super().execute_model( + scheduler_output, + intermediate_tensors=intermediate_tensors, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + is_profile=is_profile, + context_len=context_len, + valid_dummy_state_slots=valid_dummy_state_slots, + **profile_kwargs, + ) + forward_failed = False + finally: + finish_execution = getattr(self.model_state, "finish_execution", None) + if finish_execution is not None: + finish_execution(failed=forward_failed) + if not dummy_run and not is_profile: + self._maybe_send_draft_kv(scheduler_output) + self.model_state.kvpp_is_dummy_run = False + if dummy_run and lmhead_tp_enable() and not is_profile and self.is_last_pp_rank: + # lmhead TP: idle ranks never call sample(); join the target head + # here at capacity, before _dummy_run replays the dummy propose. + if self.execute_model_state is None: + raise RuntimeError( + "lmhead TP dummy join expects execute_model_state published by the upstream dummy execute_model." ) - self.model.compute_logits(self.execute_model_state.hidden_states[dummy_indices]) - self.kvpp.complete_forward() - - if dump_forward: - if dummy_run: - self._finalize_dump_data(dump=False) - elif isinstance(output, IntermediateTensors) and self.is_pooling_model: - # Non-last pooling ranks do not reach pool(). Decoding ranks - # finalize after sample_tokens(), including non-last PP ranks. - self._finalize_dump_data() - - self._cpp_execution_time_ms = _finish_profiling_chunk_timing( - profiling_config, - execution_start_time, + dummy_indices = torch.zeros( + self._lmhead_tp_max_num_logits(), + dtype=torch.int64, + device=self.device, ) + self.model.compute_logits(self.execute_model_state.hidden_states[dummy_indices]) + self.kvpp.complete_forward() + + if dump_forward: + if dummy_run: + self._finalize_dump_data(dump=False) + elif isinstance(output, IntermediateTensors) and self.is_pooling_model: + # Non-last pooling ranks do not reach pool(). Decoding ranks + # finalize after sample_tokens(), including non-last PP ranks. + self._finalize_dump_data() + + self._cpp_execution_time_ms = _finish_profiling_chunk_timing( + profiling_config, + execution_start_time, + ) - # TODO: Remove this workaround once upstream vLLM restores idle DP ranks' - # logical request counts before draft DP synchronization. - # In adaptive verification, idle DP ranks dummy-run at target graph capacity. - # Restore their request counts before draft DP sync to avoid inflating all ranks. - if ( - dummy_run - and not is_profile - and context_len == 0 - and self.adaptive_verification is not None - and self.execute_model_state is not None - ): - state = self.execute_model_state - self.execute_model_state = state._replace( - input_batch=replace(state.input_batch, num_reqs=len(scheduler_output.num_scheduled_tokens)) - ) - return output - except BaseException: - if mirror is not None: - mirror.join_pending() - raise + # TODO: Remove this workaround once upstream vLLM restores idle DP ranks' + # logical request counts before draft DP synchronization. + # In adaptive verification, idle DP ranks dummy-run at target graph capacity. + # Restore their request counts before draft DP sync to avoid inflating all ranks. + if ( + dummy_run + and not is_profile + and context_len == 0 + and self.adaptive_verification is not None + and self.execute_model_state is not None + ): + state = self.execute_model_state + self.execute_model_state = state._replace( + input_batch=replace(state.input_batch, num_reqs=len(scheduler_output.num_scheduled_tokens)) + ) + return output def _maybe_send_draft_kv(self, scheduler_output: SchedulerOutput) -> None: """Write draft KV on P, then let D pull the exact allocated cache pages.""" From 822b6c3988202fcef304086f4bf26fc08851bc82 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 10:39:28 +0800 Subject: [PATCH 094/120] refactor(hisparse): use Core grouped invalidation interface Keep the existing production pointer-table launch for both one and many groups. Remove the obsolete direct-table specialization and exercise the supported ABI in the graph regression. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 3 +- vllm_ascend/hisparse/cache_ops.py | 51 +++++++------------ vllm_ascend/ops/triton/hisparse.py | 6 +-- 3 files changed, 22 insertions(+), 38 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index adeb88c99370..2926820e3e80 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -718,6 +718,7 @@ def test_plan_gather_respects_request_state_padding(): @pytest.mark.parametrize("hot_capacity", [4, 10240]) def test_invalidate_hot_capacity_graph(hot_capacity): dgi = torch.arange(2 * hot_capacity, dtype=torch.int32, device="npu").view(2, hot_capacity) + group_table = torch.tensor([dgi.data_ptr()], dtype=torch.int64, device=dgi.device) initial = dgi.cpu() states = torch.tensor([1, 0], dtype=torch.int32, device="npu") requests = torch.tensor([0, 1, 1, -1], dtype=torch.int32, device="npu") @@ -726,7 +727,7 @@ def test_invalidate_hot_capacity_graph(hot_capacity): expected[1, -1] = expected[0, 0] = expected[0, -1] = -1 def forward(): - HiSparseNPUOps.hisparse_invalidate_written_slots(dgi, states, requests, written) + HiSparseNPUOps.hisparse_invalidate_written_slots(group_table, states, requests, written, *dgi.shape) forward() torch.npu.synchronize() diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 3211812ab0cf..4211bd0c6f1f 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -321,45 +321,31 @@ def _gather( @staticmethod def hisparse_invalidate_written_slots( - device_global_indices, + group_table, request_state_indices, req_id_per_token, written_slots, - num_state_rows=None, - region_stride=None, + num_state_rows, + region_stride, ) -> None: - device = device_global_indices.device + device = group_table.device if device.type != "npu": raise ValueError("HiSparse invalidation requires NPU state") - grouped = num_state_rows is not None or region_stride is not None - if grouped: - if num_state_rows is None or region_stride is None or num_state_rows < 0 or region_stride < 0: - raise ValueError("Grouped HiSparse invalidation requires nonnegative state and region sizes") - _check_tensor(device_global_indices, "group_table", device, (torch.int64,), ndim=1) - groups = device_global_indices.numel() - region = region_stride - state_rows = num_state_rows - else: - _check_tensor(device_global_indices, "device_global_indices", device, (torch.int32,), ndim=2) - groups = 1 - region = device_global_indices.size(1) - state_rows = device_global_indices.size(0) + if num_state_rows < 0 or region_stride < 0: + raise ValueError("HiSparse invalidation requires nonnegative state and region sizes") + _check_tensor(group_table, "group_table", device, (torch.int64,), ndim=1) + groups = group_table.numel() _check_tensor(request_state_indices, "request_state_indices", device, (torch.int32,), ndim=1) _check_tensor(written_slots, "written_slots", device, (torch.int32, torch.int64), ndim=1) _check_tensor(req_id_per_token, "req_id_per_token", device, (torch.int32,), shape=written_slots.shape) num_tokens = written_slots.numel() if num_tokens == 0: return - if region == 0 or state_rows == 0 or groups == 0: + if region_stride == 0 or num_state_rows == 0 or groups == 0: return # The 8192-row Prefill chunk is the measured large shape. Keep scratch # bounded and preserve the existing path for capture and other shapes. - if ( - grouped - and num_tokens == _SORTED_WRITTEN_SLOT_COUNT - and is_950() - and not torch.npu.is_current_stream_capturing() - ): + if num_tokens == _SORTED_WRITTEN_SLOT_COUNT and is_950() and not torch.npu.is_current_stream_capturing(): keys = torch.empty((num_tokens,), dtype=torch.int64, device=device) sorted_keys = torch.empty_like(keys) order = torch.empty_like(keys) @@ -370,31 +356,30 @@ def hisparse_invalidate_written_slots( keys, ROWS=num_tokens, REQUESTS=request_state_indices.numel(), - STATES=state_rows, + STATES=num_state_rows, TILE=256, multibuffer=False, ) torch.sort(keys, out=(sorted_keys, order)) - invalidate_written_slots_sorted[(groups, state_rows, triton.cdiv(region, 256))]( - device_global_indices, + invalidate_written_slots_sorted[(groups, num_state_rows, triton.cdiv(region_stride, 256))]( + group_table, sorted_keys, ROWS=num_tokens, - HOT=region, + HOT=region_stride, TILE=256, STEPS=num_tokens.bit_length(), multibuffer=False, ) return - invalidate_kernel[(num_tokens, triton.cdiv(region, 256), groups)]( - device_global_indices, + invalidate_kernel[(num_tokens, triton.cdiv(region_stride, 256), groups)]( + group_table, request_state_indices, req_id_per_token, written_slots, REQUESTS=request_state_indices.numel(), - STATE_ROWS=state_rows, - REGION=region, + STATE_ROWS=num_state_rows, + REGION=region_stride, BLOCK=256, - GROUPED=grouped, multibuffer=False, ) diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 4752d12b33cd..53b5aa24c003 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -80,7 +80,7 @@ def gather_rows_kernel( @triton.jit def invalidate_kernel( - dgi, + tables, states, requests, written, @@ -88,10 +88,8 @@ def invalidate_kernel( STATE_ROWS: tl.constexpr, REGION: tl.constexpr, BLOCK: tl.constexpr, - GROUPED: tl.constexpr = False, ): - if GROUPED: - dgi = tl.load(dgi + tl.program_id(2)).to(tl.pointer_type(tl.int32)) + dgi = tl.load(tables + tl.program_id(2)).to(tl.pointer_type(tl.int32)) token = tl.program_id(0) offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) request = tl.load(requests + token) From 3f48a579357185bca8b91cfe9ab56b82937b45d8 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 11:14:55 +0800 Subject: [PATCH 095/120] refactor(hisparse): install row adapters without source hashes Use the existing binding and transfer lifecycle directly. Remove whole-file source allowlists and the optional recycled-adapter installation path. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/row_dma.py | 40 ++++++++++----------------------- 1 file changed, 12 insertions(+), 28 deletions(-) diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index 265de6c4bbd8..4f74223c96ac 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -4,17 +4,10 @@ from dataclasses import dataclass from functools import wraps -from hashlib import sha256 -from pathlib import Path import numpy as np import torch -_WORKER_SHA256 = "864e2c7ce4e2cff320b4f68f8dccc9691adb829f0e50fdf6a27dc9c7885040e5" -_RECYCLED_RUNTIME_SHA256S = { - "16e33814eb5d6b44093429fbb74aaf436e8994b11b43b07d1c39851baf563759", # Core 24420b7a - "935ad5669ccdcbc0ffa45d66cfd28145cee0039fe979bb1c637f325373133325", # Core 39dbf518 -} # Large prefill batches keep the scalar path to bound temporary memory. _MAX_DESCRIPTORS = 32768 @@ -111,7 +104,7 @@ def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: return True -def _install_worker(worker_type, backend_type, bind_recycled=None) -> None: +def _install_worker(worker_type, backend_type, bind_recycled) -> None: original_init = worker_type._init_dma original_enqueue = worker_type._enqueue_row_dma if getattr(original_enqueue, "_ascend_row_dma", False): @@ -133,8 +126,7 @@ def init(worker): staging.slots.shape, dtype=slots.dtype, device="cpu", pin_memory=True ) worker._ascend_row_dma_geometry = _Geometry.bind(worker) - if bind_recycled is not None: - worker._ascend_recycled_invalidator = bind_recycled(worker) + worker._ascend_recycled_invalidator = bind_recycled(worker) @wraps(original_enqueue) def enqueue(worker, layer_indices, ready_event=None): @@ -145,29 +137,21 @@ def enqueue(worker, layer_indices, ready_event=None): enqueue._ascend_row_dma = True worker_type._init_dma = init worker_type._enqueue_row_dma = enqueue - if bind_recycled is not None: - # Keep the same pinned adapter until Core exposes this backend hook. - original_invalidate = worker_type.invalidate_blocks + original_invalidate = worker_type.invalidate_blocks - @wraps(original_invalidate) - def invalidate(worker, block_ids, request_state_indices): - owner = getattr(worker, "_ascend_recycled_invalidator", None) - if owner is None or not owner.try_call(worker, block_ids, request_state_indices): - return original_invalidate(worker, block_ids, request_state_indices) + @wraps(original_invalidate) + def invalidate(worker, block_ids, request_state_indices): + owner = getattr(worker, "_ascend_recycled_invalidator", None) + if owner is None or not owner.try_call(worker, block_ids, request_state_indices): + return original_invalidate(worker, block_ids, request_state_indices) - worker_type.invalidate_blocks = invalidate + worker_type.invalidate_blocks = invalidate def install(backend_type) -> None: - """Only adapt the exact Core worker implementation reviewed for this backend.""" + """Install the Ascend row-mirror and recycled-slot adapters.""" from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker - from vllm.v1.hisparse import runtime - if sha256(Path(worker.__file__).read_bytes()).hexdigest() != _WORKER_SHA256: - return - bind_recycled = None - if sha256(Path(runtime.__file__).read_bytes()).hexdigest() in _RECYCLED_RUNTIME_SHA256S: - from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator + from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator - bind_recycled = RecycledGroupInvalidator.bind - _install_worker(worker.HiSparseConnectorWorker, backend_type, bind_recycled) + _install_worker(worker.HiSparseConnectorWorker, backend_type, RecycledGroupInvalidator.bind) From 383ba6b6b6cee21b9ad10db5fa01f14db91b2869 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 11:25:24 +0800 Subject: [PATCH 096/120] refactor(hisparse): collapse the native resolver forwarding helper Call the native resolver directly from the existing owner-aware dispatch entry. Preserve all native arguments and runtime-owned scratch views. Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/resolver.py | 73 +++++++++++++------------------- 1 file changed, 30 insertions(+), 43 deletions(-) diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py index 2e6466ce324f..6a49ad49f32a 100644 --- a/vllm_ascend/hisparse/resolver.py +++ b/vllm_ascend/hisparse/resolver.py @@ -34,56 +34,43 @@ def _load_resolver(): return namespace -def resolve_tensors(native, arguments, resolved_scratch, errors, *, attention_capacity=0, attention_block_stride=None): - """Pass live tensors; C++ derives geometry and protects queued storage. +def resolve(owner, arguments): + """Dispatch live tensors to the native request-union resolver. During replay, external tensors remain owned by the runtime/workspace. Callers must destroy graphs before releasing or rebinding those owners. """ - a = arguments - native.hisparse_resolve_residency( - a["global_indices"], - a["request_ids"], - a["request_state_indices"], - a["source_block_table"], - a["resident_block_table"], - a["hot_block_table"], - a["device_global_indices"], - a["lru_slots"], - a["hot_indices"], - a["attention_indices"], - resolved_scratch, - a["valid_counts"], - a["swap_host_physical_rows"], - a["swap_device_physical_rows"], - a["swap_counts"], - a["miss_mask"], - a["stats"], - errors, - a["num_valid_rows"], - a["host_cache"].shape[0], - a["hot_cache"].shape[1], - a["source_block_size"], - a["resident_block_size"], - a["resident_null_block"], - a["attention_block_stride"] if attention_block_stride is None else attention_block_stride, - a["max_union_rows"], - attention_capacity, - ) - - -def resolve(owner, scope): - """Resolve with the native request-union kernel; no Triton fallback.""" - rows, topk = scope["global_indices"].shape + rows, topk = arguments["global_indices"].shape workspace = owner.workspace - resolved = scope["resolved_global_indices"] + resolved = arguments["resolved_global_indices"] if resolved is None: resolved = workspace.first_refs[: rows * topk].view(rows, topk) - resolve_tensors( - workspace.native, - scope, + workspace.native.hisparse_resolve_residency( + arguments["global_indices"], + arguments["request_ids"], + arguments["request_state_indices"], + arguments["source_block_table"], + arguments["resident_block_table"], + arguments["hot_block_table"], + arguments["device_global_indices"], + arguments["lru_slots"], + arguments["hot_indices"], + arguments["attention_indices"], resolved, + arguments["valid_counts"], + arguments["swap_host_physical_rows"], + arguments["swap_device_physical_rows"], + arguments["swap_counts"], + arguments["miss_mask"], + arguments["stats"], workspace.errors[:rows], - attention_capacity=owner.attention_capacity, - attention_block_stride=owner.attention_block_stride, + arguments["num_valid_rows"], + arguments["host_cache"].shape[0], + arguments["hot_cache"].shape[1], + arguments["source_block_size"], + arguments["resident_block_size"], + arguments["resident_null_block"], + arguments["attention_block_stride"] if owner.attention_block_stride is None else owner.attention_block_stride, + arguments["max_union_rows"], + owner.attention_capacity, ) From b22dbf6a0c6f9c874ce8b1fd80ad15ef405788ce Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 11:46:07 +0800 Subject: [PATCH 097/120] refactor(hisparse): use Core copy streams directly Remove the capture-only stream wrapper after validating the fused resolver and refill with generic torch.Stream on Ascend950. Both wrapped and generic cross-stream cases pass, and the removal passes six resolver/refill and workspace rebind Graph cases. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 4 +-- .../one_card/test_hisparse_prefill_staging.py | 3 +- vllm_ascend/hisparse/backend.py | 36 ------------------- 3 files changed, 3 insertions(+), 40 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 2926820e3e80..141ce79af5ff 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -6,7 +6,7 @@ import torch import torch_npu # noqa: F401 -from vllm_ascend.hisparse.backend import AscendHiSparseBackend, capture_aware_stream +from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, ResolverWorkspace, @@ -393,7 +393,7 @@ def test_topk_2048_speculative_union_graph(query_rows, hot_size, cross_stream): sh, sd, sc = torch.empty_like(out), torch.empty_like(out), torch.empty_like(counts) stats = torch.zeros(2, dtype=torch.int64, device="npu") - stream = capture_aware_stream(torch.Stream(device="npu")) if cross_stream else None + stream = torch.Stream(device="npu") if cross_stream else None ready = torch.Event() if cross_stream else None def resolve_and_gather(): diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index ac0bcc809743..57faea4a6cff 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -208,7 +208,6 @@ def resolve(): def test_profiled_workspace_rebind_preserves_prior_length_graphs(pages, stride_rows): from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager - from vllm_ascend.hisparse.backend import capture_aware_stream from vllm_ascend.hisparse.backend import prepare_hisparse_workspaces as actual_prepare prepare_hisparse_workspaces = torch.inference_mode()(actual_prepare) @@ -227,7 +226,7 @@ def test_profiled_workspace_rebind_preserves_prior_length_graphs(pages, stride_r def runtime(): group = SimpleNamespace( - copy_stream=capture_aware_stream(torch.Stream(device=device)), + copy_stream=torch.Stream(device=device), shared_topk=SimpleNamespace(device_topk_rows=torch.empty((4, 64), dtype=torch.int32, device=device)), ) owner = SimpleNamespace( diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index f48e783fb37e..0ff602703f7f 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -66,39 +66,6 @@ def prepare_hisparse_workspaces(forward_context: dict[str, Any]) -> None: runtime._ascend_profiled_cache_ops = workspace, _allocate_attention_lengths(runtime) -class _CaptureAwareStream(torch.Stream): - """Propagate a captured dependency before PTA allocates CANN workspace. - - PTA's native NPUEvent wait drains Host submission queues during capture; - its generic event binding does not. Keep generic event points and stream - identity, adding the missing submission boundary only during capture. - Remove when the supported PTA generic binding has the same behavior. - """ - - def _finish_capture_propagation(self) -> None: - if torch.npu.is_current_stream_capturing(): - native = torch.npu.Stream( - stream_id=self.stream_id, device_index=self.device_index, device_type=self.device_type - ) - # Obtaining the native handle drains its Host task queue. This is - # not a device synchronization and does not run during graph replay. - _ = native.npu_stream - - def wait_event(self, event) -> None: - super().wait_event(event) - self._finish_capture_propagation() - - def wait_stream(self, stream) -> None: - super().wait_stream(stream) - self._finish_capture_propagation() - - -def capture_aware_stream(stream: torch.Stream) -> torch.Stream: - if stream.device.type != "npu" or isinstance(stream, _CaptureAwareStream): - return stream - return _CaptureAwareStream(stream.stream_id, stream.device_index, stream.device_type) - - class AscendHiSparseBackend(HiSparseBackend): # Keep private pools until NPU IPC events interoperate with Core's generic streams. stats_dtype = torch.int64 @@ -155,9 +122,6 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: from vllm_ascend.hisparse.cache_ops import HiSparseNPUOps group = runtime.index_group - # Runtime copied the model group's stream reference at layer creation. - # Adapt the live runtime binding before any graph captures use it. - group.copy_stream = capture_aware_stream(group.copy_stream) prepared = getattr(runtime, "_ascend_profiled_cache_ops", None) if runtime is group.leader: workspace = prepared[0] if prepared is not None else _allocate_resolver_workspace(runtime) From f1007cd14ac0636621c64e3af6596ce76f93dd96 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 16:27:24 +0800 Subject: [PATCH 098/120] perf(hisparse): generalize sorted slot invalidation Choose sorted invalidation from the table and write geometry, with a bounded sort buffer. Use the same selection during graph capture and verify changing writes on both single-stream and cross-stream replay. Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_ops.py | 71 ++++++++++++++----- vllm_ascend/hisparse/cache_ops.py | 35 ++++++--- vllm_ascend/ops/triton/hisparse.py | 12 ++-- 3 files changed, 83 insertions(+), 35 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 141ce79af5ff..cd432ff70b9a 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -751,46 +751,81 @@ def forward(): graph.reset() -def test_common_group_invalidation_changed_graph_replay(): - """Exercise the current common pointer-table caller with changing writes.""" +@pytest.mark.parametrize( + "rows,hot_capacity,cross_stream,expect_sorted", + [(4, 16, False, False), (1537, 8193, False, True), (1537, 8193, True, True)], +) +def test_common_group_invalidation_changed_graph_replay(rows, hot_capacity, cross_stream, expect_sorted, monkeypatch): + """Exercise both invalidation routes with changing writes in eager and Graph.""" from types import SimpleNamespace from vllm.v1.hisparse.runtime import HiSparseSlotInvalidator groups = 3 - initial = torch.arange(48, dtype=torch.int32).view(3, 16) + initial = torch.arange(3 * hot_capacity, dtype=torch.int32).view(3, hot_capacity) tables = [(initial + group * 7).npu() for group in range(groups)] states = torch.tensor([2, 0, 1], dtype=torch.int32, device="npu") - requests = torch.tensor([0, 1, 2, -1], dtype=torch.int32, device="npu") - written = torch.tensor([35, 4, 19, 1], dtype=torch.int64, device="npu") + requests = torch.full((rows,), -1, dtype=torch.int32, device="npu") + requests[:3] = torch.arange(3, dtype=requests.dtype, device=requests.device) + written = torch.zeros(rows, dtype=torch.int64, device="npu") leaders = [ SimpleNamespace(index_group=SimpleNamespace(device_global_indices=table), cache_ops=HiSparseNPUOps) for table in tables ] invalidator = HiSparseSlotInvalidator(states) + consumer = torch.Stream(device="npu") if cross_stream else None + ready = torch.Event() if cross_stream else None def forward(): - invalidator(leaders, written, requests) + if consumer is None: + invalidator(leaders, written, requests) + return + current = torch.accelerator.current_stream() + consumer.wait_stream(current) + with consumer: + invalidator(leaders, written, requests) + ready.record(consumer) + current.wait_event(ready) - forward() - torch.npu.synchronize() - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - graph = torch.npu.NPUGraph() - try: - with torch.npu.graph(graph, stream=stream): - forward() - for values in [[35, 4, 19, 1], [47, 8, 23, 1], [-1, -1, -1, 1]]: - written.copy_(torch.tensor(values, dtype=torch.int64)) + def check_writes(run): + for values in ( + [2 * hot_capacity + 3, 4, hot_capacity + 3], + [3 * hot_capacity - 1, 8, hot_capacity + 7], + [-1] * 3, + ): + written[:3].copy_(torch.tensor(values, dtype=written.dtype)) for group, table in enumerate(tables): table.copy_(initial + group * 7) - graph.replay() + run() torch.npu.synchronize() for group, table in enumerate(tables): expected = initial + group * 7 - for row, value in zip([2, 0, 1], values[:3]): + for row, value in zip([2, 0, 1], values): expected[row][expected[row] == value] = -1 torch.testing.assert_close(table.cpu(), expected) + + check_writes(forward) + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + torch.npu.synchronize() + + sort = torch.sort + sort_calls = 0 + + def count_sort(*args, **kwargs): + nonlocal sort_calls + sort_calls += 1 + return sort(*args, **kwargs) + + monkeypatch.setattr(torch, "sort", count_sort) + graph = torch.npu.NPUGraph() + try: + with torch.npu.graph(graph, stream=stream): + forward() + assert sort_calls == int(expect_sorted) + check_writes(graph.replay) finally: torch.npu.synchronize() graph.reset() diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 4211bd0c6f1f..fba3948927a5 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -25,7 +25,8 @@ scatter_packed_kv_kernel, ) -_SORTED_WRITTEN_SLOT_COUNT = 8192 +_SORT_MIN_DIRECT_TILES = 64 * 1024 +_SORT_BUFFER_LIMIT_BYTES = 768 * 1024 class ResolverWorkspace: @@ -343,13 +344,25 @@ def hisparse_invalidate_written_slots( return if region_stride == 0 or num_state_rows == 0 or groups == 0: return - # The 8192-row Prefill chunk is the measured large shape. Keep scratch - # bounded and preserve the existing path for capture and other shapes. - if num_tokens == _SORTED_WRITTEN_SLOT_COUNT and is_950() and not torch.npu.is_current_stream_capturing(): + tile = 256 + hot_tiles = triton.cdiv(region_stride, tile) + table_tiles = groups * hot_tiles + search_steps = num_tokens.bit_length() + # Direct invalidation scans the Hot tiles for each write. Sorting + # instead searches once per state and Hot tile. Keep a margin for + # sorting and cap the three int64 buffers before launching it. + direct_tiles = num_tokens * table_tiles + search_probes = num_state_rows * table_tiles * search_steps + sort_buffer_bytes = 3 * num_tokens * torch.int64.itemsize + if ( + direct_tiles >= 2 * search_probes + _SORT_MIN_DIRECT_TILES + and sort_buffer_bytes <= _SORT_BUFFER_LIMIT_BYTES + and is_950() + ): keys = torch.empty((num_tokens,), dtype=torch.int64, device=device) sorted_keys = torch.empty_like(keys) order = torch.empty_like(keys) - pack_written_slot_keys[(triton.cdiv(num_tokens, 256),)]( + pack_written_slot_keys[(triton.cdiv(num_tokens, tile),)]( request_state_indices, req_id_per_token, written_slots, @@ -357,21 +370,21 @@ def hisparse_invalidate_written_slots( ROWS=num_tokens, REQUESTS=request_state_indices.numel(), STATES=num_state_rows, - TILE=256, + TILE=tile, multibuffer=False, ) torch.sort(keys, out=(sorted_keys, order)) - invalidate_written_slots_sorted[(groups, num_state_rows, triton.cdiv(region_stride, 256))]( + invalidate_written_slots_sorted[(groups, num_state_rows, hot_tiles)]( group_table, sorted_keys, ROWS=num_tokens, HOT=region_stride, - TILE=256, - STEPS=num_tokens.bit_length(), + TILE=tile, + STEPS=search_steps, multibuffer=False, ) return - invalidate_kernel[(num_tokens, triton.cdiv(region_stride, 256), groups)]( + invalidate_kernel[(num_tokens, hot_tiles, groups)]( group_table, request_state_indices, req_id_per_token, @@ -379,7 +392,7 @@ def hisparse_invalidate_written_slots( REQUESTS=request_state_indices.numel(), STATE_ROWS=num_state_rows, REGION=region_stride, - BLOCK=256, + BLOCK=tile, multibuffer=False, ) diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index 53b5aa24c003..f0c4fb3f4e66 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -101,15 +101,15 @@ def invalidate_kernel( tl.store(ptr, -1, mask=valid & (cached == value)) -@triton.jit +@triton.jit(do_not_specialize=["ROWS", "REQUESTS", "STATES"]) def pack_written_slot_keys( states, requests, written, keys, - ROWS: tl.constexpr, - REQUESTS: tl.constexpr, - STATES: tl.constexpr, + ROWS, + REQUESTS, + STATES, TILE: tl.constexpr, ): row = tl.program_id(0) * TILE + tl.arange(0, TILE) @@ -124,11 +124,11 @@ def pack_written_slot_keys( tl.store(keys + row, tl.where(valid, key, 9223372036854775807), mask=row < ROWS) -@triton.jit +@triton.jit(do_not_specialize=["ROWS"]) def invalidate_written_slots_sorted( tables, keys, - ROWS: tl.constexpr, + ROWS, HOT: tl.constexpr, TILE: tl.constexpr, STEPS: tl.constexpr, From ec9338ed2a1053a937820510b1b378b80e2240f3 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 20:24:44 +0800 Subject: [PATCH 099/120] refactor(hisparse): share native resolver UB budget checks Signed-off-by: zhenwenqi2024 --- csrc/hisparse/hisparse_resolver_torch_adpt.h | 33 +++++++++++-------- csrc/torch_binding.cpp | 6 ++-- .../one_card/test_hisparse_ops.py | 10 ++++++ .../test_hisparse_foundation_cpu.py | 31 +++-------------- vllm_ascend/hisparse/cache_ops.py | 4 +-- vllm_ascend/hisparse/resolver.py | 23 +------------ 6 files changed, 40 insertions(+), 67 deletions(-) diff --git a/csrc/hisparse/hisparse_resolver_torch_adpt.h b/csrc/hisparse/hisparse_resolver_torch_adpt.h index 0177fd280aed..b8d29877518c 100644 --- a/csrc/hisparse/hisparse_resolver_torch_adpt.h +++ b/csrc/hisparse/hisparse_resolver_torch_adpt.h @@ -22,6 +22,23 @@ at::Tensor tensor_or_undefined(const std::optional& tensor) { return tensor.value_or(at::Tensor()); } +inline int64_t resolver_hash_size(int64_t hot, int64_t union_rows) { + TORCH_CHECK(union_rows > 0 && union_rows <= hot && hot <= 32768, + "HiSparse union must fit the hot cache and int16 LRU slots"); + int64_t hash_size = 1; + while (hash_size <= union_rows) hash_size <<= 1; + const uint64_t tiles = (hot + 31) / 32; + const uint64_t budget = hisparse_resolver_scratch_capacity(); + const uint64_t fixed_words = hot + 2 * tiles + 64; + const uint64_t minimum_bytes = (2 * (union_rows + 1) + fixed_words + 15) / 16 * 64 + 512; + TORCH_CHECK(minimum_bytes <= budget, + "HiSparse hot_size=", hot, ", max_union_rows=", union_rows, + " needs at least ", minimum_bytes, " UB bytes; the compiled target allows ", budget, + ". Lower device_buffer_size or speculative tokens."); + const uint64_t available_words = (budget - 512) / 64 * 16; + return std::min(hash_size, int64_t((available_words - fixed_words) / 2)); +} + void resolve_residency( const at::Tensor& logical, const std::optional& requests, @@ -50,8 +67,7 @@ void resolve_residency( const auto rows = logical.size(0), topk = logical.size(1), hot = dgi.size(1); TORCH_CHECK(rows > 0 && topk > 0 && rows <= ((int64_t(1) << 30) - 1) / topk, "HiSparse first-reference positions exceed the int32 tag representation"); - TORCH_CHECK(union_rows > 0 && union_rows <= hot && hot <= 32768 && topk <= hot, - "HiSparse union must fit the hot cache and int16 LRU slots"); + TORCH_CHECK(topk <= hot, "HiSparse TopK must fit the hot cache"); TORCH_CHECK(host_rows > 0 && host_rows < (int64_t(1) << 31), "Invalid HiSparse host row capacity"); TORCH_CHECK(hot_block_size > 0 && attention_block_stride >= hot_block_size, "Invalid hot/attention block strides"); @@ -67,18 +83,7 @@ void resolve_residency( TORCH_CHECK(!source || source->dim() == 2, "HiSparse source table must be a matrix"); TORCH_CHECK(!resident || resident->dim() == 2, "HiSparse resident table must be a matrix"); TORCH_CHECK(lru.sizes() == dgi.sizes(), "HiSparse LRU/tag geometry mismatch"); - int64_t hash_size = 1; - while (hash_size <= union_rows) hash_size <<= 1; - const uint64_t tiles = (hot + 31) / 32; - const uint64_t budget = hisparse_resolver_scratch_capacity(); - const uint64_t available_words = (budget - 512) / 64 * 16; - const uint64_t fixed_words = hot + 2 * tiles + 64; - const uint64_t minimum_bytes = (2 * (union_rows + 1) + fixed_words + 15) / 16 * 64 + 512; - TORCH_CHECK(minimum_bytes <= budget, - "HiSparse hot_size=", hot, ", max_union_rows=", union_rows, - " needs at least ", minimum_bytes, " UB bytes; the compiled target allows ", budget, - ". Lower device_buffer_size or speculative tokens."); - hash_size = std::min(hash_size, int64_t((available_words - fixed_words) / 2)); + const int64_t hash_size = resolver_hash_size(hot, union_rows); const std::array tensors = { logical, tensor_or_undefined(requests), tensor_or_undefined(states), diff --git a/csrc/torch_binding.cpp b/csrc/torch_binding.cpp index 47a3143656cf..a32ed500093c 100644 --- a/csrc/torch_binding.cpp +++ b/csrc/torch_binding.cpp @@ -2854,8 +2854,10 @@ TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) "int attention_block_stride, int union_rows, int attention_capacity=0) -> ()"); ops.impl("hisparse_resolve_residency", torch::kPrivateUse1, &vllm_ascend::hisparse::resolve_residency); ops.impl("hisparse_resolve_residency", torch::kMeta, &vllm_ascend::hisparse::resolve_residency_meta); - ops.def("hisparse_resolver_scratch_capacity() -> int", - []() -> int64_t { return hisparse_resolver_scratch_capacity(); }); + ops.def("hisparse_resolver_validate_scratch(int hot, int union_rows) -> ()", + [](int64_t hot, int64_t union_rows) { + vllm_ascend::hisparse::resolver_hash_size(hot, union_rows); + }); #endif ops.def("get_npu_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor?"); ops.impl("get_npu_view_from_cpu_tensor", c10::DispatchKey::CompositeExplicitAutograd, diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index cd432ff70b9a..babfe0ae4dfc 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -26,6 +26,16 @@ def _require_mapped_host_device(): AscendHiSparseBackend().check_available() +def test_native_resolver_validates_ub_before_allocating_workspace(): + # This is the same validator called by workspace construction and launch. + validate = torch.ops._C_ascend.hisparse_resolver_validate_scratch + validate(10240, 2048) + with pytest.raises(RuntimeError, match="UB bytes"): + validate(32768, 32768) + with pytest.raises(RuntimeError, match="int16 LRU slots"): + validate(16, 17) + + @pytest.mark.parametrize("descriptor_count", [0, 1, 16, 17, 312]) @pytest.mark.parametrize("graph_producer", [False, True]) def test_private_host_pool_retains_pinned_storage_for_dma_writeback(descriptor_count, graph_producer): diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 16014f58e1e2..6e6b31a40118 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -178,7 +178,7 @@ def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups # Loading a device extension is the only mocked boundary; the group-owned # tensors below are real CPU storage, not a replacement resolver algorithm. - native = NS(hisparse_resolver_scratch_capacity=lambda: 175616) + native = NS(hisparse_resolver_validate_scratch=lambda hot, union: None) monkeypatch.setattr(resolver, "_load_resolver", lambda: native) backend = AscendHiSparseBackend() @@ -233,31 +233,6 @@ def test_resolver_scratch_rejects_topk_outside_hot_capacity(topk): hisparse_ops.ResolverWorkspace(4, 16, 8, torch.device("cpu"), max_topk=topk) -def test_native_resolver_uses_available_ub_without_truncating_the_union(): - from vllm_ascend.hisparse.resolver import scratch_layout - - hot, union, capacity = 12288, 8192, 175616 - buckets, size = scratch_layout(hot, union, capacity) - assert buckets > union - assert 0 < size <= capacity and size % 64 == 0 - # Find the actual supported boundary without pinning a particular hash - # bucket count or the precise UB layout/error-message arithmetic. - too_small, fits = 0, capacity - while fits - too_small > 1: - candidate = (too_small + fits) // 2 - try: - scratch_layout(hot, union, candidate) - except RuntimeError: - too_small = candidate - else: - fits = candidate - minimum_buckets, minimum_size = scratch_layout(hot, union, fits) - assert minimum_buckets > union and minimum_size <= fits - assert minimum_size % 64 == 0 - with pytest.raises(RuntimeError, match="UB bytes"): - scratch_layout(hot, union, fits - 1) - - def test_missing_native_resolver_raises_instead_of_falling_back(monkeypatch): from vllm_ascend.hisparse import resolver @@ -392,7 +367,9 @@ def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind( ): from vllm_ascend.hisparse import resolver - monkeypatch.setattr(resolver, "_load_resolver", lambda: NS(hisparse_resolver_scratch_capacity=lambda: 175616)) + monkeypatch.setattr( + resolver, "_load_resolver", lambda: NS(hisparse_resolver_validate_scratch=lambda hot, union: None) + ) backend = AscendHiSparseBackend() group = NS(copy_stream=NS(device=torch.device("cpu"))) runtime = NS( diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index fba3948927a5..4bb08274a58c 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -43,10 +43,10 @@ def __init__( self.max_topk = hot_size if max_topk is None else max_topk if not 0 < self.max_topk <= hot_size: raise ValueError("HiSparse resolver TopK capacity must be positive and fit the Hot cache") - from vllm_ascend.hisparse.resolver import _load_resolver, scratch_layout + from vllm_ascend.hisparse.resolver import _load_resolver self.native = _load_resolver() - scratch_layout(hot_size, max_union_rows, self.native.hisparse_resolver_scratch_capacity()) + self.native.hisparse_resolver_validate_scratch(hot_size, max_union_rows) # One resolved Host row per selected token, not per Hot-cache slot. self.first_refs = torch.empty(max_rows * self.max_topk, dtype=torch.int32, device=device) self.errors = torch.zeros((max_rows, 8), dtype=torch.int32, device=device) diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py index 6a49ad49f32a..6f268a1ba5de 100644 --- a/vllm_ascend/hisparse/resolver.py +++ b/vllm_ascend/hisparse/resolver.py @@ -4,31 +4,10 @@ import torch -def scratch_layout(hot, max_union_rows, capacity=None): - """Size hash/LRU scratch from the request capacity, not a model shape.""" - if not 0 < max_union_rows <= hot <= 32768: - raise ValueError("HiSparse union must fit the hot cache and int16 LRU slot representation") - hash_size = 1 << max_union_rows.bit_length() - tiles = (hot + 31) // 32 - fixed_words = hot + 2 * tiles + 64 - if capacity is not None: - available_words = max(0, (capacity - 512) // 64) * 16 - max_hash_size = (available_words - fixed_words) // 2 - if max_hash_size <= max_union_rows: - minimum_bytes = ((2 * (max_union_rows + 1) + fixed_words + 15) // 16) * 64 + 512 - raise RuntimeError( - f"HiSparse hot_size={hot}, max_union_rows={max_union_rows} needs at least {minimum_bytes} UB bytes; " - f"the compiled target allows {capacity}. Lower device_buffer_size or speculative tokens." - ) - hash_size = min(hash_size, max_hash_size) - scratch_ints = ((2 * hash_size + hot + 2 * tiles + 64 + 15) // 16) * 16 - return hash_size, scratch_ints * 4 + 512 - - def _load_resolver(): namespace = torch.ops._C_ascend if not all( - hasattr(namespace, name) for name in ("hisparse_resolver_scratch_capacity", "hisparse_resolve_residency") + hasattr(namespace, name) for name in ("hisparse_resolver_validate_scratch", "hisparse_resolve_residency") ): raise RuntimeError("HiSparse requires its compiled AscendC resolver. Rebuild vllm-ascend for this target.") return namespace From 068ed15cc0f32681cab2d00eb88041dd1bd19076 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 20:21:33 +0800 Subject: [PATCH 100/120] refactor(hisparse): reuse SFA KV preprocessing Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_kv_preprocess.py | 26 +++++++++--------- vllm_ascend/attention/sfa_hisparse.py | 27 ++----------------- 2 files changed, 14 insertions(+), 39 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py index 1e6213b1a9a3..4679a72b7005 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py @@ -42,26 +42,24 @@ def test_preprocess_writes_hisparse_mapping_once(cache_dtype, capture): torch.empty(1, block, 1, 64, dtype=values.dtype, device="npu"), torch.arange(block, dtype=torch.int64, device="npu"), ) - impl = SimpleNamespace( - enable_sparse_sfa_c8=c8, - c8_cache_dtype=torch.float8_e4m3fn, - sfa_qsfa_tile_size=128, - kv_lora_rank=512, - qk_rope_head_dim=64, - num_kv_heads=1, - _kv_staging_capacity=block, - _kv_workspace=lambda dtype: scratch, - kv_a_layernorm=SimpleNamespace(weight=gamma, variance_epsilon=1e-5), - hisparse_cache=SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)), - ) + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl.enable_sparse_sfa_c8 = c8 + impl.enable_sparse_sfa_turboquant = False + impl.c8_cache_dtype = torch.float8_e4m3fn + impl.sfa_qsfa_tile_size = 128 + impl.kv_lora_rank, impl.qk_rope_head_dim, impl.num_kv_heads = 512, 64, 1 + impl._kv_staging_capacity = block + impl._kv_workspace = lambda dtype: scratch + impl.kv_a_layernorm = SimpleNamespace(weight=gamma, variance_epsilon=1e-5) + impl.hisparse_cache = SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)) # The ordinary source cache has no Resident/Hot addressing. It must never # receive the parent's second C8 store after exec_kv has used write_target. source_cache = (torch.empty(0, device="npu"),) def forward(): - rope, latent, *scales = AscendHiSparseImpl.exec_kv(impl, values, cos, sin, source_cache, slots, metadata) + rope, latent, *scales = impl.exec_kv(values, cos, sin, source_cache, slots, metadata) scales = scales[0] if scales else None - AscendHiSparseImpl._store_parallel_kv(impl, rope, latent, scales, source_cache, slots, metadata) + impl._store_parallel_kv(rope, latent, scales, source_cache, slots, metadata) graph = None try: diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 0a899cb65245..088da8019233 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -6,7 +6,6 @@ from dataclasses import dataclass import torch -import torch_npu from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import logger @@ -276,21 +275,9 @@ def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): assert self.kv_a_layernorm is not None if self.enable_sparse_sfa_c8: - from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope from vllm_ascend.hisparse.c8 import pack_c8_components - values = kv_no_split.reshape(-1, self.num_kv_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim) - rope, latent, scales = custom_kv_rmsnorm_rope( - values, - self.kv_a_layernorm.weight, - cos, - sin, - self.kv_lora_rank, - self.qk_rope_head_dim, - epsilon=self.kv_a_layernorm.variance_epsilon, - dst_type=self.c8_cache_dtype, - tile_size=self.sfa_qsfa_tile_size, - ) + rope, latent, scales = super().exec_kv(kv_no_split, cos, sin, kv_cache, slots, attn_metadata) rows = latent.numel() // self.kv_lora_rank packed = pack_c8_components(latent.contiguous(), rope.contiguous(), scales.contiguous()) cache = self.hisparse_cache @@ -312,17 +299,7 @@ def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): # Match ordinary SFA's fused RMSNorm/RoPE rounding. Separate RMSNorm # can differ by a BF16 ULP near a rounding midpoint. Local slots # isolate padding from the real resident/hot destination below. - torch_npu.npu_kv_rmsnorm_rope_cache( - values.view(num_tokens, self.num_kv_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim), - self.kv_a_layernorm.weight, - cos, - sin, - local_slots[:num_tokens], - rope_cache, - latent_cache, - epsilon=self.kv_a_layernorm.variance_epsilon, - cache_mode="PA", - ) + super().exec_kv(values, cos, sin, (latent_cache, rope_cache), local_slots[:num_tokens], attn_metadata) latent = latent_cache.view(-1, self.kv_lora_rank)[:num_tokens] rope = rope_cache.view(-1, self.qk_rope_head_dim)[:num_tokens] else: From 78ab28526bdc7568f16cf5170731ba84927c0415 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 20:27:09 +0800 Subject: [PATCH 101/120] refactor(hisparse): keep RoPE sharing out of the foundation Signed-off-by: zhenwenqi2024 --- tests/ut/worker/test_attn_utils_v2.py | 133 ------------------ .../attention/context_parallel/sfa_cp.py | 7 +- vllm_ascend/attention/indexer.py | 9 +- vllm_ascend/attention/sfa_hisparse.py | 5 +- vllm_ascend/attention/sfa_v1.py | 15 +- vllm_ascend/ops/rotary_embedding.py | 53 ------- vllm_ascend/worker/v2/attn_utils.py | 22 --- vllm_ascend/worker/v2/model_states/default.py | 6 +- 8 files changed, 11 insertions(+), 239 deletions(-) diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index debd5326ee31..3e234cdde5db 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -9,7 +9,6 @@ from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache -from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, @@ -69,138 +68,6 @@ from vllm_ascend.worker.v2.utils import AscendV2KVBlockZeroer -@pytest.mark.parametrize("full_graph,for_capture", [(False, False), (True, False), (True, True)]) -@pytest.mark.parametrize("boundary", ["none", "builder", "callback", "provider"]) -@pytest.mark.parametrize("inference", [False, True]) -def test_plain_sfa_and_indexer_share_rope_only_within_one_compatible_build( - monkeypatch, full_graph, for_capture, boundary, inference -): - with torch.inference_mode(inference): - from vllm_ascend.attention.indexer import AscendSFAIndexerMetadataBuilder - from vllm_ascend.ops import rotary_embedding as rotary - - table = torch.arange(128, dtype=torch.float32).reshape(32, 4) - monkeypatch.setattr(rotary, "_cos_cache", table) - monkeypatch.setattr(rotary, "_sin_cache", -table) - monkeypatch.setattr(rotary, "_cos_sin_flat", None) - monkeypatch.setattr(rotary, "_cos_mla", torch.empty(3, 1, 1, 4)) - monkeypatch.setattr(rotary, "_sin_mla", torch.empty(3, 1, 1, 4)) - lookup = rotary.get_cos_and_sin_mla - lookups = [] - - def counted_lookup(positions, use_cache=False): - lookups.append(positions.clone()) - return lookup(positions, use_cache=use_cache) - - monkeypatch.setattr(rotary, "get_cos_and_sin_mla", counted_lookup) - - class Recorder: - def __init__(self, offset=0): - self.offset = offset - self.nope = False - - def build(self, common_prefix_len, common_attn_metadata, shared_mla_rope=None, **kwargs): - common = common_attn_metadata - positions = common.positions[: common.num_input_tokens] - if self.offset: - positions = positions + self.offset - if shared_mla_rope is None: - cos, sin = counted_lookup(positions, use_cache=True) - else: - cos, sin = shared_mla_rope.get(common.positions, positions) - torch.testing.assert_close(cos, table[positions].unsqueeze(1).unsqueeze(2)) - torch.testing.assert_close(sin, -table[positions].unsqueeze(1).unsqueeze(2)) - # The unrelated consumer owns its copies, like a separate builder. - return SimpleNamespace(cos=cos.clone() if self.offset else cos, sin=sin.clone() if self.offset else sin) - - class PlainSFA(Recorder, AscendSFAMetadataBuilder): - pass - - class PlainIndexer(Recorder, AscendSFAIndexerMetadataBuilder): - pass - - class UnrelatedBuilder(Recorder, AttentionMetadataBuilder): - pass - - builders = [("sfa", PlainSFA())] - if boundary == "builder": - builders.append(("other", UnrelatedBuilder(5))) - builders.append(("indexer", PlainIndexer())) - groups = [ - [SimpleNamespace(layer_names=[name], get_metadata_builder=lambda _, builder=builder: builder)] - for name, builder in builders - ] - positions = torch.tensor([1, 2, 3], dtype=torch.int64) - config = SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=None) for _ in builders]) - - def update_positions(group, common=None): - if group: - positions.add_(1) - return {} - - provider = SimpleNamespace( - get_extra_common_attn_kwargs=update_positions, - get_extra_attn_kwargs=lambda *args: {}, - ) - for _ in range(2): - lookups.clear() - metadata = attn_utils.build_attn_metadata( - attn_groups=groups, - num_reqs=1, - num_tokens=3, - query_start_loc_gpu=torch.tensor([0, 3], dtype=torch.int32), - query_start_loc_cpu=torch.tensor([0, 3], dtype=torch.int32), - max_query_len=3, - seq_lens=torch.tensor([10], dtype=torch.int32), - max_seq_len=10, - block_tables=tuple(torch.zeros(1, 1, dtype=torch.int32) for _ in builders), - slot_mappings=tuple(torch.zeros(3, dtype=torch.int32) for _ in builders), - kv_cache_config=config, - positions=positions, - full_graph_mode=full_graph, - for_cudagraph_capture=for_capture, - offload_dummy=for_capture, - is_prefilling=np.array([True]), - common_metadata_callback=update_positions if boundary == "callback" else None, - model_specific_attn_metadata=provider if boundary == "provider" else None, - parallel_config=SimpleNamespace(prefill_context_parallel_size=1, decode_context_parallel_size=1), - ) - expected = table[positions].unsqueeze(1).unsqueeze(2) - for name in ("sfa", "indexer"): - torch.testing.assert_close(metadata[name].cos, expected) - torch.testing.assert_close(metadata[name].sin, -expected) - expected_lookups = 1 if boundary == "none" and not for_capture else len(builders) - assert len(lookups) == expected_lookups - # The same position tensor is updated between calls; no prior-build hit. - positions.add_(1) - - -@pytest.mark.parametrize("change", ["positions", "table", "destination", "rebind"]) -def test_mla_rope_cache_refreshes_current_values_after_mutation(monkeypatch, change): - from vllm_ascend.ops import rotary_embedding as rotary - - table = torch.arange(128, dtype=torch.float32).reshape(32, 4) - monkeypatch.setattr(rotary, "_cos_cache", table) - monkeypatch.setattr(rotary, "_sin_cache", -table) - monkeypatch.setattr(rotary, "_cos_sin_flat", None) - monkeypatch.setattr(rotary, "_cos_mla", torch.empty(3, 1, 1, 4)) - monkeypatch.setattr(rotary, "_sin_mla", torch.empty(3, 1, 1, 4)) - positions = torch.tensor([1, 2, 3]) - cache = rotary.MLARopeMetadataCache() - cache.get(positions, positions) - if change == "positions": - positions.add_(1) - elif change == "table": - table.add_(1) - elif change == "destination": - rotary._cos_mla.zero_() - else: - monkeypatch.setattr(rotary, "_cos_cache", table + 1) - cos, sin = cache.get(positions, positions) - torch.testing.assert_close(cos, rotary._cos_cache[positions].unsqueeze(1).unsqueeze(2)) - torch.testing.assert_close(sin, rotary._sin_cache[positions].unsqueeze(1).unsqueeze(2)) - - class SparseAttentionBackend(AscendAttentionBackend): @classmethod def is_sparse(cls) -> bool: diff --git a/vllm_ascend/attention/context_parallel/sfa_cp.py b/vllm_ascend/attention/context_parallel/sfa_cp.py index b0813d10f711..f41078d65aed 100644 --- a/vllm_ascend/attention/context_parallel/sfa_cp.py +++ b/vllm_ascend/attention/context_parallel/sfa_cp.py @@ -781,12 +781,7 @@ def build( global_dcp_num_blocks = pcp_context.global_block_table_num_blocks[pcp_cache_group_idx] metadata = self._build_with_metadata_view( common_attn_metadata, - lambda: self._build( - common_attn_metadata, - draft_index=None, - pcp_context=pcp_context, - shared_mla_rope=kwargs.get("shared_mla_rope"), - ), + lambda: self._build(common_attn_metadata, draft_index=None, pcp_context=pcp_context), global_dcp_block_table=global_dcp_block_table, global_dcp_num_blocks=global_dcp_num_blocks, ) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index 158ebf0c978b..3ea0f6cf9159 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -36,7 +36,7 @@ ) from vllm_ascend.attention.utils import split_decodes_and_prefills, uses_pcp_sharded_layout from vllm_ascend.device.device_op import DeviceOperator -from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache, get_cos_and_sin_mla +from vllm_ascend.ops.rotary_embedding import get_cos_and_sin_mla from vllm_ascend.ops.triton.rope import rope_forward_triton_siso from vllm_ascend.ops.triton.sfa_indexer_store import store_indexer_key_scale from vllm_ascend.utils import ( @@ -928,8 +928,6 @@ def _build( buffer_key: object, use_cached_rope: bool, copy_rope: bool, - *, - shared_mla_rope: MLARopeMetadataCache | None = None, **kwargs, ) -> AscendSFAIndexerMetadata: num_reqs = common_attn_metadata.num_reqs @@ -955,10 +953,7 @@ def _build( input_positions = common_attn_metadata.positions[:num_input_tokens].long() block_size = self.kernel_block_size - if shared_mla_rope is not None and use_cached_rope: - cos, sin = shared_mla_rope.get(common_attn_metadata.positions, input_positions) - else: - cos, sin = get_cos_and_sin_mla(input_positions, use_cache=use_cached_rope) + cos, sin = get_cos_and_sin_mla(input_positions, use_cache=use_cached_rope) cos = cos[:num_input_tokens] sin = sin[:num_input_tokens] cum_query_lens = common_attn_metadata.query_start_loc[1 : num_reqs + 1] diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 088da8019233..8d994ea0f978 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -19,7 +19,6 @@ build_hisparse_prefill_staging_plan, stage_prefill_rows, ) -from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache # Bound the temporary D2D packing workspace independently of prefill length. _ATTENTION_CHUNK_ROWS = 32 @@ -68,10 +67,8 @@ def _build( common_attn_metadata, draft_index=None, pcp_context=None, - *, - shared_mla_rope: MLARopeMetadataCache | None = None, ): - metadata = super()._build(common_attn_metadata, draft_index, pcp_context, shared_mla_rope=shared_mla_rope) + metadata = super()._build(common_attn_metadata, draft_index, pcp_context) metadata.num_reqs = common_attn_metadata.num_reqs if draft_index not in self._request_row_buffers: self._request_row_buffers[draft_index] = torch.empty( diff --git a/vllm_ascend/attention/sfa_v1.py b/vllm_ascend/attention/sfa_v1.py index 9b328493d034..c8b4fc01121c 100644 --- a/vllm_ascend/attention/sfa_v1.py +++ b/vllm_ascend/attention/sfa_v1.py @@ -56,7 +56,7 @@ OFFLOAD_KV_CACHE_TUPLE_LEN, OFFLOAD_V_CACHE_NPU_INDEX, ) -from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache, get_cos_and_sin_mla +from vllm_ascend.ops.rotary_embedding import get_cos_and_sin_mla from vllm_ascend.quantization.methods import ( AscendW8A8DynamicLinearMethod, AscendW8A8LinearMethod, @@ -588,19 +588,12 @@ def build( common_prefix_len: int, common_attn_metadata: AscendCommonAttentionMetadata, fast_build: bool = False, - *, - shared_mla_rope: MLARopeMetadataCache | None = None, **kwargs, ) -> AscendSFAMetadata: # common_prefix_len / fast_build are unused; kept for API compatibility. metadata = self._build_with_metadata_view( common_attn_metadata, - lambda: self._build( - common_attn_metadata, - draft_index=None, - pcp_context=common_attn_metadata.pcp_context, - shared_mla_rope=shared_mla_rope, - ), + lambda: self._build(common_attn_metadata, draft_index=None, pcp_context=common_attn_metadata.pcp_context), ) if ( self.use_pcp @@ -654,8 +647,6 @@ def _build( common_attn_metadata: AscendCommonAttentionMetadata, draft_index: int | None = None, pcp_context: Any | None = None, - *, - shared_mla_rope: MLARopeMetadataCache | None = None, ) -> AscendSFAMetadata: num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens @@ -685,8 +676,6 @@ def _build( if self.nope: cos, sin = None, None - elif shared_mla_rope is not None and draft_index is None: - cos, sin = shared_mla_rope.get(common_attn_metadata.positions, input_positions) else: cos, sin = get_cos_and_sin_mla(input_positions, use_cache=(draft_index is None)) diff --git a/vllm_ascend/ops/rotary_embedding.py b/vllm_ascend/ops/rotary_embedding.py index cda198c35c31..c88b76c875d9 100644 --- a/vllm_ascend/ops/rotary_embedding.py +++ b/vllm_ascend/ops/rotary_embedding.py @@ -148,59 +148,6 @@ def get_cos_and_sin_mla(positions, use_cache=False): return _cos_mla[:num_tokens, ...], _sin_mla[:num_tokens, ...] -class MLARopeMetadataCache: - """Share RoPE within one metadata build, between position/storage writers. - - The caller invalidates sharing across custom providers and other builders. - Inference tensors have no version counter, so these write boundaries also - protect their contents. A new cache is used for every metadata build. - """ - - def __init__(self) -> None: - self._num_tokens = 0 - self._value: tuple[torch.Tensor, torch.Tensor] | None = None - self._owners: tuple = () - self._versions: tuple = () - - @staticmethod - def _versions_of(tensors: tuple) -> tuple: - return tuple(None if tensor is None or tensor.is_inference() else tensor._version for tensor in tensors) - - @staticmethod - def _uses_persistent_buffers(value: tuple[torch.Tensor, torch.Tensor]) -> bool: - cos, sin = value - return ( - _cos_mla is not None - and _sin_mla is not None - and cos.numel() > 0 - and sin.numel() > 0 - and cos.device == _cos_mla.device - and sin.device == _sin_mla.device - and cos.data_ptr() == _cos_mla.data_ptr() - and sin.data_ptr() == _sin_mla.data_ptr() - ) - - def get(self, positions: torch.Tensor, input_positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - num_tokens = input_positions.shape[0] - owners = (positions, _cos_cache, _sin_cache, _cos_sin_flat, _cos_mla, _sin_mla) - if ( - self._num_tokens == num_tokens - and self._value is not None - and all(previous is current for previous, current in zip(self._owners, owners)) - and self._versions == self._versions_of(owners) - and self._uses_persistent_buffers(self._value) - ): - return self._value - value = get_cos_and_sin_mla(input_positions, use_cache=True) - self._num_tokens = num_tokens - self._owners = owners - self._versions = self._versions_of(owners) - # Preserve the original fallback for absent, empty or mismatched graph - # buffers. Only slices of the stable destinations may be reused. - self._value = value if self._uses_persistent_buffers(value) else None - return value - - def _record_cos_sin_cache(cos_sin_cache): global _cos_sin_cache if _cos_sin_cache is not None: diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 562de2499e1a..04a3cb50d06c 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -69,7 +69,6 @@ from vllm_ascend.attention.dsa_v1 import AscendDSAMetadataBuilder from vllm_ascend.attention.dsa_v41 import AscendDSAV41MetadataBuilder from vllm_ascend.attention.fa_v1 import AscendFlashAttentionBackend -from vllm_ascend.attention.indexer import AscendSFAIndexerMetadataBuilder from vllm_ascend.attention.sfa_v1 import AscendSFAMetadataBuilder from vllm_ascend.attention.utils import ( AscendCommonAttentionMetadata, @@ -110,7 +109,6 @@ make_qsa_kernel_page_cache, make_slab_view, ) -from vllm_ascend.ops.rotary_embedding import MLARopeMetadataCache from vllm_ascend.quantization.methods.kv_cache.turboquant import TURBOQUANT_CACHE_DTYPE from vllm_ascend.quantization.methods.kv_cache.turboquant.cache import uses_turboquant_groups from vllm_ascend.quantization.utils import enable_fa_quant @@ -549,8 +547,6 @@ def build_attn_metadata( if num_actual_reqs is None: num_actual_reqs = num_reqs - shared_mla_rope = None - # positions will not be used directly in graph modes, so it is saft to create it here. if positions is None: positions = torch.zeros(num_input_tokens, dtype=torch.int64, device=query_start_loc_gpu.device) @@ -601,9 +597,6 @@ def build_attn_metadata( group_causal = causal if isinstance(causal, bool) else causal.get(i, True) common_v41_metadata: dict[str, Any] = {} - if model_specific_attn_metadata is not None: - # Custom providers may update positions or persistent RoPE storage. - shared_mla_rope = None common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) if model_specific_attn_metadata is not None @@ -647,7 +640,6 @@ def build_attn_metadata( ) if common_metadata_callback is not None: common_metadata_callback(i, common_attn_metadata) - shared_mla_rope = None for attn_group in attn_groups[i]: attn_metadata_builder = attn_group.get_metadata_builder(0) @@ -673,8 +665,6 @@ def build_attn_metadata( metadata_variants = ([(target_names, common_attn_metadata)] if target_names else []) + [ (draft_names, draft_common) ] - if not for_cudagraph_capture and model_specific_attn_metadata is not None: - shared_mla_rope = None attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_attn_kwargs( attn_metadata_builder, @@ -683,18 +673,6 @@ def build_attn_metadata( if not for_cudagraph_capture and model_specific_attn_metadata is not None else {} ) - # Core's capture entry point does not accept the shared RoPE context. - # Keep sharing on the runtime build interface that declares it. - if not for_cudagraph_capture and isinstance( - attn_metadata_builder, (AscendSFAMetadataBuilder, AscendSFAIndexerMetadataBuilder) - ): - if shared_mla_rope is None: - shared_mla_rope = MLARopeMetadataCache() - attn_metadata_extra_kwargs["shared_mla_rope"] = shared_mla_rope - else: - # An unrelated builder may update the same persistent RoPE - # storage. Share only within runs of known compatible builders. - shared_mla_rope = None if is_dsa_builder: # DSA cache groups share request-level metadata during replay. attn_metadata_extra_kwargs.update( diff --git a/vllm_ascend/worker/v2/model_states/default.py b/vllm_ascend/worker/v2/model_states/default.py index a2964fd780a3..5f99b5546915 100644 --- a/vllm_ascend/worker/v2/model_states/default.py +++ b/vllm_ascend/worker/v2/model_states/default.py @@ -561,7 +561,11 @@ def prepare_attn( copy_sfa_restore_tails=restore_copy_sfa_tails, draft_layer_names=getattr(self, "_offload_draft_layer_names", None), common_metadata_callback=( - self._decoder_replay_common_metadata.__setitem__ if self.decoder_replay is not None else None + lambda group_id, common: ( + self._decoder_replay_common_metadata.__setitem__(group_id, common) + if self.decoder_replay is not None + else None + ) ), ) self._prepare_decoder_replay(input_batch, cudagraph_mode, attn_groups, kv_cache_config) From 77923814e4cb94fba6c4b0fd38baa6f04ffcaccd Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 21:19:17 +0800 Subject: [PATCH 102/120] chore(hisparse): drop unrelated model-state formatting Signed-off-by: zhenwenqi2024 --- vllm_ascend/worker/v2/model_states/default.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm_ascend/worker/v2/model_states/default.py b/vllm_ascend/worker/v2/model_states/default.py index 5f99b5546915..73a290369f97 100644 --- a/vllm_ascend/worker/v2/model_states/default.py +++ b/vllm_ascend/worker/v2/model_states/default.py @@ -561,11 +561,9 @@ def prepare_attn( copy_sfa_restore_tails=restore_copy_sfa_tails, draft_layer_names=getattr(self, "_offload_draft_layer_names", None), common_metadata_callback=( - lambda group_id, common: ( - self._decoder_replay_common_metadata.__setitem__(group_id, common) - if self.decoder_replay is not None - else None - ) + lambda group_id, common: self._decoder_replay_common_metadata.__setitem__(group_id, common) + if self.decoder_replay is not None + else None ), ) self._prepare_decoder_replay(input_batch, cudagraph_mode, attn_groups, kv_cache_config) From 56ec03370a21cef16b7f762e2c8710a31cb93657 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 22:12:20 +0800 Subject: [PATCH 103/120] refactor(hisparse): simplify Prefill staging and Indexer validation Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 42 ----- .../test_hisparse_foundation_cpu.py | 3 +- vllm_ascend/attention/sfa_hisparse.py | 30 ++-- vllm_ascend/hisparse/backend.py | 5 +- vllm_ascend/hisparse/binding.py | 4 +- vllm_ascend/hisparse/cache_ops.py | 6 +- vllm_ascend/hisparse/prefill.py | 153 ------------------ vllm_ascend/ops/triton/hisparse.py | 6 +- 8 files changed, 21 insertions(+), 228 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 57faea4a6cff..3c37600e45a1 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -12,7 +12,6 @@ from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, - gather_resident_rows, scatter_packed_kv, ) from vllm_ascend.hisparse.host_pool import allocate_mapped_host @@ -133,47 +132,6 @@ def test_backend_resident_payload_graph_replay(dtype, width, index_dtype, num_ro graph.reset() -@pytest.mark.parametrize("row_bytes", [37, 656]) -@pytest.mark.parametrize("capture", [False, True]) -def test_prefill_resident_patch_changing_extents(row_bytes, capture): - if not torch.npu.is_available(): - pytest.skip("NPU required") - torch.npu.set_device(0) - block = 4 - for rows, source_blocks in ((8, 4), (24, 6), (40, 8)): - source_cpu = (torch.arange(source_blocks * 5 * block * row_bytes) % 251).to(torch.uint8) - source_cpu = source_cpu.view(source_blocks, 5, block, row_bytes) - source = source_cpu.npu() - destination = torch.full((rows // block, 2, block, row_bytes), 7, dtype=torch.uint8, device="npu:0") - indices = torch.empty(rows, dtype=torch.int32, device="npu:0") - source_view, destination_view = source[:, 0], destination[:, 0] - - indices.fill_(-1) - gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) - torch.npu.synchronize() - graph = None - if capture: - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph): - gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) - for iteration in range(2): - values = (torch.arange(rows, dtype=torch.int32) + iteration * 3) % (source_blocks * block + 2) - 1 - current_source = source_cpu + iteration * 13 - source.copy_(current_source) - indices.copy_(values) - destination.fill_(7) - if graph is None: - gather_resident_rows(source_view, destination_view, indices, skip_invalid=True) - else: - graph.replay() - torch.npu.synchronize() - expected = torch.full(destination.shape, 7, dtype=torch.uint8) - for row, selected in enumerate(values.tolist()): - if 0 <= selected < source_blocks * block: - expected[row // block, 0, row % block] = current_source[selected // block, 0, selected % block] - assert torch.equal(destination.cpu(), expected) - - def test_prefill_singleton_column_mapping_tracks_current_residency(): if not torch.npu.is_available(): pytest.skip("HiSparse prefill mapping requires an NPU") diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 6e6b31a40118..eea8a284ea2f 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -24,7 +24,6 @@ from vllm_ascend.hisparse.binding import ( _bind_indexer_views, align_hisparse_indexer_specs, - get_hisparse_indexer_specs, ) from vllm_ascend.hisparse.c8 import ( c8_storage_view, @@ -311,7 +310,7 @@ def test_unsupported_indexer_formats_are_rejected_before_pool_binding(indexer_sp spec = replace(indexer_spec, **{field: 1}) config = NS(kv_cache_groups=[NS(layer_names=["indexer"], kv_cache_spec=spec)]) with pytest.raises(ValueError, match="C4|C8|scales"): - get_hisparse_indexer_specs(config) + AscendHiSparseBackend().validate_kv_cache_config(config) def test_capture_metadata_prevents_warmup_host_mirrors(monkeypatch): diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 8d994ea0f978..d67e54bbf2cc 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -17,7 +17,6 @@ from vllm_ascend.hisparse.prefill import ( AscendHiSparsePrefillStagingPlan, build_hisparse_prefill_staging_plan, - stage_prefill_rows, ) # Bound the temporary D2D packing workspace independently of prefill length. @@ -140,7 +139,6 @@ def __init__(self, *args, **kwargs): if not isinstance(self.index_group, HiSparseMLAIndexGroup): raise RuntimeError("Ascend HiSparse was initialized without a HiSparse index group") self.hisparse_cache = self.index_group.cache(self.index_group_index) - self.hisparse_cache._ascend_prefill_layer_name = self.layer_name self.topk = self.index_group.logical_topk_indices.shape[-1] if self.topk % 128: raise ValueError("Ascend HiSparse SFA requires top-k to be divisible by 128") @@ -531,8 +529,8 @@ def _execute_sparse_flash_attention_process( self.index_group.set_logical_topk_ready(self.index_group_index) outputs = [] num_decode = metadata.num_decode_tokens + cache = self.hisparse_cache if num_decode: - cache = self.hisparse_cache assert cache.source_block_table is not None # Keep mixed prefill/decode batches on the NPU residency resolver. # The generic index-group resident shortcut uses the CUDA sparse @@ -565,21 +563,19 @@ def _execute_sparse_flash_attention_process( query, query_rope, logical[num_decode:], metadata, actual_seq_lengths_query, actual_seq_lengths_key ) if prefill_output is None: - staged, staged_table, req_ids, slot = stage_prefill_rows( - self.index_group, self.index_group_index, metadata + host = cache.runtime.host_cache + host = host.view(-1, metadata.prefill.host_staging_plan.block_size, host.shape[-1]) + staged, staged_table, req_ids = self.index_group.stage_prefill_rows( + self.index_group_index, host, metadata ) - try: - selected = logical[num_decode:] - nonnegative = selected.clamp_min(0) - block_ids = nonnegative.to(torch.int64) // metadata.block_size - request_tables = staged_table[req_ids.to(torch.int64)] - physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size - physical = physical + nonnegative % metadata.block_size - physical = physical.masked_fill(selected < 0, -1).to(torch.int32) - prefill_output = self._attention_physical(query, query_rope, staged, physical) - finally: - if slot is not None: - slot.finish() + selected = logical[num_decode:] + nonnegative = selected.clamp_min(0) + block_ids = nonnegative.to(torch.int64) // metadata.block_size + request_tables = staged_table[req_ids.to(torch.int64)] + physical = torch.gather(request_tables, 1, block_ids) * metadata.block_size + physical = physical + nonnegative % metadata.block_size + physical = physical.masked_fill(selected < 0, -1).to(torch.int32) + prefill_output = self._attention_physical(query, query_rope, staged, physical) outputs.append(prefill_output) if not outputs: return torch.zeros_like(ql_nope) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 0ff602703f7f..6cf9ac8ff506 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -169,9 +169,10 @@ def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: return ops def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: - from .binding import get_hisparse_indexer_specs + from .binding import _validate_indexer_spec, get_hisparse_indexer_specs - get_hisparse_indexer_specs(kv_cache_config) + for spec in get_hisparse_indexer_specs(kv_cache_config).values(): + _validate_indexer_spec(spec) def prepare_kv_cache_views( self, diff --git a/vllm_ascend/hisparse/binding.py b/vllm_ascend/hisparse/binding.py index 3abbe55f33ea..f9b6ccb5a3f7 100644 --- a/vllm_ascend/hisparse/binding.py +++ b/vllm_ascend/hisparse/binding.py @@ -56,7 +56,7 @@ def align_hisparse_indexer_specs( def get_hisparse_indexer_specs( kv_cache_config: KVCacheConfig, ) -> dict[str, AscendSFAIndexerCacheSpec]: - """Validate the static indexer format before common storage allocation.""" + """Collect the Ascend indexer specs keyed by layer.""" indexer_specs = {} for group in kv_cache_config.kv_cache_groups: for name in group.layer_names: @@ -64,7 +64,6 @@ def get_hisparse_indexer_specs( if isinstance(spec, UniformTypeKVCacheSpecs): spec = spec.kv_cache_specs[name] if isinstance(spec, AscendSFAIndexerCacheSpec): - _validate_indexer_spec(spec) indexer_specs[name] = spec return indexer_specs @@ -90,7 +89,6 @@ def _validate_indexer_spec(spec: AscendSFAIndexerCacheSpec) -> None: def _quantized_indexer_views(cache: torch.Tensor, spec: AscendSFAIndexerCacheSpec): """Alias the current Core allocation; retain no additional storage owner.""" - _validate_indexer_spec(spec) pages, heads, block, width = cache.shape stride = cache.stride(0) key_bytes = block * heads * width diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index 4bb08274a58c..abe66a77ecd9 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -397,12 +397,9 @@ def hisparse_invalidate_written_slots( ) -def gather_resident_rows( - source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor, *, skip_invalid: bool = False -) -> None: +def gather_resident_rows(source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor) -> None: """Copy packed Resident rows using Core's negative-ID-to-row0 convention. - With ``skip_invalid``, retain staged Host bytes at invalid IDs instead. Source and destination may use different dtypes with the same byte width. """ _check_packed_cache(source) @@ -438,7 +435,6 @@ def gather_resident_rows( ROW_BYTES=row_bytes, BLOCK_ROWS=block_rows, BLOCK_BYTES=block_bytes, - SKIP_INVALID=skip_invalid, multibuffer=False, ) diff --git a/vllm_ascend/hisparse/prefill.py b/vllm_ascend/hisparse/prefill.py index 2c169fdedb9d..269546bcbadd 100644 --- a/vllm_ascend/hisparse/prefill.py +++ b/vllm_ascend/hisparse/prefill.py @@ -4,19 +4,11 @@ from __future__ import annotations -from dataclasses import dataclass, field - import torch from torch.overrides import TorchFunctionMode -from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import get_forward_context from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan -from vllm_ascend.hisparse.cache_ops import gather_resident_rows - -_PREFILL_LOOKAHEAD = "ascend_hisparse_prefill_lookahead" - class _PrefillColumnSelection(TorchFunctionMode): """Keep a local workaround for the NPU singleton-column indexing defect. @@ -51,151 +43,6 @@ def ensure_gpu_sources(self, resident_state_rows, state_indices, resident_block_ super().ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size) -@dataclass(eq=False) -class _StagingSlot: - tensor: torch.Tensor - available: torch.npu.Event - in_use: bool = False - - def finish(self) -> None: - # This event follows the last Attention read on the consumer stream. - self.available = torch.npu.Event() - self.available.record() - self.in_use = False - - -@dataclass -class _PrefillLookahead: - slots: list[_StagingSlot] = field(default_factory=list) - pending: tuple | None = None - - -def _host_view(runtime, block_size): - cache = runtime.host_cache - return cache.view(-1, block_size, cache.shape[-1]) - - -def _stage_next(state, cache, plan) -> None: - runtime = cache.runtime - host = _host_view(runtime, plan.block_size) - rows = plan.row_ids.size(1) - if rows == 0: - return - shape = (rows // plan.block_size, plan.block_size, host.size(2)) - slot = next( - ( - slot - for slot in state.slots - if not slot.in_use - and slot.tensor.shape == shape - and slot.tensor.dtype == host.dtype - and slot.tensor.device == plan.block_table.device - ), - None, - ) - if slot is None and len(state.slots) == 2: - return - stream = runtime.index_group.copy_stream - side = torch.npu.Stream( - stream_id=stream.stream_id, device_index=stream.device_index, device_type=stream.device_type - ) - # The scheduler's prior Host publication and this plan's device writes - # precede this boundary on the compute stream. - side.wait_stream(torch.npu.current_stream()) - if slot is not None: - side.wait_event(slot.available) - # A failed enqueue must not make the possibly-written slot reusable. - slot.in_use = True - ready = torch.npu.Event() - with torch.npu.stream(side): - for tensor in (plan.row_ids, plan.dst_rows, plan.miss_mask): - tensor.record_stream(side) - if slot is None: - staged = runtime.gather_prefill_cache(host, plan, resident_cache=None) - else: - slot.tensor.record_stream(side) - runtime.cache_ops.hisparse_gather_plan( - host.view(-1, host.size(2)), - slot.tensor, - plan.row_ids, - plan.dst_rows, - plan.miss_mask, - None, - None, - 0, - ) - ready.record(side) - if slot is None: - slot = _StagingSlot(staged, ready) - state.slots.append(slot) - else: - slot.available = ready - slot.in_use = False - state.pending = (cache, plan, plan.gpu_row_ids, plan.miss_mask, slot) - - -def stage_prefill_rows(group, layer_index, metadata): - """Pipeline historical Host misses for one same-indexer follower.""" - cache = group.cache(layer_index) - plan = metadata.prefill.host_staging_plan - host = _host_view(cache.runtime, plan.block_size) - if cache.all_context_pages_resident or metadata.hisparse_for_capture: - return (*group.stage_prefill_rows(layer_index, host, metadata), None) - context = get_forward_context() - if ( - context.cudagraph_runtime_mode != CUDAGraphMode.NONE - or not isinstance(context.attn_metadata, dict) - or torch.npu.is_current_stream_capturing() - ): - return (*group.stage_prefill_rows(layer_index, host, metadata), None) - - state = context.additional_kwargs.get(_PREFILL_LOOKAHEAD) - pending = None - if state is not None: - pending, state.pending = state.pending, None - slot = None - if pending is not None: - owner, owner_plan, gpu_rows, miss_mask, prefetched = pending - if owner is cache and owner_plan is plan and gpu_rows is plan.gpu_row_ids and miss_mask is plan.miss_mask: - slot = prefetched - if slot is not None: - slot.in_use = True - main = torch.npu.current_stream() - main.wait_event(slot.available) - slot.tensor.record_stream(main) - # The follower has now written its current KV. Patch its Resident - # rows without overwriting the prefetched Host bytes. - gather_resident_rows(cache.view.cache, slot.tensor, plan.gpu_row_ids[0], skip_invalid=True) - staged, table = slot.tensor, plan.block_table - req_ids = metadata.req_id_per_token[metadata.num_decode_tokens :] - if metadata.num_decodes: - req_ids = req_ids - metadata.num_decodes - else: - staged, table, req_ids = group.stage_prefill_rows(layer_index, host, metadata) - - next_index = layer_index + 1 - if next_index < len(group.caches) and plan.gpu_row_ids is not None and cache.view is not None: - follower = group.cache(next_index) - if ( - context.attn_metadata.get(getattr(follower, "_ascend_prefill_layer_name", None)) is metadata - and follower.view is not None - and follower.block_table is not None - and cache.block_table is not None - and follower.block_table.data_ptr() == cache.block_table.data_ptr() - and follower.view.block_size == cache.view.block_size - and plan.gpu_source_key == (cache.block_table.data_ptr(), cache.view.block_size) - and cache.source_block_table is not None - and follower.source_block_table is cache.source_block_table - and follower.runtime.request_state_indices is cache.runtime.request_state_indices - and follower.runtime.index_group is cache.runtime.index_group - ): - if state is None: - state = _PrefillLookahead() - context.additional_kwargs[_PREFILL_LOOKAHEAD] = state - _stage_next(state, follower, plan) - return staged, table, req_ids, slot - - def build_hisparse_prefill_staging_plan( block_table: torch.Tensor, seq_lens: torch.Tensor, diff --git a/vllm_ascend/ops/triton/hisparse.py b/vllm_ascend/ops/triton/hisparse.py index f0c4fb3f4e66..890b76838ade 100644 --- a/vllm_ascend/ops/triton/hisparse.py +++ b/vllm_ascend/ops/triton/hisparse.py @@ -226,16 +226,14 @@ def gather_resident_rows_kernel( ROW_BYTES: tl.constexpr, BLOCK_ROWS: tl.constexpr, BLOCK_BYTES: tl.constexpr, - SKIP_INVALID: tl.constexpr, ): row = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64) src = tl.load(row_ids + row, mask=row < ROWS, other=-1).to(tl.int64) - if not SKIP_INVALID: - src = tl.maximum(src, 0) + src = tl.maximum(src, 0) src_offset = src // SOURCE_BLOCK * SOURCE_STRIDE + src % SOURCE_BLOCK * ROW_BYTES dst_offset = row // DEST_BLOCK * DEST_STRIDE + row % DEST_BLOCK * ROW_BYTES byte = tl.program_id(1) * BLOCK_BYTES + tl.arange(0, BLOCK_BYTES) - valid = (row < ROWS) & (src >= 0) & (src < SOURCE_ROWS) + valid = (row < ROWS) & (src < SOURCE_ROWS) mask = valid[:, None] & (byte[None, :] < ROW_BYTES) payload = tl.load(source + src_offset[:, None] + byte[None, :], mask=mask, other=0) tl.store(destination + dst_offset[:, None] + byte[None, :], payload, mask=mask) From 554334ca4960661e2cc8d2f725ed87a7a07871d7 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 22:36:56 +0800 Subject: [PATCH 104/120] refactor(hisparse): consolidate resolver dispatch and SFA tests Signed-off-by: zhenwenqi2024 --- .github/workflows/scripts/test_config.yaml | 3 +- .../one_card/test_hisparse_attention.py | 240 ++++++++++++++++- .../one_card/test_hisparse_kv_preprocess.py | 242 ------------------ .../one_card/test_hisparse_ops.py | 41 ++- .../test_hisparse_foundation_cpu.py | 72 +++--- vllm_ascend/hisparse/backend.py | 4 +- vllm_ascend/hisparse/cache_ops.py | 54 +++- vllm_ascend/hisparse/resolver.py | 55 ---- 8 files changed, 346 insertions(+), 365 deletions(-) delete mode 100644 tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py delete mode 100644 vllm_ascend/hisparse/resolver.py diff --git a/.github/workflows/scripts/test_config.yaml b/.github/workflows/scripts/test_config.yaml index 56aa0c225171..1be9594360a7 100644 --- a/.github/workflows/scripts/test_config.yaml +++ b/.github/workflows/scripts/test_config.yaml @@ -49,7 +49,6 @@ curated_tests: - tests/e2e/pull_request/one_card/test_flash_attn_precision.py - tests/e2e/pull_request/one_card/test_swiglu_mx_quant.py - tests/e2e/pull_request/one_card/test_hisparse_attention.py - - tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py - tests/e2e/pull_request/one_card/test_hisparse_ops.py - tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -125,7 +124,7 @@ runner_mapping: default: a5-4 tests/e2e/pull_request/one_card/test_swiglu_mx_quant\.py$: default: a5-4 - tests/e2e/pull_request/one_card/test_hisparse_(attention|kv_preprocess|ops|prefill_staging)\.py$: + tests/e2e/pull_request/one_card/test_hisparse_(attention|ops|prefill_staging)\.py$: default: a5-4 tests/e2e/pull_request/one_card: default: a2-1 diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 6dad6291d0f2..c4d9a395ed49 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -8,10 +8,10 @@ import pytest import torch -import torch_npu # noqa: F401 +import torch_npu -@pytest.fixture(autouse=True) +@pytest.fixture def _require_npu_sfa(): from vllm_ascend.utils import bootstrap_custom_op_env @@ -32,6 +32,7 @@ def _require_npu_sfa(): @pytest.mark.parametrize("capture", [False, True]) @pytest.mark.parametrize("strided", [False, True]) @pytest.mark.parametrize("num_queries", [1, 7, 35]) +@pytest.mark.usefixtures("_require_npu_sfa") def test_sparse_attention_reference(num_queries, strided, capture): from vllm_ascend.attention.sfa_hisparse import _ATTENTION_CHUNK_ROWS, AscendHiSparseImpl @@ -114,6 +115,7 @@ def test_sparse_attention_reference(num_queries, strided, capture): @pytest.mark.parametrize("strided", [False, True]) @pytest.mark.parametrize("index_mode", ["virtual", "mapped", "mapped_lengths"]) @pytest.mark.parametrize("empty_batch", [False, True]) +@pytest.mark.usefixtures("_require_npu_sfa") def test_c8_attention_uses_existing_native_operator(strided, capture, index_mode, empty_batch): from vllm_ascend.device.device_config import is_950 @@ -216,3 +218,237 @@ def forward(): torch.npu.synchronize() if graph is not None: graph.reset() + + +@pytest.fixture +def _require_a5_preprocess(): + if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): + pytest.skip("HiSparse integration currently targets Ascend 950") + torch.npu.set_device(0) + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.usefixtures("_require_a5_preprocess") +def test_preprocess_writes_hisparse_mapping_once(cache_dtype, capture): + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope + + rows, block = 5, 128 + c8 = cache_dtype == torch.float8_e4m3fn + width = 656 if c8 else 576 + torch.manual_seed(71) + values = torch.randn(rows, 2624, dtype=torch.bfloat16, device="npu")[:, 2048:] + gamma = torch.ones(512, dtype=torch.bfloat16, device="npu") + cos = torch.ones(rows, 1, 1, 64, dtype=torch.bfloat16, device="npu") + sin = torch.zeros_like(cos) + raw = torch.full((3, 2, block, width * cache_dtype.itemsize), 90, dtype=torch.uint8, device="npu") + destination = raw[:, 1].view(cache_dtype) + logical = [-1, 129, 0, 257, 128] + slots = torch.tensor(logical, dtype=torch.int64, device="npu") + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="npu"), hisparse_for_capture=True + ) + scratch = ( + torch.empty(1, block, 1, 512, dtype=values.dtype, device="npu"), + torch.empty(1, block, 1, 64, dtype=values.dtype, device="npu"), + torch.arange(block, dtype=torch.int64, device="npu"), + ) + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl.enable_sparse_sfa_c8 = c8 + impl.enable_sparse_sfa_turboquant = False + impl.c8_cache_dtype = torch.float8_e4m3fn + impl.sfa_qsfa_tile_size = 128 + impl.kv_lora_rank, impl.qk_rope_head_dim, impl.num_kv_heads = 512, 64, 1 + impl._kv_staging_capacity = block + impl._kv_workspace = lambda dtype: scratch + impl.kv_a_layernorm = SimpleNamespace(weight=gamma, variance_epsilon=1e-5) + impl.hisparse_cache = SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)) + # The ordinary source cache has no Resident/Hot addressing. It must never + # receive the parent's second C8 store after exec_kv has used write_target. + source_cache = (torch.empty(0, device="npu"),) + + def forward(): + rope, latent, *scales = impl.exec_kv(values, cos, sin, source_cache, slots, metadata) + scales = scales[0] if scales else None + impl._store_parallel_kv(rope, latent, scales, source_cache, slots, metadata) + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + forward() + for count in (4, 2): + values.neg_() + metadata.query_start_loc[-1:].fill_(count) + raw.fill_(90) + if graph is None: + forward() + else: + graph.replay() + torch.npu.synchronize() + if c8: + rope, latent, scales = custom_kv_rmsnorm_rope( + values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 + ) + components = (latent, rope, scales) + else: + # Reuse native SFA as the oracle; test only the new destination, + # strides, row mapping and padding, not RMSNorm arithmetic. + reference_k, reference_r = torch.empty_like(scratch[0]), torch.empty_like(scratch[1]) + torch_npu.npu_kv_rmsnorm_rope_cache( + values.view(rows, 1, 1, 576), + gamma, + cos, + sin, + scratch[2][:rows], + reference_r, + reference_k, + epsilon=1e-5, + cache_mode="PA", + ) + components = (reference_k.view(-1, 512)[:rows], reference_r.view(-1, 64)[:rows]) + payload = torch.cat( + [part.contiguous().view(torch.uint8).reshape(rows, -1).cpu() for part in components], -1 + ) + expected = torch.full(raw.shape, 90, dtype=torch.uint8) + for row, slot in enumerate(logical): + if row < count and slot >= 0: + expected[slot // block, 1, slot % block] = payload[row] + torch.testing.assert_close(raw.cpu(), expected, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("rows", [4, 16]) +@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) +@pytest.mark.usefixtures("_require_a5_preprocess") +def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture, slot_dtype): + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.v1.hisparse.runtime import HiSparseCacheHandle + + from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl + from vllm_ascend.attention.sfa_v1 import AscendSFAImpl + from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod + + torch.manual_seed(71) + device, block, width = "npu", 128, 656 + impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) + impl._quant_type = AscendW8A8MXFP8DynamicLinearMethod + impl.enable_sparse_sfa_c8 = True + impl.c8_cache_dtype = torch.float8_e4m3fn + impl.q_lora_rank, impl.kv_lora_rank, impl.qk_rope_head_dim = 2048, 512, 64 + impl.local_num_heads, impl.has_indexer = 8, True + impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False + impl.sfa_qsfa_tile_size = block + + def mx_layer(inputs, outputs): + weight, scale = torch_npu.npu_dynamic_mx_quant( + torch.randn(outputs, inputs, dtype=torch.bfloat16, device=device) * 0.02, + dst_type=torch.float8_e4m3fn, + ) + return SimpleNamespace(weight=weight.T.contiguous(), weight_scale=scale.transpose(0, 1).contiguous()) + + impl.fused_qkv_a_proj = mx_layer(6144, 2624) + impl.q_proj = mx_layer(2048, 2048) + impl.sfa_qsfa_k_nope_clip_alpha = None + impl.W_UK_T = torch.randn(8, 192, 512, dtype=torch.bfloat16, device=device) * 0.02 + impl.q_a_layernorm = SimpleNamespace( + weight=torch.ones(2048, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 + ) + impl.kv_a_layernorm = SimpleNamespace( + weight=torch.ones(512, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 + ) + # Use the parent's production quantized-weight preparation and Prolog as + # the oracle. This case tests ownership/addressing, not Prolog arithmetic. + impl._process_weights_for_fused_prolog_v3() + impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16, device=device) + reference = torch.empty((1, block, 1, width), dtype=impl.c8_cache_dtype, device=device) + local_slots = torch.arange(block, dtype=torch.int64, device=device) + raw = torch.full((3, 2, block, width), 90, dtype=torch.uint8, device=device) + logical = [-1, 129, 0, 257, 128][:rows] + list(range(16, 16 + max(0, rows - 5))) + slots = torch.tensor(logical, dtype=slot_dtype, device=device) + cache = HiSparseCacheHandle(SimpleNamespace(max_num_reqs=rows, kv_dtype=impl.c8_cache_dtype, row_width=width)) + cache.bind_cache( + raw.view(-1), + byte_offset=block * width, + block_stride=2 * block * width, + num_blocks=3, + block_size=block, + block_table=torch.empty(0, dtype=torch.int32, device=device), + slot_mapping=slots, + ) + cache.dummy_batch = True # This component test has no Host-mirror connector. + cache.num_actual_tokens = rows + impl.hisparse_cache = cache + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, rows], dtype=torch.int32, device=device), hisparse_for_capture=capture + ) + x = torch.randn(rows, 6144, dtype=torch.bfloat16, device=device) + cos = torch.ones(rows, 1, 1, 64, dtype=x.dtype, device=device) + sin = torch.zeros_like(cos) + + def forward(): + context = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + with override_forward_context(context): + return impl._sfa_preprocess_prolog_v3(x, (), cos, sin, slots, attn_metadata=metadata) + + def compare(actual, expected): + if isinstance(actual, tuple): + assert len(actual) == len(expected) + for a, e in zip(actual, expected): + compare(a, e) + elif actual is not None: + if actual.dtype == torch.float8_e4m3fn: + actual, expected = actual.view(torch.uint8), expected.view(torch.uint8) + torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0) + else: + assert expected is None + + graph = None + try: + forward() + torch.npu.synchronize() + if capture: + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + captured = forward() + for valid in (rows - 1, 2, 0, rows): + x.neg_() + metadata.query_start_loc[-1:].fill_(valid) + # The real slot producer marks dummy and padded rows with -1. + # Replay consumes these changed device values at fixed addresses. + slots.copy_( + torch.tensor( + [slot if row < valid else -1 for row, slot in enumerate(logical)], dtype=slot_dtype, device=device + ) + ) + if not capture: + cache.num_actual_tokens = valid + raw.fill_(90) + if graph is None: + actual = forward() + else: + graph.replay() + actual = captured + expected = AscendSFAImpl._sfa_preprocess_prolog_v3( + impl, x, (reference,), cos, sin, local_slots[:rows], attn_metadata=metadata + ) + torch.npu.synchronize() + compare(actual, expected) + payload = reference.view(torch.uint8).reshape(block, width).cpu() + wanted = torch.full(raw.shape, 90, dtype=torch.uint8) + for row, slot in enumerate(logical): + if row < valid and slot >= 0: + wanted[slot // block, 1, slot % block] = payload[row] + torch.testing.assert_close(raw.cpu(), wanted, rtol=0, atol=0) + finally: + torch.npu.synchronize() + if graph is not None: + graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py b/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py deleted file mode 100644 index 4679a72b7005..000000000000 --- a/tests/e2e/pull_request/one_card/test_hisparse_kv_preprocess.py +++ /dev/null @@ -1,242 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Check HiSparse destinations against the existing SFA preprocessing ops.""" - -from types import SimpleNamespace - -import pytest -import torch -import torch_npu - - -def _setup(): - if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): - pytest.skip("HiSparse integration currently targets Ascend 950") - torch.npu.set_device(0) - - -@pytest.mark.parametrize("capture", [False, True]) -@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) -def test_preprocess_writes_hisparse_mapping_once(cache_dtype, capture): - _setup() - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - from vllm_ascend.attention.sfa_v1 import custom_kv_rmsnorm_rope - - rows, block = 5, 128 - c8 = cache_dtype == torch.float8_e4m3fn - width = 656 if c8 else 576 - torch.manual_seed(71) - values = torch.randn(rows, 2624, dtype=torch.bfloat16, device="npu")[:, 2048:] - gamma = torch.ones(512, dtype=torch.bfloat16, device="npu") - cos = torch.ones(rows, 1, 1, 64, dtype=torch.bfloat16, device="npu") - sin = torch.zeros_like(cos) - raw = torch.full((3, 2, block, width * cache_dtype.itemsize), 90, dtype=torch.uint8, device="npu") - destination = raw[:, 1].view(cache_dtype) - logical = [-1, 129, 0, 257, 128] - slots = torch.tensor(logical, dtype=torch.int64, device="npu") - metadata = SimpleNamespace( - query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="npu"), hisparse_for_capture=True - ) - scratch = ( - torch.empty(1, block, 1, 512, dtype=values.dtype, device="npu"), - torch.empty(1, block, 1, 64, dtype=values.dtype, device="npu"), - torch.arange(block, dtype=torch.int64, device="npu"), - ) - impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) - impl.enable_sparse_sfa_c8 = c8 - impl.enable_sparse_sfa_turboquant = False - impl.c8_cache_dtype = torch.float8_e4m3fn - impl.sfa_qsfa_tile_size = 128 - impl.kv_lora_rank, impl.qk_rope_head_dim, impl.num_kv_heads = 512, 64, 1 - impl._kv_staging_capacity = block - impl._kv_workspace = lambda dtype: scratch - impl.kv_a_layernorm = SimpleNamespace(weight=gamma, variance_epsilon=1e-5) - impl.hisparse_cache = SimpleNamespace(write_target=lambda count, slot_count: (destination, slots, count)) - # The ordinary source cache has no Resident/Hot addressing. It must never - # receive the parent's second C8 store after exec_kv has used write_target. - source_cache = (torch.empty(0, device="npu"),) - - def forward(): - rope, latent, *scales = impl.exec_kv(values, cos, sin, source_cache, slots, metadata) - scales = scales[0] if scales else None - impl._store_parallel_kv(rope, latent, scales, source_cache, slots, metadata) - - graph = None - try: - forward() - torch.npu.synchronize() - if capture: - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph): - forward() - for count in (4, 2): - values.neg_() - metadata.query_start_loc[-1:].fill_(count) - raw.fill_(90) - if graph is None: - forward() - else: - graph.replay() - torch.npu.synchronize() - if c8: - rope, latent, scales = custom_kv_rmsnorm_rope( - values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 - ) - components = (latent, rope, scales) - else: - # Reuse native SFA as the oracle; test only the new destination, - # strides, row mapping and padding, not RMSNorm arithmetic. - reference_k, reference_r = torch.empty_like(scratch[0]), torch.empty_like(scratch[1]) - torch_npu.npu_kv_rmsnorm_rope_cache( - values.view(rows, 1, 1, 576), - gamma, - cos, - sin, - scratch[2][:rows], - reference_r, - reference_k, - epsilon=1e-5, - cache_mode="PA", - ) - components = (reference_k.view(-1, 512)[:rows], reference_r.view(-1, 64)[:rows]) - payload = torch.cat( - [part.contiguous().view(torch.uint8).reshape(rows, -1).cpu() for part in components], -1 - ) - expected = torch.full(raw.shape, 90, dtype=torch.uint8) - for row, slot in enumerate(logical): - if row < count and slot >= 0: - expected[slot // block, 1, slot % block] = payload[row] - torch.testing.assert_close(raw.cpu(), expected, rtol=0, atol=0) - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() - - -@pytest.mark.parametrize("capture", [False, True]) -@pytest.mark.parametrize("rows", [4, 16]) -@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) -def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture, slot_dtype): - _setup() - from vllm.forward_context import ForwardContext, override_forward_context - from vllm.v1.hisparse.runtime import HiSparseCacheHandle - - from vllm_ascend.attention.sfa_hisparse import AscendHiSparseImpl - from vllm_ascend.attention.sfa_v1 import AscendSFAImpl - from vllm_ascend.quantization.methods import AscendW8A8MXFP8DynamicLinearMethod - - torch.manual_seed(71) - device, block, width = "npu", 128, 656 - impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) - impl._quant_type = AscendW8A8MXFP8DynamicLinearMethod - impl.enable_sparse_sfa_c8 = True - impl.c8_cache_dtype = torch.float8_e4m3fn - impl.q_lora_rank, impl.kv_lora_rank, impl.qk_rope_head_dim = 2048, 512, 64 - impl.local_num_heads, impl.has_indexer = 8, True - impl.is_kv_consumer, impl.rl_weight_update_enabled = False, False - impl.sfa_qsfa_tile_size = block - - def mx_layer(inputs, outputs): - weight, scale = torch_npu.npu_dynamic_mx_quant( - torch.randn(outputs, inputs, dtype=torch.bfloat16, device=device) * 0.02, - dst_type=torch.float8_e4m3fn, - ) - return SimpleNamespace(weight=weight.T.contiguous(), weight_scale=scale.transpose(0, 1).contiguous()) - - impl.fused_qkv_a_proj = mx_layer(6144, 2624) - impl.q_proj = mx_layer(2048, 2048) - impl.sfa_qsfa_k_nope_clip_alpha = None - impl.W_UK_T = torch.randn(8, 192, 512, dtype=torch.bfloat16, device=device) * 0.02 - impl.q_a_layernorm = SimpleNamespace( - weight=torch.ones(2048, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 - ) - impl.kv_a_layernorm = SimpleNamespace( - weight=torch.ones(512, dtype=torch.bfloat16, device=device), variance_epsilon=1e-5 - ) - # Use the parent's production quantized-weight preparation and Prolog as - # the oracle. This case tests ownership/addressing, not Prolog arithmetic. - impl._process_weights_for_fused_prolog_v3() - impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16, device=device) - reference = torch.empty((1, block, 1, width), dtype=impl.c8_cache_dtype, device=device) - local_slots = torch.arange(block, dtype=torch.int64, device=device) - raw = torch.full((3, 2, block, width), 90, dtype=torch.uint8, device=device) - logical = [-1, 129, 0, 257, 128][:rows] + list(range(16, 16 + max(0, rows - 5))) - slots = torch.tensor(logical, dtype=slot_dtype, device=device) - cache = HiSparseCacheHandle(SimpleNamespace(max_num_reqs=rows, kv_dtype=impl.c8_cache_dtype, row_width=width)) - cache.bind_cache( - raw.view(-1), - byte_offset=block * width, - block_stride=2 * block * width, - num_blocks=3, - block_size=block, - block_table=torch.empty(0, dtype=torch.int32, device=device), - slot_mapping=slots, - ) - cache.dummy_batch = True # This component test has no Host-mirror connector. - cache.num_actual_tokens = rows - impl.hisparse_cache = cache - metadata = SimpleNamespace( - query_start_loc=torch.tensor([0, rows], dtype=torch.int32, device=device), hisparse_for_capture=capture - ) - x = torch.randn(rows, 6144, dtype=torch.bfloat16, device=device) - cos = torch.ones(rows, 1, 1, 64, dtype=x.dtype, device=device) - sin = torch.zeros_like(cos) - - def forward(): - context = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) - with override_forward_context(context): - return impl._sfa_preprocess_prolog_v3(x, (), cos, sin, slots, attn_metadata=metadata) - - def compare(actual, expected): - if isinstance(actual, tuple): - assert len(actual) == len(expected) - for a, e in zip(actual, expected): - compare(a, e) - elif actual is not None: - if actual.dtype == torch.float8_e4m3fn: - actual, expected = actual.view(torch.uint8), expected.view(torch.uint8) - torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0) - else: - assert expected is None - - graph = None - try: - forward() - torch.npu.synchronize() - if capture: - graph = torch.npu.NPUGraph() - with torch.npu.graph(graph): - captured = forward() - for valid in (rows - 1, 2, 0, rows): - x.neg_() - metadata.query_start_loc[-1:].fill_(valid) - # The real slot producer marks dummy and padded rows with -1. - # Replay consumes these changed device values at fixed addresses. - slots.copy_( - torch.tensor( - [slot if row < valid else -1 for row, slot in enumerate(logical)], dtype=slot_dtype, device=device - ) - ) - if not capture: - cache.num_actual_tokens = valid - raw.fill_(90) - if graph is None: - actual = forward() - else: - graph.replay() - actual = captured - expected = AscendSFAImpl._sfa_preprocess_prolog_v3( - impl, x, (reference,), cos, sin, local_slots[:rows], attn_metadata=metadata - ) - torch.npu.synchronize() - compare(actual, expected) - payload = reference.view(torch.uint8).reshape(block, width).cpu() - wanted = torch.full(raw.shape, 90, dtype=torch.uint8) - for row, slot in enumerate(logical): - if row < valid and slot >= 0: - wanted[slot // block, 1, slot % block] = payload[row] - torch.testing.assert_close(raw.cpu(), wanted, rtol=0, atol=0) - finally: - torch.npu.synchronize() - if graph is not None: - graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index babfe0ae4dfc..def07a97d901 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Real-NPU HiSparse state and mapped-payload tests against a CPU oracle.""" +from contextlib import contextmanager + import pytest import torch import torch_npu # noqa: F401 @@ -725,6 +727,21 @@ def test_plan_gather_respects_request_state_padding(): mapping.close() +@contextmanager +def _invalidation_graph(forward): + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + forward() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + try: + yield graph, stream + finally: + torch.npu.synchronize() + graph.reset() + + @pytest.mark.parametrize("hot_capacity", [4, 10240]) def test_invalidate_hot_capacity_graph(hot_capacity): dgi = torch.arange(2 * hot_capacity, dtype=torch.int32, device="npu").view(2, hot_capacity) @@ -742,13 +759,7 @@ def forward(): forward() torch.npu.synchronize() torch.testing.assert_close(dgi.cpu(), expected) - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): - forward() - torch.npu.synchronize() - graph = torch.npu.NPUGraph() - try: + with _invalidation_graph(forward) as (graph, stream): with torch.npu.graph(graph, stream=stream): forward() for _ in range(10): @@ -756,9 +767,6 @@ def forward(): graph.replay() torch.npu.synchronize() torch.testing.assert_close(dgi.cpu(), expected) - finally: - torch.npu.synchronize() - graph.reset() @pytest.mark.parametrize( @@ -815,11 +823,6 @@ def check_writes(run): torch.testing.assert_close(table.cpu(), expected) check_writes(forward) - stream = torch.npu.Stream() - stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(stream): - forward() - torch.npu.synchronize() sort = torch.sort sort_calls = 0 @@ -829,16 +832,12 @@ def count_sort(*args, **kwargs): sort_calls += 1 return sort(*args, **kwargs) - monkeypatch.setattr(torch, "sort", count_sort) - graph = torch.npu.NPUGraph() - try: + with _invalidation_graph(forward) as (graph, stream): + monkeypatch.setattr(torch, "sort", count_sort) with torch.npu.graph(graph, stream=stream): forward() assert sort_calls == int(expect_sorted) check_writes(graph.replay) - finally: - torch.npu.synchronize() - graph.reset() @pytest.mark.parametrize("capture", [False, True]) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index eea8a284ea2f..4d77da1254a4 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -173,12 +173,10 @@ def test_mapped_addresses_retain_offsets_and_reject_expired_or_external_views(): def test_resolver_scratch_follows_group_geometry_and_is_not_shared_across_groups( monkeypatch, max_rows, hot_rows, union_rows, prepare_workspaces ): - from vllm_ascend.hisparse import resolver - # Loading a device extension is the only mocked boundary; the group-owned # tensors below are real CPU storage, not a replacement resolver algorithm. native = NS(hisparse_resolver_validate_scratch=lambda hot, union: None) - monkeypatch.setattr(resolver, "_load_resolver", lambda: native) + monkeypatch.setattr(hisparse_ops, "_load_resolver", lambda: native) backend = AscendHiSparseBackend() def group(): @@ -233,11 +231,9 @@ def test_resolver_scratch_rejects_topk_outside_hot_capacity(topk): def test_missing_native_resolver_raises_instead_of_falling_back(monkeypatch): - from vllm_ascend.hisparse import resolver - - monkeypatch.setattr(resolver.torch.ops, "_C_ascend", NS()) + monkeypatch.setattr(hisparse_ops.torch.ops, "_C_ascend", NS()) with pytest.raises(RuntimeError, match="requires its compiled AscendC resolver"): - resolver._load_resolver() + hisparse_ops._load_resolver() @pytest.fixture @@ -275,34 +271,40 @@ def test_expanded_indexer_alias_preserves_shared_slab_offsets(indexer_spec, bloc assert original[-1, 0, 1, 2] == -73 -def test_alignment_updates_only_indexer_specs_without_changing_inputs(indexer_spec): +@pytest.mark.parametrize("c8", [False, True]) +def test_indexer_page_alignment_preserves_input_specs(indexer_spec, c8): source = MLAAttentionSpec(block_size=4, num_kv_heads=1, head_size=4, dtype=torch.float32) - indexer = indexer_spec - values = {"source": source, "indexer": indexer} - result = align_hisparse_indexer_specs(values) + specs = {"source": source, "indexer": indexer_spec} + packed_row_bytes = None + if c8: + indexer_spec = replace( + indexer_spec, + block_size=128, + head_size=128, + dtype=torch.float8_e4m3fn, + scale_dim=1, + scale_dtype=torch.float32, + cache_sparse_li_c8=True, + ) + packed_row_bytes = 656 + specs.update( + indexer=indexer_spec, + resident=HiSparseResidentSpec(block_size=128, page_size=128 * 656, block_stride_alignment=128), + hot=HiSparseHotSpec(block_size=128, page_size=128 * 656, blocks_per_request=1, block_stride_alignment=128), + ) + result = align_hisparse_indexer_specs(specs, packed_row_bytes=packed_row_bytes) assert result["source"] is source - assert result["indexer"].get_block_stride_alignment() == math.lcm(3, indexer.page_size_bytes) - assert values["indexer"] is indexer and indexer.block_stride_alignment == 3 - - -def test_c8_page_alignment_reserves_scale_bytes_and_whole_main_rows(indexer_spec): - resident = HiSparseResidentSpec(block_size=128, page_size=128 * 656, block_stride_alignment=128) - hot = HiSparseHotSpec(block_size=128, page_size=128 * 656, blocks_per_request=1, block_stride_alignment=128) - indexer = replace( - indexer_spec, - block_size=128, - head_size=128, - dtype=torch.float8_e4m3fn, - scale_dim=1, - scale_dtype=torch.float32, - cache_sparse_li_c8=True, - ) - result = align_hisparse_indexer_specs({"resident": resident, "hot": hot, "indexer": indexer}, packed_row_bytes=656) - stride_alignment = math.lcm(*(spec.get_block_stride_alignment() for spec in result.values())) - assert stride_alignment % 656 == 0 and stride_alignment % 128 == 0 - assert result["indexer"].page_size_padded == 128 * (128 + 4) - assert resident.block_stride_alignment == hot.block_stride_alignment == 128 - assert indexer.page_size_padded is None + assert specs["indexer"] is indexer_spec and indexer_spec.block_stride_alignment == 3 + if c8: + alignment = math.lcm(*(spec.get_block_stride_alignment() for name, spec in result.items() if name != "source")) + assert alignment % 656 == 0 and alignment % 128 == 0 + assert result["indexer"].page_size_padded == 128 * (128 + 4) + assert indexer_spec.page_size_padded is None + for name in ("resident", "hot"): + assert specs[name].block_stride_alignment == 128 + assert result[name].block_stride_alignment == math.lcm(128, packed_row_bytes) + else: + assert result["indexer"].get_block_stride_alignment() == math.lcm(3, indexer_spec.page_size_bytes) @pytest.mark.parametrize("field", ["scale_dim", "cache_sparse_li_c8", "cache_sparse_li_c4"]) @@ -364,10 +366,8 @@ def test_c8_storage_mapping_preserves_shared_slab_offsets(stride_rows): def test_c8_decode_lengths_follow_binding_geometry_and_survive_rebind( monkeypatch, pages, stride_rows, max_rows, prepare_workspaces ): - from vllm_ascend.hisparse import resolver - monkeypatch.setattr( - resolver, "_load_resolver", lambda: NS(hisparse_resolver_validate_scratch=lambda hot, union: None) + hisparse_ops, "_load_resolver", lambda: NS(hisparse_resolver_validate_scratch=lambda hot, union: None) ) backend = AscendHiSparseBackend() group = NS(copy_stream=NS(device=torch.device("cpu"))) diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 6cf9ac8ff506..fe27a5bb1076 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -77,14 +77,12 @@ def __init__(self) -> None: self._failed_dma_batches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] def check_available(self) -> None: + from vllm_ascend.hisparse.cache_ops import _load_resolver from vllm_ascend.hisparse.host_pool import install - from vllm_ascend.hisparse.resolver import _load_resolver from vllm_ascend.hisparse.row_dma import install as install_row_dma from vllm_ascend.utils import bootstrap_custom_op_env install() - # Import the device kernels before the connector allocates Host storage. - importlib.import_module("vllm_ascend.hisparse.cache_ops") if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): # A5 disables the global custom-op selection policy. HiSparse still # needs its explicit DMA dependency; load the extension lazily here diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index abe66a77ecd9..ee859e5264e8 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -29,6 +29,15 @@ _SORT_BUFFER_LIMIT_BYTES = 768 * 1024 +def _load_resolver(): + namespace = torch.ops._C_ascend + if not all( + hasattr(namespace, name) for name in ("hisparse_resolver_validate_scratch", "hisparse_resolve_residency") + ): + raise RuntimeError("HiSparse requires its compiled AscendC resolver. Rebuild vllm-ascend for this target.") + return namespace + + class ResolverWorkspace: """Fixed scratch shared by one index group on its ordered copy stream.""" @@ -43,8 +52,6 @@ def __init__( self.max_topk = hot_size if max_topk is None else max_topk if not 0 < self.max_topk <= hot_size: raise ValueError("HiSparse resolver TopK capacity must be positive and fit the Hot cache") - from vllm_ascend.hisparse.resolver import _load_resolver - self.native = _load_resolver() self.native.hisparse_resolver_validate_scratch(hot_size, max_union_rows) # One resolved Host row per selected token, not per Hot-cache slot. @@ -240,9 +247,48 @@ def hisparse_resolve_residency( != 1 ): raise ValueError("All three compact HiSparse swap outputs are required together") - from vllm_ascend.hisparse.resolver import resolve + self._resolve(locals()) - resolve(self, locals()) + def _resolve(self, arguments): + """Dispatch live tensors to the native request-union resolver. + + During replay, external tensors remain owned by the runtime/workspace. + Callers must destroy graphs before releasing or rebinding those owners. + """ + rows, topk = arguments["global_indices"].shape + workspace = self.workspace + resolved = arguments["resolved_global_indices"] + if resolved is None: + resolved = workspace.first_refs[: rows * topk].view(rows, topk) + workspace.native.hisparse_resolve_residency( + arguments["global_indices"], + arguments["request_ids"], + arguments["request_state_indices"], + arguments["source_block_table"], + arguments["resident_block_table"], + arguments["hot_block_table"], + arguments["device_global_indices"], + arguments["lru_slots"], + arguments["hot_indices"], + arguments["attention_indices"], + resolved, + arguments["valid_counts"], + arguments["swap_host_physical_rows"], + arguments["swap_device_physical_rows"], + arguments["swap_counts"], + arguments["miss_mask"], + arguments["stats"], + workspace.errors[:rows], + arguments["num_valid_rows"], + arguments["host_cache"].shape[0], + arguments["hot_cache"].shape[1], + arguments["source_block_size"], + arguments["resident_block_size"], + arguments["resident_null_block"], + arguments["attention_block_stride"] if self.attention_block_stride is None else self.attention_block_stride, + arguments["max_union_rows"], + self.attention_capacity, + ) def hisparse_gather_compact( self, host_cache, hot_cache, miss_global_indices, miss_hot_indices, miss_counts diff --git a/vllm_ascend/hisparse/resolver.py b/vllm_ascend/hisparse/resolver.py deleted file mode 100644 index 6f268a1ba5de..000000000000 --- a/vllm_ascend/hisparse/resolver.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Tensor-only native resolver dispatch; no address-keyed launch ownership.""" - -import torch - - -def _load_resolver(): - namespace = torch.ops._C_ascend - if not all( - hasattr(namespace, name) for name in ("hisparse_resolver_validate_scratch", "hisparse_resolve_residency") - ): - raise RuntimeError("HiSparse requires its compiled AscendC resolver. Rebuild vllm-ascend for this target.") - return namespace - - -def resolve(owner, arguments): - """Dispatch live tensors to the native request-union resolver. - - During replay, external tensors remain owned by the runtime/workspace. - Callers must destroy graphs before releasing or rebinding those owners. - """ - rows, topk = arguments["global_indices"].shape - workspace = owner.workspace - resolved = arguments["resolved_global_indices"] - if resolved is None: - resolved = workspace.first_refs[: rows * topk].view(rows, topk) - workspace.native.hisparse_resolve_residency( - arguments["global_indices"], - arguments["request_ids"], - arguments["request_state_indices"], - arguments["source_block_table"], - arguments["resident_block_table"], - arguments["hot_block_table"], - arguments["device_global_indices"], - arguments["lru_slots"], - arguments["hot_indices"], - arguments["attention_indices"], - resolved, - arguments["valid_counts"], - arguments["swap_host_physical_rows"], - arguments["swap_device_physical_rows"], - arguments["swap_counts"], - arguments["miss_mask"], - arguments["stats"], - workspace.errors[:rows], - arguments["num_valid_rows"], - arguments["host_cache"].shape[0], - arguments["hot_cache"].shape[1], - arguments["source_block_size"], - arguments["resident_block_size"], - arguments["resident_null_block"], - arguments["attention_block_stride"] if owner.attention_block_stride is None else owner.attention_block_stride, - arguments["max_union_rows"], - owner.attention_capacity, - ) From 72218c964bec727abd7b6039b1c97ae9bdb5dc55 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 23:01:01 +0800 Subject: [PATCH 105/120] refactor(hisparse): consolidate Prefill and Host allocation ownership Signed-off-by: zhenwenqi2024 --- .../one_card/test_hisparse_prefill_staging.py | 2 +- vllm_ascend/attention/sfa_hisparse.py | 56 ++++++++++++++++-- vllm_ascend/hisparse/backend.py | 9 --- vllm_ascend/hisparse/host_pool.py | 6 +- vllm_ascend/hisparse/prefill.py | 59 ------------------- 5 files changed, 58 insertions(+), 74 deletions(-) delete mode 100644 vllm_ascend/hisparse/prefill.py diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index 3c37600e45a1..a0cf49482372 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -9,13 +9,13 @@ import torch_npu # noqa: F401 from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView +from vllm_ascend.attention.sfa_hisparse import build_hisparse_prefill_staging_plan from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, scatter_packed_kv, ) from vllm_ascend.hisparse.host_pool import allocate_mapped_host -from vllm_ascend.hisparse.prefill import build_hisparse_prefill_staging_plan @pytest.mark.parametrize("resident_block_size", [128, 64, 32]) diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index d67e54bbf2cc..076d0169e957 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -6,18 +6,17 @@ from dataclasses import dataclass import torch +from torch.overrides import TorchFunctionMode from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import logger from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup, SparseMLAIndexGroupBuilder from vllm.v1.attention.ops.metadata import compute_token_to_req_indices +from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan +from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan from vllm.v1.worker.workspace import current_workspace_manager from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder -from vllm_ascend.hisparse.prefill import ( - AscendHiSparsePrefillStagingPlan, - build_hisparse_prefill_staging_plan, -) # Bound the temporary D2D packing workspace independently of prefill length. _ATTENTION_CHUNK_ROWS = 32 @@ -34,6 +33,55 @@ def finish_hisparse_kv_update(cache, metadata) -> None: cache.finish_kv_update() +class _PrefillColumnSelection(TorchFunctionMode): + """Keep a local workaround for the NPU singleton-column indexing defect. + + torch-npu 2.10.0.post7 / CANN 9.2.1 can interleave rows for ``x[:, ids]`` + when x has one column. index_select has the same semantics and returns + correct rows. Remove this scope when the supported device stack is fixed. + """ + + def __torch_function__(self, func, types, args=(), kwargs=None): + if func is torch.Tensor.__getitem__: + tensor, index = args + if ( + tensor.device.type == "npu" + and tensor.ndim == 2 + and isinstance(index, tuple) + and len(index) == 2 + and isinstance(index[0], slice) + and index[0] == slice(None) + and isinstance(index[1], torch.Tensor) + and index[1].ndim == 1 + ): + return torch.index_select(tensor, 1, index[1]) + return func(*args, **(kwargs or {})) + + +class AscendHiSparsePrefillStagingPlan(HiSparsePrefillStagingPlan): + def ensure_gpu_sources(self, resident_state_rows, state_indices, resident_block_size) -> None: + # Core retains source-key caching, representative-page selection and + # current Resident/Host masks. The override affects only this call. + with _PrefillColumnSelection(): + super().ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size) + + +def build_hisparse_prefill_staging_plan( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + staging_block_capacity: int, +) -> AscendHiSparsePrefillStagingPlan: + plan = build_core_prefill_plan(block_table, seq_lens, block_size, staging_block_capacity) + return AscendHiSparsePrefillStagingPlan( + block_table=plan.block_table, + row_ids=plan.row_ids, + dst_rows=plan.dst_rows, + miss_mask=plan.miss_mask, + block_size=plan.block_size, + ) + + @dataclass class AscendHiSparsePrefillMetadata: host_staging_plan: AscendHiSparsePrefillStagingPlan diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index fe27a5bb1076..dc5ff86a2ddf 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -97,15 +97,6 @@ def check_available(self) -> None: _load_resolver() install_row_dma(AscendHiSparseBackend) - def allocate_private_host_pool(self, size: int) -> tuple[torch.Tensor, torch.Tensor]: - from vllm_ascend.hisparse.host_pool import allocate_mapped_host - - payload = allocate_mapped_host((size,), torch.int8) - registration = payload._hisparse_registration - registered = registration.tensor.view(torch.int8) - registered._hisparse_registration = registration - return registered[:size], registered - def pin_tensor(self, tensor: torch.Tensor) -> None: from vllm_ascend.hisparse.host_pool import MappedHostTensor diff --git a/vllm_ascend/hisparse/host_pool.py b/vllm_ascend/hisparse/host_pool.py index ba318e9ae172..fba0a9132f72 100644 --- a/vllm_ascend/hisparse/host_pool.py +++ b/vllm_ascend/hisparse/host_pool.py @@ -120,7 +120,11 @@ def install() -> None: def allocate(size): backend = get_hisparse_backend() if isinstance(backend, AscendHiSparseBackend): - return backend.allocate_private_host_pool(size) + payload = allocate_mapped_host((size,), torch.int8) + registration = payload._hisparse_registration + registered = registration.tensor.view(torch.int8) + registered._hisparse_registration = registration + return registered[:size], registered return original(size) allocate._ascend_pinned_private_pool = True diff --git a/vllm_ascend/hisparse/prefill.py b/vllm_ascend/hisparse/prefill.py deleted file mode 100644 index 269546bcbadd..000000000000 --- a/vllm_ascend/hisparse/prefill.py +++ /dev/null @@ -1,59 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reuse Core's prefill plan with NPU-safe column selection.""" - -from __future__ import annotations - -import torch -from torch.overrides import TorchFunctionMode -from vllm.v1.hisparse.runtime import HiSparsePrefillStagingPlan -from vllm.v1.hisparse.runtime import build_hisparse_prefill_staging_plan as build_core_prefill_plan - - -class _PrefillColumnSelection(TorchFunctionMode): - """Keep a local workaround for the NPU singleton-column indexing defect. - - torch-npu 2.10.0.post7 / CANN 9.2.1 can interleave rows for ``x[:, ids]`` - when x has one column. index_select has the same semantics and returns - correct rows. Remove this scope when the supported device stack is fixed. - """ - - def __torch_function__(self, func, types, args=(), kwargs=None): - if func is torch.Tensor.__getitem__: - tensor, index = args - if ( - tensor.device.type == "npu" - and tensor.ndim == 2 - and isinstance(index, tuple) - and len(index) == 2 - and isinstance(index[0], slice) - and index[0] == slice(None) - and isinstance(index[1], torch.Tensor) - and index[1].ndim == 1 - ): - return torch.index_select(tensor, 1, index[1]) - return func(*args, **(kwargs or {})) - - -class AscendHiSparsePrefillStagingPlan(HiSparsePrefillStagingPlan): - def ensure_gpu_sources(self, resident_state_rows, state_indices, resident_block_size) -> None: - # Core retains source-key caching, representative-page selection and - # current Resident/Host masks. The override affects only this call. - with _PrefillColumnSelection(): - super().ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size) - - -def build_hisparse_prefill_staging_plan( - block_table: torch.Tensor, - seq_lens: torch.Tensor, - block_size: int, - staging_block_capacity: int, -) -> AscendHiSparsePrefillStagingPlan: - plan = build_core_prefill_plan(block_table, seq_lens, block_size, staging_block_capacity) - return AscendHiSparsePrefillStagingPlan( - block_table=plan.block_table, - row_ids=plan.row_ids, - dst_rows=plan.dst_rows, - miss_mask=plan.miss_mask, - block_size=plan.block_size, - ) From fea0b7517d027c043cb1018032a3d385f1d15fba Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Wed, 11 Nov 2026 23:46:26 +0800 Subject: [PATCH 106/120] refactor(hisparse): reuse Core row DMA descriptor construction Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/row_dma.py | 113 +------------------------------- 1 file changed, 3 insertions(+), 110 deletions(-) diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py index 4f74223c96ac..13c4b0d2926b 100644 --- a/vllm_ascend/hisparse/row_dma.py +++ b/vllm_ascend/hisparse/row_dma.py @@ -2,118 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project """Prepare HiSparse row-mirror DMA through Core's descriptor ownership.""" -from dataclasses import dataclass from functools import wraps -import numpy as np import torch -# Large prefill batches keep the scalar path to bound temporary memory. -_MAX_DESCRIPTORS = 32768 - - -@dataclass(frozen=True) -class _Geometry: - sources: tuple - destinations: tuple - block_size: int - source_columns: np.ndarray - source_addresses: np.ndarray - source_page_bytes: np.ndarray - source_pages: np.ndarray - destination_addresses: np.ndarray - destination_rows: np.ndarray - row_bytes: np.ndarray - - @classmethod - def bind(cls, worker): - sources, destinations = worker.resident_caches, worker.host_caches - block_size = worker.kernel_block_size - facts = [] - for handle, source, destination in zip(worker.cache_handles, sources, destinations, strict=True): - column = handle.runtime.resident_source_index - if column < 0: - return None - item_size = source.element_size() - facts.append( - ( - column, - source.data_ptr(), - source.stride(0) * item_size, - source.shape[0], - destination.data_ptr(), - destination.shape[0], - source.shape[-1] * item_size, - ) - ) - if not facts: - return None - columns = np.ascontiguousarray(np.asarray(facts, dtype=np.int64).T) - columns.flags.writeable = False - # Core _init_dma already validated the rows and bound the cache owners. - return cls(sources, destinations, block_size, *columns) - - -def _try_enqueue(worker, geometry, layer_indices, ready_event) -> bool: - if not layer_indices or not worker._row_mirror_num_rows or not worker.is_host_writer: - return True - num_layers = len(layer_indices) - num_mirrors = len(worker._row_mirrors) - count = num_mirrors * num_layers - if ( - geometry is None - or num_layers < 2 - or count > _MAX_DESCRIPTORS - or geometry.sources is not worker.resident_caches - or geometry.destinations is not worker.host_caches - or geometry.block_size != worker.kernel_block_size - ): - return False - - layers = np.asarray(layer_indices, dtype=np.intp) - starts = worker._row_mirror_source_starts - columns = geometry.source_columns[layers] - if np.any(columns >= starts.shape[1]): - return False - pages, offsets = np.divmod(starts[:, columns], geometry.block_size) - row_counts = worker._row_mirror_counts[:, None] - destinations = worker._row_mirror_destination_starts[:, None] - # Delegate invalid inputs to the original path for its checks and errors. - if ( - np.any(pages < 0) - or np.any(pages >= geometry.source_pages[layers]) - or np.any(offsets + row_counts > geometry.block_size) - or np.any(destinations < 0) - or np.any(destinations + row_counts > geometry.destination_rows[layers]) - ): - return False - - descriptors = worker._acquire_dma_descriptors(count) - shape = (num_mirrors, num_layers) - src = descriptors.src_np[:count].reshape(shape) - dst = descriptors.dst_np[:count].reshape(shape) - row_bytes = geometry.row_bytes[layers] - # The native submission consumes mirror-major descriptors in this order. - np.multiply(pages, geometry.source_page_bytes[layers], out=src) - src += geometry.source_addresses[layers] - src += offsets * row_bytes - np.multiply(destinations, row_bytes, out=dst) - dst += geometry.destination_addresses[layers] - descriptors.sizes_np[:count].reshape(shape)[:] = row_counts * row_bytes - worker._submit_dma_descriptors(descriptors, count, ready_event=ready_event) - return True - def _install_worker(worker_type, backend_type, bind_recycled) -> None: original_init = worker_type._init_dma - original_enqueue = worker_type._enqueue_row_dma - if getattr(original_enqueue, "_ascend_row_dma", False): + if getattr(original_init, "_ascend_row_dma", False): return @wraps(original_init) def init(worker): original_init(worker) - worker._ascend_row_dma_geometry = None if isinstance(worker.backend, backend_type): if worker.is_host_writer and worker.shared_host_region is None: # Core allocates int64 staging; the bound Ascend slot producer uses int32. @@ -125,18 +26,10 @@ def init(worker): staging.slots = torch.empty( staging.slots.shape, dtype=slots.dtype, device="cpu", pin_memory=True ) - worker._ascend_row_dma_geometry = _Geometry.bind(worker) worker._ascend_recycled_invalidator = bind_recycled(worker) - @wraps(original_enqueue) - def enqueue(worker, layer_indices, ready_event=None): - geometry = getattr(worker, "_ascend_row_dma_geometry", None) - if geometry is None or not _try_enqueue(worker, geometry, layer_indices, ready_event): - return original_enqueue(worker, layer_indices, ready_event=ready_event) - - enqueue._ascend_row_dma = True + init._ascend_row_dma = True worker_type._init_dma = init - worker_type._enqueue_row_dma = enqueue original_invalidate = worker_type.invalidate_blocks @wraps(original_invalidate) @@ -149,7 +42,7 @@ def invalidate(worker, block_ids, request_state_indices): def install(backend_type) -> None: - """Install the Ascend row-mirror and recycled-slot adapters.""" + """Match Ascend slot dtype and install the recycled-slot adapter.""" from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator From aed45cef49077d18d17efddea5960fa835cac105 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 00:38:28 +0800 Subject: [PATCH 107/120] refactor(hisparse): consolidate worker DMA hooks into backend Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/backend.py | 43 ++++++++++++++++++++++++++-- vllm_ascend/hisparse/row_dma.py | 50 --------------------------------- 2 files changed, 41 insertions(+), 52 deletions(-) delete mode 100644 vllm_ascend/hisparse/row_dma.py diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index dc5ff86a2ddf..3f37f45ace18 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -5,6 +5,7 @@ from __future__ import annotations import importlib +from functools import wraps from typing import TYPE_CHECKING, Any import torch @@ -16,6 +17,45 @@ from vllm.v1.kv_cache_interface import KVCacheConfig +def _install_worker_dma_hooks() -> None: + """Bind Ascend staging dtype and grouped invalidation to the Core worker.""" + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import HiSparseConnectorWorker + + from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator + + original_init = HiSparseConnectorWorker._init_dma + if getattr(original_init, "_ascend_row_dma", False): + return + + @wraps(original_init) + def init(worker): + original_init(worker) + if isinstance(worker.backend, AscendHiSparseBackend): + if worker.is_host_writer and worker.shared_host_region is None: + # Core allocates int64 staging; the bound Ascend slot producer uses int32. + # Cache binding is complete here and precedes graph capture. + staging = worker._slot_mapping_staging + slots = worker.cache_handles[0].slot_mapping + if staging is not None and not staging.num_tokens and slots.device.type == "npu": + if slots.dtype == torch.int32: + staging.slots = torch.empty( + staging.slots.shape, dtype=slots.dtype, device="cpu", pin_memory=True + ) + worker._ascend_recycled_invalidator = RecycledGroupInvalidator.bind(worker) + + init._ascend_row_dma = True + HiSparseConnectorWorker._init_dma = init + original_invalidate = HiSparseConnectorWorker.invalidate_blocks + + @wraps(original_invalidate) + def invalidate(worker, block_ids, request_state_indices): + owner = getattr(worker, "_ascend_recycled_invalidator", None) + if owner is None or not owner.try_call(worker, block_ids, request_state_indices): + return original_invalidate(worker, block_ids, request_state_indices) + + HiSparseConnectorWorker.invalidate_blocks = invalidate + + def _allocate_resolver_workspace(runtime: HiSparseRuntime): from vllm_ascend.hisparse.cache_ops import ResolverWorkspace @@ -79,7 +119,6 @@ def __init__(self) -> None: def check_available(self) -> None: from vllm_ascend.hisparse.cache_ops import _load_resolver from vllm_ascend.hisparse.host_pool import install - from vllm_ascend.hisparse.row_dma import install as install_row_dma from vllm_ascend.utils import bootstrap_custom_op_env install() @@ -95,7 +134,7 @@ def check_available(self) -> None: if not hasattr(torch.ops._C_ascend, "swap_blocks_batch"): raise RuntimeError("Ascend HiSparse requires the Ascend swap_blocks_batch operator.") _load_resolver() - install_row_dma(AscendHiSparseBackend) + _install_worker_dma_hooks() def pin_tensor(self, tensor: torch.Tensor) -> None: from vllm_ascend.hisparse.host_pool import MappedHostTensor diff --git a/vllm_ascend/hisparse/row_dma.py b/vllm_ascend/hisparse/row_dma.py deleted file mode 100644 index 13c4b0d2926b..000000000000 --- a/vllm_ascend/hisparse/row_dma.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project -"""Prepare HiSparse row-mirror DMA through Core's descriptor ownership.""" - -from functools import wraps - -import torch - - -def _install_worker(worker_type, backend_type, bind_recycled) -> None: - original_init = worker_type._init_dma - if getattr(original_init, "_ascend_row_dma", False): - return - - @wraps(original_init) - def init(worker): - original_init(worker) - if isinstance(worker.backend, backend_type): - if worker.is_host_writer and worker.shared_host_region is None: - # Core allocates int64 staging; the bound Ascend slot producer uses int32. - # Cache binding is complete here and precedes graph capture. - staging = worker._slot_mapping_staging - slots = worker.cache_handles[0].slot_mapping - if staging is not None and not staging.num_tokens and slots.device.type == "npu": - if slots.dtype == torch.int32: - staging.slots = torch.empty( - staging.slots.shape, dtype=slots.dtype, device="cpu", pin_memory=True - ) - worker._ascend_recycled_invalidator = bind_recycled(worker) - - init._ascend_row_dma = True - worker_type._init_dma = init - original_invalidate = worker_type.invalidate_blocks - - @wraps(original_invalidate) - def invalidate(worker, block_ids, request_state_indices): - owner = getattr(worker, "_ascend_recycled_invalidator", None) - if owner is None or not owner.try_call(worker, block_ids, request_state_indices): - return original_invalidate(worker, block_ids, request_state_indices) - - worker_type.invalidate_blocks = invalidate - - -def install(backend_type) -> None: - """Match Ascend slot dtype and install the recycled-slot adapter.""" - from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import worker - - from vllm_ascend.hisparse.recycled_invalidate import RecycledGroupInvalidator - - _install_worker(worker.HiSparseConnectorWorker, backend_type, RecycledGroupInvalidator.bind) From c4d8bc97a76ee18108661bbf80403c0fb38c5116 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:01:17 +0800 Subject: [PATCH 108/120] test(hisparse): use real attention defaults in worker fixtures Signed-off-by: zhenwenqi2024 --- tests/ut/worker/test_attn_utils_v2.py | 19 +++++++------- tests/ut/worker/test_model_runner_v2.py | 3 +++ tests/ut/worker/test_worker_multi_instance.py | 5 +++- tests/ut/worker/test_worker_v1.py | 25 +++++++++++-------- 4 files changed, 31 insertions(+), 21 deletions(-) diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index 56d1f2097a60..04449bdccdea 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -6,6 +6,7 @@ import numpy as np import pytest import torch +from vllm.config.attention import AttentionConfig from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache @@ -1132,7 +1133,7 @@ def test_gqa_pcp_dcp_cache_keeps_tp_local_kv_heads(monkeypatch, dcp_size): prefill_context_parallel_size=2, decode_context_parallel_size=dcp_size, ), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(indexer_kv_dtype="auto"), cache_config=SimpleNamespace(cache_dtype="auto", block_size=128), ) monkeypatch.setattr( @@ -1178,7 +1179,7 @@ def test_draft_layer_specs_are_tagged_for_eagle_groups(monkeypatch, is_draft_lay prefill_context_parallel_size=1, decode_context_parallel_size=1, ), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(indexer_kv_dtype="auto"), cache_config=SimpleNamespace(cache_dtype="auto", block_size=128), ) monkeypatch.setattr( @@ -1230,7 +1231,7 @@ def test_mla_cache_spec_uses_physical_heads_for_kimi_k3( decode_context_parallel_size=1, prefill_context_parallel_size=1, ), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(indexer_kv_dtype="auto"), cache_config=SimpleNamespace(block_size=128, cache_dtype="auto"), ) monkeypatch.setattr( @@ -1290,7 +1291,7 @@ def test_sfa_indexer_cache_spec_runtime_ownership_and_dcp_replication( additional_config={}, parallel_config=SimpleNamespace(decode_context_parallel_size=4), cache_config=SimpleNamespace(block_size=128, cache_dtype="auto"), - attention_config=SimpleNamespace(indexer_kv_dtype="int8", hisparse_config=None), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), model_config=SimpleNamespace( dtype=torch.bfloat16, hf_text_config=SimpleNamespace(index_head_dim=128), @@ -1372,7 +1373,7 @@ def test_mrv2_initializes_dsv4_cache_only_layer( ) vllm_config = SimpleNamespace( additional_config={}, - attention_config=SimpleNamespace(indexer_kv_dtype="int8", hisparse_config=None), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), model_config=SimpleNamespace( dtype=torch.bfloat16, hf_config=SimpleNamespace( @@ -2366,7 +2367,7 @@ def test_kpool_sfa_c8_main_cache_spec(monkeypatch): layer.get_kv_cache_spec = lambda _cfg: SimpleNamespace(**{**vars(source_spec), "head_size": 512}) vllm_config = SimpleNamespace( additional_config={}, - attention_config=SimpleNamespace(indexer_kv_dtype="int8"), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), cache_config=SimpleNamespace(block_size=16, cache_dtype="int8"), kv_transfer_config=None, model_config=SimpleNamespace( @@ -2401,7 +2402,7 @@ def test_v2_sparse_offload_marks_main_mla_cache_host_resident(monkeypatch): vllm_config = SimpleNamespace( parallel_config=SimpleNamespace(decode_context_parallel_size=1), cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), - attention_config=SimpleNamespace(indexer_kv_dtype="int8"), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), model_config=SimpleNamespace( dtype=torch.bfloat16, hf_text_config=SimpleNamespace(kv_lora_rank=128, qk_rope_head_dim=64), @@ -2717,7 +2718,7 @@ def test_attn_state_mla_spec_and_metadata_wrappers(monkeypatch): vllm_config = SimpleNamespace( parallel_config=SimpleNamespace(decode_context_parallel_size=1), cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), - attention_config=SimpleNamespace(indexer_kv_dtype="int8", hisparse_config=None), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), model_config=SimpleNamespace( dtype=torch.bfloat16, hf_text_config=SimpleNamespace(kv_lora_rank=128, qk_rope_head_dim=64), @@ -2884,7 +2885,7 @@ def _c8_mxfp_vllm_config(cache_dtype: str = "mxfp8") -> SimpleNamespace: return SimpleNamespace( kv_transfer_config=None, cache_config=SimpleNamespace(cache_dtype=cache_dtype), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(indexer_kv_dtype="auto"), parallel_config=SimpleNamespace(decode_context_parallel_size=1), model_config=SimpleNamespace(hf_text_config=SimpleNamespace(model_type="qwen3_next")), quant_config=None, diff --git a/tests/ut/worker/test_model_runner_v2.py b/tests/ut/worker/test_model_runner_v2.py index 8c6b67605cab..b834e82fe6ee 100644 --- a/tests/ut/worker/test_model_runner_v2.py +++ b/tests/ut/worker/test_model_runner_v2.py @@ -9,6 +9,7 @@ import pytest import torch from vllm.config import CompilationMode, CUDAGraphMode +from vllm.config.attention import AttentionConfig from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.offloader.base import NoopOffloader @@ -1237,6 +1238,7 @@ def test_initialize_kv_cache_installs_aclgraph_factory_and_pcp(a5, architecture) runner.input_buffers = object() runner.compilation_config = SimpleNamespace(static_forward_context={}) runner.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), compilation_config=runner.compilation_config, scheduler_config=SimpleNamespace(max_num_seqs=8), kv_transfer_config=None, @@ -1304,6 +1306,7 @@ def test_initialize_kv_cache_forwards_allocation_context(): runner = _make_runner() runner.compilation_config = SimpleNamespace(static_forward_context={}) runner.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), compilation_config=runner.compilation_config, scheduler_config=SimpleNamespace(max_num_seqs=8), kv_transfer_config=None, diff --git a/tests/ut/worker/test_worker_multi_instance.py b/tests/ut/worker/test_worker_multi_instance.py index 4d11725a1946..e610667070c0 100644 --- a/tests/ut/worker/test_worker_multi_instance.py +++ b/tests/ut/worker/test_worker_multi_instance.py @@ -18,6 +18,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +from vllm.config.attention import AttentionConfig from vllm.utils.mem_constants import GiB_bytes from tests.ut.base import TestBase @@ -53,7 +54,9 @@ def _make_worker( with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() - worker.vllm_config = SimpleNamespace(kv_transfer_config=None, additional_config={}) + worker.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), kv_transfer_config=None, additional_config={} + ) worker._kvpp_cache_allocation_plan = None worker.model_runner = MagicMock() worker.model_runner.model_memory_usage = model_memory_usage diff --git a/tests/ut/worker/test_worker_v1.py b/tests/ut/worker/test_worker_v1.py index 11ece7a8ce6f..e1603ebaec09 100644 --- a/tests/ut/worker/test_worker_v1.py +++ b/tests/ut/worker/test_worker_v1.py @@ -5,6 +5,7 @@ import torch from vllm.config import CacheConfig, ModelConfig, ParallelConfig, ProfilerConfig, VllmConfig +from vllm.config.attention import AttentionConfig from vllm.sequence import IntermediateTensors from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -348,7 +349,7 @@ def test_deepseek_v4_shared_tuple_layout_does_not_scale_budget(self): ) ] worker = NPUWorker.__new__(NPUWorker) - worker.vllm_config = SimpleNamespace() + worker.vllm_config = SimpleNamespace(attention_config=AttentionConfig()) worker.get_kv_cache_spec = MagicMock(return_value={"attn": spec}) with patch("vllm_ascend.worker.worker.get_kv_cache_groups", return_value=groups): @@ -377,6 +378,7 @@ def test_hybrid_budget_scaling_follows_runner_backing_layout(self): cache_config = SimpleNamespace(get_resolved_kv_cache_layout=lambda: layout) worker = NPUWorker.__new__(NPUWorker) worker.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), cache_config=cache_config, # Connector identity must not change standardized shared-backing # allocation or its corresponding budget calculation. @@ -422,6 +424,7 @@ def test_pure_attention_multi_group_budget_scales_for_private_layout(self): ] worker = NPUWorker.__new__(NPUWorker) worker.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: SimpleNamespace( is_layer_compact=True, @@ -1251,7 +1254,7 @@ def test_explicit_kv_cache_memory_applies_sparse_offload_constraints(self): worker = NPUWorker.__new__(NPUWorker) worker._kvpp_cache_allocation_plan = None worker.cache_config = SimpleNamespace(kv_cache_memory_bytes=8192) - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.init_snapshot = SimpleNamespace(free_memory=16384) worker._scale_kv_cache_memory_for_multi_group = MagicMock(return_value=6144) worker._apply_kv_offload_decode_memory_constraints = MagicMock(return_value=4096) @@ -1279,7 +1282,7 @@ def test_explicit_kv_cache_memory_converts_kvpp_budget_after_offload(self): worker = NPUWorker.__new__(NPUWorker) worker.cache_config = SimpleNamespace(kv_cache_memory_bytes=8192) - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.init_snapshot = SimpleNamespace(free_memory=16384) worker.model_runner = MagicMock() worker._scale_kv_cache_memory_for_multi_group = MagicMock(return_value=6144) @@ -1346,7 +1349,7 @@ def test_determine_available_memory_normal_case( with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() worker._kvpp_cache_allocation_plan = None - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.init_snapshot = mock_init_snapshot worker.requested_memory = 10000 * 0.8 worker.model_runner = MagicMock() @@ -1391,7 +1394,7 @@ def test_determine_available_memory_applies_startup_plan_before_fast_path( with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() worker._kvpp_cache_allocation_plan = None - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.cache_config = MagicMock() worker.cache_config.kv_cache_memory_bytes = None worker.init_snapshot = MagicMock(free_memory=8 << 30) @@ -1452,7 +1455,7 @@ def test_determine_available_memory_with_non_torch_allocations( with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() worker._kvpp_cache_allocation_plan = None - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.init_snapshot = mock_init_snapshot worker.requested_memory = 10000 * 0.9 worker.model_runner = MagicMock() @@ -1506,7 +1509,7 @@ def test_determine_available_memory_memory_profiling_error( worker.init_snapshot = mock_init_snapshot worker.requested_memory = 10000 * 0.8 worker.model_runner = MagicMock() - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.cache_config = MagicMock() worker.cache_config.gpu_memory_utilization = 0.8 worker.cache_config.kv_cache_memory_bytes = None @@ -1561,7 +1564,7 @@ def test_determine_available_memory_negative_result( with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() worker._kvpp_cache_allocation_plan = None - worker.vllm_config = MagicMock(kv_transfer_config=None) + worker.vllm_config = MagicMock(attention_config=AttentionConfig(), kv_transfer_config=None) worker.init_snapshot = mock_init_snapshot worker.requested_memory = 10000 * 0.8 worker.model_runner = MagicMock() @@ -1903,7 +1906,7 @@ def test_compile_or_warm_up_model_with_eager_mode( ): worker = NPUWorker() worker.model_runner = MagicMock() - worker.vllm_config = MagicMock() + worker.vllm_config = MagicMock(attention_config=AttentionConfig()) worker.model_config = MagicMock() worker.model_config.enforce_eager = True worker.model_config.seed = 12345 @@ -1969,7 +1972,7 @@ def test_compile_or_warm_up_model_with_graph_capture( ): worker = NPUWorker() worker.model_runner = MagicMock() - worker.vllm_config = MagicMock() + worker.vllm_config = MagicMock(attention_config=AttentionConfig()) worker.model_config = MagicMock() worker.model_config.enforce_eager = False # Enable graph capture worker.model_config.seed = 67890 @@ -2031,7 +2034,7 @@ def test_compile_or_warm_up_model_saves_startup_plan_after_graph_capture(self): worker.model_runner = MagicMock() worker.model_runner.model_memory_usage = model_memory worker.model_runner.capture_model.return_value = npugraph_memory - worker.vllm_config = MagicMock() + worker.vllm_config = MagicMock(attention_config=AttentionConfig()) worker.vllm_config.compilation_config.compile_sizes = [] worker.vllm_config.compilation_config.cudagraph_capture_sizes = [] worker.vllm_config.compilation_config.get_compile_ranges.return_value = [] From 29e6b32cec31169df7c6a9054ddc3f4b71c48b8f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:06:41 +0800 Subject: [PATCH 109/120] test(hisparse): cover the native miss loop and supported Core interface Signed-off-by: zhenwenqi2024 --- csrc/hisparse/hisparse_resolver.cpp | 2 ++ .../user_guide/feature_guide/hisparse.md | 17 +++------- .../one_card/test_hisparse_attention.py | 2 ++ .../one_card/test_hisparse_ops.py | 32 ++++++++++++------- .../one_card/test_hisparse_prefill_staging.py | 4 ++- .../test_hisparse_foundation_cpu.py | 4 ++- 6 files changed, 35 insertions(+), 26 deletions(-) diff --git a/csrc/hisparse/hisparse_resolver.cpp b/csrc/hisparse/hisparse_resolver.cpp index 404b21cb259b..084b8b88c9e4 100644 --- a/csrc/hisparse/hisparse_resolver.cpp +++ b/csrc/hisparse/hisparse_resolver.cpp @@ -487,6 +487,8 @@ __simt_vf__ __launch_bounds__(kThreads) inline void resolve_vf( if (swap_device) swap_device[int64_t(first) * topk + rank] = physical; store_miss(d, int64_t(row) * topk + col, 1); } + // Finish this tile's reads and miss-slot writes before the next tile + // reuses the block prefix scratch or the output pass reads the slots. asc_syncthreads(); consumed += tile_misses; } diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index 0e1433d6bcc7..94402a340343 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -18,19 +18,10 @@ Attention interface. Attention core arithmetic is unchanged. C4, TurboQuant, INT8-format C8, context parallelism and external P/D transport are outside this backend's current support. -## Build the resolver - -Use the paired merged vLLM version in `.github/vllm-main-verified.commit` and -install vllm-ascend normally in the matching Torch/torch-npu/CANN environment. -The standard CMake build compiles `csrc/hisparse/hisparse_resolver.cpp` using -CANN's ASC integration when the Ascend 950 SIMT headers are available. It links -the kernel into `vllm_ascend_C` and registers -`torch.ops._C_ascend.hisparse_resolve_residency` with the other Ascend operators; -there is no separate resolver build or library-loading step. - -An unsupported capacity or missing extension raises an error before execution; -there is no Triton resolver fallback. Other build targets retain their existing -operators without registering this resolver. +Use the vLLM revision pinned in `.github/vllm-main-verified.commit` and install +vLLM Ascend in the matching Torch, torch-npu and CANN environment. Ascend +HiSparse needs vLLM's pluggable HiSparse backend interface. The vLLM 0.33.0 +release supports HiSparse on NVIDIA CUDA only. ## Enable HiSparse diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index c4d9a395ed49..953fda74d4d1 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -13,6 +13,7 @@ @pytest.fixture def _require_npu_sfa(): + pytest.importorskip("vllm.v1.hisparse.backend") from vllm_ascend.utils import bootstrap_custom_op_env # Match platform initialization: publish the OPP path before device queries @@ -222,6 +223,7 @@ def forward(): @pytest.fixture def _require_a5_preprocess(): + pytest.importorskip("vllm.v1.hisparse.backend") if not torch.npu.is_available() or not torch.npu.get_device_name(0).startswith("Ascend950"): pytest.skip("HiSparse integration currently targets Ascend 950") torch.npu.set_device(0) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index def07a97d901..60bfefc7479a 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -8,7 +8,6 @@ import torch import torch_npu # noqa: F401 -from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, ResolverWorkspace, @@ -17,6 +16,9 @@ ) from vllm_ascend.hisparse.host_pool import MappedHostTensor, allocate_mapped_host +pytest.importorskip("vllm.v1.hisparse.backend") +from vllm_ascend.hisparse.backend import AscendHiSparseBackend # noqa: E402 + @pytest.fixture(autouse=True) def _require_mapped_host_device(): @@ -381,17 +383,21 @@ def attention_reference(physical): mapping.close() -@pytest.mark.parametrize("query_rows,hot_size", [(2, 4096), (4, 10240)]) +@pytest.mark.parametrize("query_rows,hot_size", [(2, 4096), (4, 10240), (3, 4096)]) @pytest.mark.parametrize("cross_stream", [False, True]) def test_topk_2048_speculative_union_graph(query_rows, hot_size, cross_stream): topk, width, block = 2048, 576, 128 - union = query_rows * topk + positions = query_rows * topk + union = min(positions, hot_size) + # The 3 x 2048 case has more positions than the 4096-row mask/scan + # capacity. Its duplicates cross tiles of the iterative miss scan. blocks = hot_size // block + 1 host = allocate_mapped_host((hot_size + block, width), torch.bfloat16) mapping = host._hisparse_registration host.copy_((torch.arange(host.size(0)) % 97).to(torch.bfloat16)[:, None].expand_as(host)) hot = torch.zeros((blocks, block, width), dtype=torch.bfloat16, device="npu") - indices = torch.arange(union, dtype=torch.int32, device="npu").view(query_rows, topk) + logical = torch.arange(positions, dtype=torch.int32).view(query_rows, topk) % union + indices = logical.npu() requests = torch.zeros(query_rows, dtype=torch.int32, device="npu") states = torch.zeros(1, dtype=torch.int32, device="npu") table = torch.arange(1, blocks, dtype=torch.int32, device="npu")[None, :] @@ -452,17 +458,21 @@ def forward(): forward() torch.npu.synchronize() - expected = torch.arange(block, block + union, dtype=torch.int32).view(query_rows, topk) + expected = logical + block + expected_swaps = torch.tensor( + [min(topk, max(0, union - row * topk)) for row in range(query_rows)], dtype=torch.int32 + ) + cold_hits = positions - union torch.testing.assert_close(out.cpu(), expected) - torch.testing.assert_close(sc.cpu(), torch.full((query_rows,), topk, dtype=torch.int32)) - torch.testing.assert_close(stats.cpu(), torch.tensor([0, union], dtype=torch.int64)) + torch.testing.assert_close(sc.cpu(), expected_swaps) + torch.testing.assert_close(stats.cpu(), torch.tensor([cold_hits, union], dtype=torch.int64)) torch.testing.assert_close(hot[1:].cpu().reshape(hot_size, width)[:union], host[block : block + union]) expected_lru = torch.cat((torch.arange(union, hot_size), torch.arange(union))).to(torch.int16)[None, :] torch.testing.assert_close(lru.cpu(), expected_lru) forward() torch.npu.synchronize() torch.testing.assert_close(sc.cpu(), torch.zeros(query_rows, dtype=torch.int32)) - torch.testing.assert_close(stats.cpu(), torch.tensor([union, union], dtype=torch.int64)) + torch.testing.assert_close(stats.cpu(), torch.tensor([cold_hits + positions, union], dtype=torch.int64)) capture_stream = torch.npu.Stream() capture_stream.wait_stream(torch.npu.current_stream()) with torch.npu.stream(capture_stream): @@ -477,13 +487,13 @@ def forward(): graph.replay() torch.npu.synchronize() torch.testing.assert_close(out.cpu(), expected) - torch.testing.assert_close(sc.cpu(), torch.full((query_rows,), topk, dtype=torch.int32)) - torch.testing.assert_close(stats.cpu(), torch.tensor([0, union], dtype=torch.int64)) + torch.testing.assert_close(sc.cpu(), expected_swaps) + torch.testing.assert_close(stats.cpu(), torch.tensor([cold_hits, union], dtype=torch.int64)) for _ in range(10): graph.replay() torch.npu.synchronize() torch.testing.assert_close(sc.cpu(), torch.zeros(query_rows, dtype=torch.int32)) - torch.testing.assert_close(stats.cpu(), torch.tensor([10 * union, union], dtype=torch.int64)) + torch.testing.assert_close(stats.cpu(), torch.tensor([cold_hits + 10 * positions, union], dtype=torch.int64)) torch.testing.assert_close(hot[1:].cpu().reshape(hot_size, width)[:union], host[block : block + union]) torch.testing.assert_close(lru.cpu(), expected_lru) graph.reset() diff --git a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py index a0cf49482372..17bf118dc935 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py @@ -10,13 +10,15 @@ from vllm.v1.hisparse.runtime import HiSparseRuntime, PagedCacheView from vllm_ascend.attention.sfa_hisparse import build_hisparse_prefill_staging_plan -from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.cache_ops import ( HiSparseNPUOps, scatter_packed_kv, ) from vllm_ascend.hisparse.host_pool import allocate_mapped_host +pytest.importorskip("vllm.v1.hisparse.backend") +from vllm_ascend.hisparse.backend import AscendHiSparseBackend # noqa: E402 + @pytest.mark.parametrize("resident_block_size", [128, 64, 32]) @pytest.mark.parametrize("host_pages", [1, 2]) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 4d77da1254a4..8d4eefa5bc06 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -20,7 +20,6 @@ from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.hisparse import cache_ops as hisparse_ops from vllm_ascend.hisparse import host_pool -from vllm_ascend.hisparse.backend import AscendHiSparseBackend from vllm_ascend.hisparse.binding import ( _bind_indexer_views, align_hisparse_indexer_specs, @@ -33,6 +32,9 @@ ) from vllm_ascend.ops.triton import sfa_indexer_store +pytest.importorskip("vllm.v1.hisparse.backend") +from vllm_ascend.hisparse.backend import AscendHiSparseBackend # noqa: E402 + @pytest.fixture(params=[False, True], ids=["standalone", "profiled"]) def prepare_workspaces(request): From 8cf106dfce5d9034c1eddcaa27de017ac67b8fad Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:16:21 +0800 Subject: [PATCH 110/120] fix(hisparse): use the shared vLLM logger Signed-off-by: zhenwenqi2024 --- vllm_ascend/hisparse/host_mirror.py | 3 +-- vllm_ascend/hisparse/recycled_invalidate.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm_ascend/hisparse/host_mirror.py b/vllm_ascend/hisparse/host_mirror.py index 4ea45ca18a09..68bf5062df95 100644 --- a/vllm_ascend/hisparse/host_mirror.py +++ b/vllm_ascend/hisparse/host_mirror.py @@ -238,7 +238,7 @@ def overlap_target_host_mirror(runner): from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import HiSparseConnector from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import HiSparseConnectorWorker from vllm.distributed.parallel_state import get_tensor_model_parallel_rank - from vllm.logger import init_logger + from vllm.logger import logger from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector from vllm_ascend.hisparse.backend import AscendHiSparseBackend @@ -270,7 +270,6 @@ def overlap_target_host_mirror(runner): or output_stream == compute_stream ): return connector - logger = init_logger(__name__) rank = get_tensor_model_parallel_rank() logger.info("HiSparse direct Target Host mirror overlap enabled: rank=%d", rank) return HostMirrorConnector(connector, worker, output_stream, sources, logger, rank) diff --git a/vllm_ascend/hisparse/recycled_invalidate.py b/vllm_ascend/hisparse/recycled_invalidate.py index 69a29532ac37..25b93b9abd5b 100644 --- a/vllm_ascend/hisparse/recycled_invalidate.py +++ b/vllm_ascend/hisparse/recycled_invalidate.py @@ -6,13 +6,12 @@ import torch from vllm.distributed import get_tensor_model_parallel_rank -from vllm.logger import init_logger +from vllm.logger import logger from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import async_tensor_h2d from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type -logger = init_logger(__name__) _BLOCK_BUCKET = 8 _TILE = 256 _INT32_MAX = 2**31 - 1 From 9cd801d2acb0d761c0df18308a997475f68e7457 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:22:02 +0800 Subject: [PATCH 111/120] fix(hisparse): follow the current indexer and metadata contracts Signed-off-by: zhenwenqi2024 --- .../test_hisparse_foundation_cpu.py | 26 ++++++++++++++++++- vllm_ascend/attention/indexer.py | 10 +++++-- vllm_ascend/attention/sfa_hisparse.py | 13 ++++++---- vllm_ascend/hisparse/backend.py | 2 +- vllm_ascend/hisparse/cache_ops.py | 5 ++-- vllm_ascend/hisparse/host_pool.py | 8 ++++-- 6 files changed, 50 insertions(+), 14 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 8d4eefa5bc06..0c5ba5155265 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -423,6 +423,29 @@ def test_c8_compaction_preserves_duplicate_order_and_empty_rows(): torch.testing.assert_close(compact_c8_indices(indices), expected, rtol=0, atol=0) +def test_bf16_hisparse_indexer_passes_through_page_scatter_indices(monkeypatch): + cache = torch.empty((2, 4, 128), dtype=torch.bfloat16) + keys = torch.zeros((2, 128), dtype=torch.bfloat16) + slots = torch.tensor([2, 4], dtype=torch.int32) + scatter_indices = torch.tensor([[0, 2], [1, 0]], dtype=torch.int32) + writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) + torch.nn.Module.__init__(writer) + writer.k_cache = NS(kv_cache=(cache,)) + writer.enable_sparse_li_c8 = False + writer.enable_sparse_li_c4 = False + calls = [] + + def write_cache_tensor(cache_arg, updates, slots_arg, indices_arg): + calls.append(indices_arg) + assert cache_arg is cache + assert slots_arg is slots + torch.testing.assert_close(updates, keys) + + monkeypatch.setattr(writer, "_write_cache_tensor", write_cache_tensor) + writer.write_cache(keys, None, slots, scatter_indices) + assert len(calls) == 1 and calls[0] is scatter_indices + + @pytest.mark.parametrize("padding", [0, 256]) def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(monkeypatch, indexer_spec, padding): spec = replace( @@ -479,5 +502,6 @@ def store(quant, scales, slots, cache, scale_cache, *, block_size, key_block_str expected[start : start + 128] = values_raw[i] start = offset + page * stride + key_bytes + row * 4 expected[start : start + 4] = scale_bytes[i] - writer.write_cache(values, scales, slots) + scatter_indices = torch.stack((slots // spec.block_size, slots % spec.block_size), dim=-1) + writer.write_cache(values, scales, slots, scatter_indices) torch.testing.assert_close(raw, expected, rtol=0, atol=0) diff --git a/vllm_ascend/attention/indexer.py b/vllm_ascend/attention/indexer.py index c9f4aec976a6..0a9fbaade14b 100644 --- a/vllm_ascend/attention/indexer.py +++ b/vllm_ascend/attention/indexer.py @@ -922,9 +922,15 @@ def get_name() -> str: def get_builder_cls(): return AscendHiSparseIndexerMetadataBuilder - def write_cache(self, k_li, k_li_scale, slot_mapping) -> None: + def write_cache( + self, + k_li: torch.Tensor, + k_li_scale: torch.Tensor | None, + slot_mapping: torch.Tensor, + cache_scatter_indices: torch.Tensor | None = None, + ) -> None: if not self.enable_sparse_li_c8: - return super().write_cache(k_li, k_li_scale, slot_mapping) + return super().write_cache(k_li, k_li_scale, slot_mapping, cache_scatter_indices) if k_li_scale is None: raise ValueError("HiSparse Indexer C8 requires FP32 scales") from vllm_ascend.hisparse.binding import write_hisparse_indexer_c8 diff --git a/vllm_ascend/attention/sfa_hisparse.py b/vllm_ascend/attention/sfa_hisparse.py index 076d0169e957..011a2c060415 100644 --- a/vllm_ascend/attention/sfa_hisparse.py +++ b/vllm_ascend/attention/sfa_hisparse.py @@ -4,6 +4,7 @@ import math from dataclasses import dataclass +from typing import cast import torch from torch.overrides import TorchFunctionMode @@ -97,7 +98,7 @@ class AscendHiSparseMetadata(AscendSFAMetadata): class AscendHiSparseMetadataBuilder(AscendSFAMetadataBuilder): def build_for_cudagraph_capture(self, common_attn_metadata, **kwargs): - metadata = super().build_for_cudagraph_capture(common_attn_metadata, **kwargs) + metadata = cast(AscendHiSparseMetadata, super().build_for_cudagraph_capture(common_attn_metadata, **kwargs)) metadata.hisparse_for_capture = True return metadata @@ -115,7 +116,7 @@ def _build( draft_index=None, pcp_context=None, ): - metadata = super()._build(common_attn_metadata, draft_index, pcp_context) + metadata = cast(AscendHiSparseMetadata, super()._build(common_attn_metadata, draft_index, pcp_context)) metadata.num_reqs = common_attn_metadata.num_reqs if draft_index not in self._request_row_buffers: self._request_row_buffers[draft_index] = torch.empty( @@ -313,14 +314,14 @@ def forward(self, layer_name, hidden_states, kv_cache, attn_metadata, output=Non self.index_group.prepare_for_batch(self.index_group_index, attn_metadata) return super().forward(layer_name, hidden_states, kv_cache, attn_metadata, output) - def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): + def _exec_kv_local(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): from vllm_ascend.hisparse.cache_ops import scatter_packed_kv assert self.kv_a_layernorm is not None if self.enable_sparse_sfa_c8: from vllm_ascend.hisparse.c8 import pack_c8_components - rope, latent, scales = super().exec_kv(kv_no_split, cos, sin, kv_cache, slots, attn_metadata) + rope, latent, scales = super()._exec_kv_local(kv_no_split, cos, sin, kv_cache, slots, attn_metadata) rows = latent.numel() // self.kv_lora_rank packed = pack_c8_components(latent.contiguous(), rope.contiguous(), scales.contiguous()) cache = self.hisparse_cache @@ -342,7 +343,9 @@ def exec_kv(self, kv_no_split, cos, sin, kv_cache, slots, attn_metadata): # Match ordinary SFA's fused RMSNorm/RoPE rounding. Separate RMSNorm # can differ by a BF16 ULP near a rounding midpoint. Local slots # isolate padding from the real resident/hot destination below. - super().exec_kv(values, cos, sin, (latent_cache, rope_cache), local_slots[:num_tokens], attn_metadata) + super()._exec_kv_local( + values, cos, sin, (latent_cache, rope_cache), local_slots[:num_tokens], attn_metadata + ) latent = latent_cache.view(-1, self.kv_lora_rank)[:num_tokens] rope = rope_cache.view(-1, self.qk_rope_head_dim)[:num_tokens] else: diff --git a/vllm_ascend/hisparse/backend.py b/vllm_ascend/hisparse/backend.py index 3f37f45ace18..859cafa83bcd 100644 --- a/vllm_ascend/hisparse/backend.py +++ b/vllm_ascend/hisparse/backend.py @@ -43,7 +43,7 @@ def init(worker): ) worker._ascend_recycled_invalidator = RecycledGroupInvalidator.bind(worker) - init._ascend_row_dma = True + init.__dict__["_ascend_row_dma"] = True HiSparseConnectorWorker._init_dma = init original_invalidate = HiSparseConnectorWorker.invalidate_blocks diff --git a/vllm_ascend/hisparse/cache_ops.py b/vllm_ascend/hisparse/cache_ops.py index ee859e5264e8..0dc0917224cb 100644 --- a/vllm_ascend/hisparse/cache_ops.py +++ b/vllm_ascend/hisparse/cache_ops.py @@ -247,16 +247,15 @@ def hisparse_resolve_residency( != 1 ): raise ValueError("All three compact HiSparse swap outputs are required together") - self._resolve(locals()) + self._resolve(workspace, locals()) - def _resolve(self, arguments): + def _resolve(self, workspace: ResolverWorkspace, arguments) -> None: """Dispatch live tensors to the native request-union resolver. During replay, external tensors remain owned by the runtime/workspace. Callers must destroy graphs before releasing or rebinding those owners. """ rows, topk = arguments["global_indices"].shape - workspace = self.workspace resolved = arguments["resolved_global_indices"] if resolved is None: resolved = workspace.first_refs[: rows * topk].view(rows, topk) diff --git a/vllm_ascend/hisparse/host_pool.py b/vllm_ascend/hisparse/host_pool.py index fba0a9132f72..3f31e02f85cc 100644 --- a/vllm_ascend/hisparse/host_pool.py +++ b/vllm_ascend/hisparse/host_pool.py @@ -66,7 +66,11 @@ def __init__(self, tensor: torch.Tensor) -> None: ) if status: raise RuntimeError(f"aclrtHostRegister for HiSparse failed: {status}") - self.device_ptr = int(mapped.value) + device_ptr = mapped.value + if device_ptr is None: + self._lib.aclrtHostUnregister(self.host_address) + raise RuntimeError("aclrtHostRegister returned a null device pointer for HiSparse") + self.device_ptr = device_ptr self.closed = False def address_of(self, tensor: torch.Tensor) -> int: @@ -127,5 +131,5 @@ def allocate(size): return registered[:size], registered return original(size) - allocate._ascend_pinned_private_pool = True + allocate.__dict__["_ascend_pinned_private_pool"] = True runtime.allocate_pinned_host_pool = allocate From 08aacc702cecfe3d6a2b40caf3c13f160786a896 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:26:37 +0800 Subject: [PATCH 112/120] fix(hisparse): type the cache layout and verification paths Signed-off-by: zhenwenqi2024 --- .../pull_request/one_card/test_hisparse_attention.py | 2 +- tests/e2e/pull_request/one_card/test_hisparse_ops.py | 11 +++++++---- .../kv_transfer/test_hisparse_foundation_cpu.py | 1 + vllm_ascend/worker/v2/attn_utils.py | 2 +- 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 953fda74d4d1..6c108cbedb47 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -296,7 +296,7 @@ def forward(): rope, latent, scales = custom_kv_rmsnorm_rope( values.reshape(rows, 1, 1, 576), gamma, cos, sin, 512, 64, tile_size=128 ) - components = (latent, rope, scales) + components: tuple[torch.Tensor, ...] = (latent, rope, scales) else: # Reuse native SFA as the oracle; test only the new destination, # strides, row mapping and padding, not RMSNorm arithmetic. diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 60bfefc7479a..89d3f6ee2966 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -189,7 +189,7 @@ def reference(topk, requests, states, source, resident, hot_table, dgi, lru, blo hits = [slot for slot in order if int(dgi[state, slot]) in union] evict = [slot for slot in order if slot not in hits] mapping = {int(dgi[state, slot]): slot for slot in hits} - misses = [] + misses: list[tuple[int, int]] = [] for host in union: if host not in mapping: slot = evict[len(misses)] @@ -622,7 +622,10 @@ def forward(): sh.fill_(-9) sd.fill_(-9) stats.zero_() - graph.replay() if capture else forward() + if graph is None: + forward() + else: + graph.replay() torch.npu.synchronize() for label, actual, wanted in zip( ("indices", "counts", "swap_host", "swap_device", "swap_counts", "dgi", "lru", "stats"), @@ -1020,13 +1023,13 @@ def test_direct_target_mirror_survives_legal_draft_writes(graph_producer, shared owners.append(owner) handles.append(handle) for leader, followers in ((owners[0], [owners[1]]), (owners[2], [])): - group = SimpleNamespace( + index_group = SimpleNamespace( leader=leader, followers=followers, device_global_indices=torch.full((1, 16), -1, dtype=torch.int32, device=source.device), ) for owner in [leader, *followers]: - owner.index_group = group + owner.index_group = index_group worker = HiSparseConnectorWorker.__new__(HiSparseConnectorWorker) worker.backend, worker._initialized = backend, False worker.vllm_config = SimpleNamespace( diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 0c5ba5155265..0e7c41b6bce6 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -304,6 +304,7 @@ def test_indexer_page_alignment_preserves_input_specs(indexer_spec, c8): assert indexer_spec.page_size_padded is None for name in ("resident", "hot"): assert specs[name].block_stride_alignment == 128 + assert packed_row_bytes is not None assert result[name].block_stride_alignment == math.lcm(128, packed_row_bytes) else: assert result["indexer"].get_block_stride_alignment() == math.lcm(3, indexer_spec.page_size_bytes) diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 8adb3addd9ff..e7c6e27f047d 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -286,7 +286,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if isinstance(attn_module, MLAAttention) and vllm_config.attention_config.hisparse_config is not None: # Keep Core's row-packed source spec and actual index-group ownership. - source_layout = {} + source_layout: dict[str, Any] = {} if getattr(attn_module.impl, "enable_sparse_sfa_c8", False): from vllm_ascend.hisparse.c8 import C8_ROW_BYTES From 58f957854fbab764af8c1a38f84ae70e785a27ba Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:31:58 +0800 Subject: [PATCH 113/120] test(hisparse): initialize SFA flags in device fixtures Signed-off-by: zhenwenqi2024 --- tests/e2e/pull_request/one_card/test_hisparse_attention.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_attention.py b/tests/e2e/pull_request/one_card/test_hisparse_attention.py index 6c108cbedb47..db3ad404f7d7 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_attention.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_attention.py @@ -4,6 +4,7 @@ import math from functools import partial +from importlib import import_module from types import SimpleNamespace import pytest @@ -27,7 +28,7 @@ def _require_npu_sfa(): pytest.skip("The native SFA operation is not registered for Ascend 310P") torch.npu.set_device(0) # Load the required SFA extension without changing custom-op selection. - import vllm_ascend.vllm_ascend_C # noqa: F401 + import_module("vllm_ascend.vllm_ascend_C") @pytest.mark.parametrize("capture", [False, True]) @@ -258,6 +259,8 @@ def test_preprocess_writes_hisparse_mapping_once(cache_dtype, capture): ) impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) impl.enable_sparse_sfa_c8 = c8 + impl.pcp_enabled = False + impl.use_rope = True impl.enable_sparse_sfa_turboquant = False impl.c8_cache_dtype = torch.float8_e4m3fn impl.sfa_qsfa_tile_size = 128 @@ -343,6 +346,7 @@ def test_c8_prolog_keeps_parent_outputs_and_hisparse_destination(rows, capture, device, block, width = "npu", 128, 656 impl = AscendHiSparseImpl.__new__(AscendHiSparseImpl) impl._quant_type = AscendW8A8MXFP8DynamicLinearMethod + impl.use_rope = True impl.enable_sparse_sfa_c8 = True impl.c8_cache_dtype = torch.float8_e4m3fn impl.q_lora_rank, impl.kv_lora_rank, impl.qk_rope_head_dim = 2048, 512, 64 From b1920e75a8bc37437b31d10d15ecc7b13b7d2035 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:35:18 +0800 Subject: [PATCH 114/120] fix(hisparse): bind C8 indexer views from the full Core page Signed-off-by: zhenwenqi2024 --- .../kv_transfer/test_hisparse_foundation_cpu.py | 16 +++++++++++++--- vllm_ascend/hisparse/binding.py | 9 ++++++--- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py index 0e7c41b6bce6..8c8c6d30c379 100644 --- a/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py +++ b/tests/ut/distributed/kv_transfer/test_hisparse_foundation_cpu.py @@ -8,7 +8,15 @@ import pytest import torch -from vllm.v1.kv_cache_interface import HiSparseHotSpec, HiSparseResidentSpec, MLAAttentionSpec, SparseCacheRole +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheTensor, + MLAAttentionSpec, + SparseCacheRole, + create_kv_cache_views, +) +from vllm.v1.kv_cache_layout import KVCacheLayout from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.attention.sfa_hisparse import ( @@ -462,9 +470,11 @@ def test_c8_indexer_views_and_writes_preserve_neighboring_slab_bytes(monkeypatch key_bytes = 128 * 128 stride = spec.page_size_bytes + padding raw = torch.full((offset + pages * stride + 256,), 165, dtype=torch.uint8) - cache = raw.view(torch.float8_e4m3fn).as_strided( - (pages, 1, 128, 128), (stride, key_bytes, 128, 1), storage_offset=offset + tensor_config = KVCacheTensor( + size=raw.numel(), layers=["indexer"], layer_stride=spec.page_size_bytes, block_stride=stride, offset=offset ) + (cache,) = create_kv_cache_views(raw.view(torch.int8), spec, pages, KVCacheLayout.BLHNC, tensor_config) + assert cache.shape == (pages, 1, 128, 132) layer = NS() caches = {"indexer": cache} _bind_indexer_views(caches, {"indexer": spec}, {"indexer": layer}) diff --git a/vllm_ascend/hisparse/binding.py b/vllm_ascend/hisparse/binding.py index f9b6ccb5a3f7..3deae1b3cc82 100644 --- a/vllm_ascend/hisparse/binding.py +++ b/vllm_ascend/hisparse/binding.py @@ -89,15 +89,18 @@ def _validate_indexer_spec(spec: AscendSFAIndexerCacheSpec) -> None: def _quantized_indexer_views(cache: torch.Tensor, spec: AscendSFAIndexerCacheSpec): """Alias the current Core allocation; retain no additional storage owner.""" - pages, heads, block, width = cache.shape + pages, heads, block, row_bytes = cache.shape stride = cache.stride(0) + width = spec.head_size key_bytes = block * heads * width scale_bytes = block * heads * torch.float32.itemsize offset = cache.storage_offset() + # Core's view covers all bytes reserved for each token. The Indexer + # kernels use a page with contiguous K followed by contiguous scales. if ( cache.dtype != torch.float8_e4m3fn - or (heads, block, width) != (1, spec.block_size, spec.head_size) - or cache.stride()[1:] != (block * width, width, 1) + or (heads, block, row_bytes) != (1, spec.block_size, spec.state_content_size_bytes) + or cache.stride()[2:] != (row_bytes, 1) or pages <= 0 or stride < key_bytes + scale_bytes or stride % width From 4c813f2b174ee9958e7d85da95a7e951fb85266e Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:39:53 +0800 Subject: [PATCH 115/120] test(hisparse): exercise the Core page view on NPU Signed-off-by: zhenwenqi2024 --- tests/e2e/pull_request/one_card/test_hisparse_ops.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/e2e/pull_request/one_card/test_hisparse_ops.py b/tests/e2e/pull_request/one_card/test_hisparse_ops.py index 89d3f6ee2966..1ebe6579c20b 100644 --- a/tests/e2e/pull_request/one_card/test_hisparse_ops.py +++ b/tests/e2e/pull_request/one_card/test_hisparse_ops.py @@ -857,6 +857,9 @@ def count_sort(*args, **kwargs): def test_page_strided_indexer_write_preserves_bits_and_padding(capture): from types import SimpleNamespace + from vllm.v1.kv_cache_interface import KVCacheTensor, create_kv_cache_views + from vllm.v1.kv_cache_layout import KVCacheLayout + from vllm_ascend.attention.indexer import AscendHiSparseIndexerBackend from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.hisparse.binding import _quantized_indexer_views @@ -864,9 +867,6 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): pages, block, width, page_stride = 3, 128, 128, 419840 size = (pages - 1) * page_stride + block * width + block * 4 raw = torch.full((size,), 90, dtype=torch.uint8, device="npu") - packed = torch.as_strided( - raw.view(torch.float8_e4m3fn), (pages, 1, block, width), (page_stride, block * width, width, 1) - ) spec = AscendSFAIndexerCacheSpec( block_size=block, num_kv_heads=1, @@ -879,6 +879,10 @@ def test_page_strided_indexer_write_preserves_bits_and_padding(capture): cache_sparse_li_c4=False, sfa_dcp_replicated_indexer_size=1, ) + tensor_config = KVCacheTensor( + size=size, layers=["indexer"], layer_stride=spec.page_size_bytes, block_stride=page_stride + ) + (packed,) = create_kv_cache_views(raw.view(torch.int8), spec, pages, KVCacheLayout.BLNHC, tensor_config) key, scale = _quantized_indexer_views(packed, spec) writer = AscendHiSparseIndexerBackend.__new__(AscendHiSparseIndexerBackend) torch.nn.Module.__init__(writer) From 8c2cc0156ac59ab0441c454e790ee77e3f81bd81 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 01:53:07 +0800 Subject: [PATCH 116/120] fix(ci): reuse sources that already include the base snapshot Signed-off-by: zhenwenqi2024 --- .github/workflows/_build_csrc_cache.yaml | 2 +- .github/workflows/_selected_tests.yaml | 2 +- .github/workflows/_selected_tests_upstream.yaml | 2 +- .github/workflows/pr_test.yaml | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/_build_csrc_cache.yaml b/.github/workflows/_build_csrc_cache.yaml index 70a5473069e1..5411d07c67bb 100644 --- a/.github/workflows/_build_csrc_cache.yaml +++ b/.github/workflows/_build_csrc_cache.yaml @@ -152,7 +152,7 @@ jobs: fi git fetch "$BASE_REMOTE" "$BASE_REF" git cat-file -e "${BASE_SHA}^{commit}" - git rebase "$BASE_SHA" + git merge-base --is-ancestor "$BASE_SHA" HEAD || git rebase "$BASE_SHA" - name: Update submodules run: | diff --git a/.github/workflows/_selected_tests.yaml b/.github/workflows/_selected_tests.yaml index 0be07dfe9b55..8fef7b70ec70 100644 --- a/.github/workflows/_selected_tests.yaml +++ b/.github/workflows/_selected_tests.yaml @@ -161,7 +161,7 @@ jobs: BASE_SHA=$(git rev-parse FETCH_HEAD) fi git cat-file -e "${BASE_SHA}^{commit}" - git rebase "$BASE_SHA" + git merge-base --is-ancestor "$BASE_SHA" HEAD || git rebase "$BASE_SHA" git submodule sync --recursive git submodule update --init --recursive diff --git a/.github/workflows/_selected_tests_upstream.yaml b/.github/workflows/_selected_tests_upstream.yaml index 80db12668245..aa08ec89d253 100644 --- a/.github/workflows/_selected_tests_upstream.yaml +++ b/.github/workflows/_selected_tests_upstream.yaml @@ -146,7 +146,7 @@ jobs: BASE_SHA=$(git rev-parse FETCH_HEAD) fi git cat-file -e "${BASE_SHA}^{commit}" - git rebase "$BASE_SHA" + git merge-base --is-ancestor "$BASE_SHA" HEAD || git rebase "$BASE_SHA" git submodule sync --recursive git submodule update --init --recursive diff --git a/.github/workflows/pr_test.yaml b/.github/workflows/pr_test.yaml index 24b82e0a3cf4..5119f27accbb 100644 --- a/.github/workflows/pr_test.yaml +++ b/.github/workflows/pr_test.yaml @@ -531,7 +531,7 @@ jobs: exit 1 fi - git rebase "$BASE_SHA" + git merge-base --is-ancestor "$BASE_SHA" HEAD || git rebase "$BASE_SHA" echo "base_sha=$BASE_SHA" >> "$GITHUB_OUTPUT" - name: Validate test coverage config From b6a87beeac13d20052741ac4975045e309bc6c1b Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 02:01:12 +0800 Subject: [PATCH 117/120] test(hisparse): register device suite durations for CI selection Signed-off-by: zhenwenqi2024 --- .github/workflows/scripts/estimated_times.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/scripts/estimated_times.yaml b/.github/workflows/scripts/estimated_times.yaml index 57c3e97bbb31..8d8afe4efb97 100644 --- a/.github/workflows/scripts/estimated_times.yaml +++ b/.github/workflows/scripts/estimated_times.yaml @@ -15,6 +15,9 @@ estimated_times: tests/e2e/pull_request/one_card/test_kv_pool_accuracy.py: 410 tests/e2e/pull_request/one_card/test_flash_attn_precision.py: 30 tests/e2e/pull_request/one_card/test_swiglu_mx_quant.py: 30 + tests/e2e/pull_request/one_card/test_hisparse_attention.py: 30 + tests/e2e/pull_request/one_card/test_hisparse_ops.py: 30 + tests/e2e/pull_request/one_card/test_hisparse_prefill_staging.py: 30 tests/e2e/pull_request/four_card/model_runner_v2/test_kvpp_pp.py: 300 tests/e2e/pull_request/four_card/test_kvpp_pd.py: 230 tests/e2e/pull_request/two_card/test_kvpp.py: 290 From 1d3d79180f2167201b9e1fe583577195a80a37a1 Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 02:12:26 +0800 Subject: [PATCH 118/120] test(hisparse): use real attention defaults in the remaining CPU fixtures Signed-off-by: zhenwenqi2024 --- tests/ut/attention/test_flash_attn.py | 5 +++-- tests/ut/attention/test_sfa_cp.py | 6 +++++- tests/ut/attention/test_sfa_v1.py | 3 ++- tests/ut/models/test_glm5next_cache_config.py | 2 ++ tests/ut/worker/test_full_sliding_kv_cache.py | 2 ++ tests/ut/worker/test_model_runner_v2_mamba.py | 5 +++-- tests/ut/worker/test_worker_v2.py | 3 ++- 7 files changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/ut/attention/test_flash_attn.py b/tests/ut/attention/test_flash_attn.py index dd1648409734..e9f6c9eaf494 100644 --- a/tests/ut/attention/test_flash_attn.py +++ b/tests/ut/attention/test_flash_attn.py @@ -7,6 +7,7 @@ import pytest import torch +from vllm.config import AttentionConfig from vllm.model_executor.layers.attention.attention import Attention from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.selector import AttentionSelectorConfig @@ -127,7 +128,7 @@ def test_flash_attn_192_cache_spec_integration(monkeypatch, backend): config = SimpleNamespace( model_config=SimpleNamespace(dtype=torch.bfloat16), parallel_config=SimpleNamespace(decode_context_parallel_size=1), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(), cache_config=SimpleNamespace(cache_dtype="fp8", block_size=128, skip_page_size_padded=None), ) monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_args: {"draft.attn": layer}) @@ -478,7 +479,7 @@ def test_only_undersized_draft_swa_pages_drop_inherited_padding( config = SimpleNamespace( model_config=SimpleNamespace(dtype=torch.bfloat16), parallel_config=SimpleNamespace(decode_context_parallel_size=1), - attention_config=SimpleNamespace(indexer_kv_dtype="auto"), + attention_config=AttentionConfig(), cache_config=SimpleNamespace(cache_dtype="fp8", block_size=128, skip_page_size_padded=padding), ) original = layer.get_kv_cache_spec(config) diff --git a/tests/ut/attention/test_sfa_cp.py b/tests/ut/attention/test_sfa_cp.py index 5c5ced68da52..f9533ae7c18f 100644 --- a/tests/ut/attention/test_sfa_cp.py +++ b/tests/ut/attention/test_sfa_cp.py @@ -10,6 +10,7 @@ import pytest import torch import vllm.distributed as vllm_distributed +from vllm.config import AttentionConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.ops.pcp import _gather_prefill_cache_inputs @@ -504,7 +505,10 @@ def test_sfa_dcp_extends_v1_backend() -> None: @pytest.mark.parametrize("pcp_size,dcp_enabled", [(1, False), (2, False), (1, True), (2, True)]) def test_sfa_backend_selects_dcp_subclass_and_layout_specific_builder(pcp_size, dcp_enabled): - config = SimpleNamespace(parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size)) + config = SimpleNamespace( + attention_config=AttentionConfig(), + parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size), + ) with ( patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config", return_value=config), patch("vllm_ascend.attention.sfa_v1.get_ascend_config") as ascend_config, diff --git a/tests/ut/attention/test_sfa_v1.py b/tests/ut/attention/test_sfa_v1.py index c53cd8c2a951..3015417bab6d 100644 --- a/tests/ut/attention/test_sfa_v1.py +++ b/tests/ut/attention/test_sfa_v1.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch import torch -from vllm.config import set_current_vllm_config +from vllm.config import AttentionConfig, set_current_vllm_config from vllm.distributed.parallel_state import GroupCoordinator from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod @@ -45,6 +45,7 @@ class TestAscendSFABackend(TestBase): def setUp(self): self.mock_config = MagicMock() + self.mock_config.attention_config = AttentionConfig() mock_parallel_config = MagicMock() mock_parallel_config.prefill_context_parallel_size = 1 mock_parallel_config.decode_context_parallel_size = 1 diff --git a/tests/ut/models/test_glm5next_cache_config.py b/tests/ut/models/test_glm5next_cache_config.py index 5276ed4b4bf0..9c517abbe1cc 100644 --- a/tests/ut/models/test_glm5next_cache_config.py +++ b/tests/ut/models/test_glm5next_cache_config.py @@ -7,6 +7,7 @@ import pytest import torch +from vllm.config import AttentionConfig from vllm.v1.core.single_type_kv_cache_manager import ( register_all_kvcache_specs, ) @@ -70,6 +71,7 @@ def _ratio_kwargs(ratio: int) -> dict[str, int]: def make_config(): return SimpleNamespace( + attention_config=AttentionConfig(), model_config=SimpleNamespace(max_model_len=2048, hf_config=SimpleNamespace(model_type="glm_moe_dsa")), parallel_config=SimpleNamespace( decode_context_parallel_size=1, diff --git a/tests/ut/worker/test_full_sliding_kv_cache.py b/tests/ut/worker/test_full_sliding_kv_cache.py index 7fa36a0849b2..228860343d06 100644 --- a/tests/ut/worker/test_full_sliding_kv_cache.py +++ b/tests/ut/worker/test_full_sliding_kv_cache.py @@ -8,6 +8,7 @@ import pytest import torch +from vllm.config import AttentionConfig from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -29,6 +30,7 @@ def _config(): layout = SimpleNamespace(is_layer_compact=True, is_block_compact=True) return SimpleNamespace( + attention_config=AttentionConfig(), model_config=SimpleNamespace(hf_text_config=SimpleNamespace(model_type="gemma4_text"), use_mla=False), cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: layout, cache_dtype="auto"), parallel_config=SimpleNamespace(decode_context_parallel_size=1, prefill_context_parallel_size=1), diff --git a/tests/ut/worker/test_model_runner_v2_mamba.py b/tests/ut/worker/test_model_runner_v2_mamba.py index f5119b7140db..fc0e9f83964c 100644 --- a/tests/ut/worker/test_model_runner_v2_mamba.py +++ b/tests/ut/worker/test_model_runner_v2_mamba.py @@ -9,6 +9,7 @@ import pytest import torch import vllm.v1.worker.gpu.cp_utils as _cp_utils +from vllm.config import AttentionConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import ( @@ -49,7 +50,7 @@ def _mock_vllm_config(): # Config for get_kv_cache_spec: attn_utils reads attention_config.indexer_kv_dtype. config = MagicMock() - config.attention_config.indexer_kv_dtype = "int8" + config.attention_config = AttentionConfig(indexer_kv_dtype="int8") config.cache_config.cache_dtype = "auto" return config @@ -1181,7 +1182,7 @@ def test_hybrid_dflash_writes_preserve_other_request_blocks( original_specs = {"target.full": full, "draft.sliding": sliding, "draft.full": draft_full, "target.gdn": mamba} config = SimpleNamespace( parallel_config=SimpleNamespace(decode_context_parallel_size=1), - attention_config=SimpleNamespace(indexer_kv_dtype="int8"), + attention_config=AttentionConfig(indexer_kv_dtype="int8"), cache_config=SimpleNamespace(cache_dtype="auto"), model_config=SimpleNamespace(dtype=torch.bfloat16, hf_config=SimpleNamespace()), kv_transfer_config=None, diff --git a/tests/ut/worker/test_worker_v2.py b/tests/ut/worker/test_worker_v2.py index d1361e759da2..80b2b4a83ad1 100644 --- a/tests/ut/worker/test_worker_v2.py +++ b/tests/ut/worker/test_worker_v2.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.config import CUDAGraphMode, KVTransferConfig +from vllm.config import AttentionConfig, CUDAGraphMode, KVTransferConfig from vllm.sequence import IntermediateTensors from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, @@ -436,6 +436,7 @@ def test_mrv2_block_outermost_hybrid_budget_respects_shared_backing(with_kv_tran worker.use_v2_model_runner = True worker.model_config = SimpleNamespace(hf_text_config=SimpleNamespace(model_type="qwen4_exp")) worker.vllm_config = SimpleNamespace( + attention_config=AttentionConfig(), cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC), kv_transfer_config=( KVTransferConfig(kv_connector="MooncakeConnector", kv_role="kv_both") if with_kv_transfer else None From 48e4b4a63acd796f6a63f66aba009019007c387f Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 Date: Thu, 12 Nov 2026 02:22:56 +0800 Subject: [PATCH 119/120] fix(ci): preserve the based snapshot in csrc cache checks Signed-off-by: zhenwenqi2024 --- .github/workflows/_ensure_csrc_cache.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/_ensure_csrc_cache.yaml b/.github/workflows/_ensure_csrc_cache.yaml index f5d25a635e80..890a83aa21cd 100644 --- a/.github/workflows/_ensure_csrc_cache.yaml +++ b/.github/workflows/_ensure_csrc_cache.yaml @@ -139,7 +139,7 @@ jobs: fi git fetch "$BASE_REMOTE" "$BASE_REF" git cat-file -e "${BASE_SHA}^{commit}" - git rebase "$BASE_SHA" + git merge-base --is-ancestor "$BASE_SHA" HEAD || git rebase "$BASE_SHA" - name: Calculate csrc hash id: hash From f8110b0b40a7c56d74939b2ef92081e8654c6358 Mon Sep 17 00:00:00 2001 From: LifeOfCoffee Date: Thu, 12 Nov 2026 14:11:17 +0800 Subject: [PATCH 120/120] chore: define HiSparse source readiness interface Add the minimal worker-side Host writeback contract for the follow-up AscendStore pooling integration. This commit does not wire or change runtime behavior. Signed-off-by: LifeOfCoffee --- .../kv_transfer/kv_pool/ascend_store/metadata.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py index 2a29497efe95..3f203335f610 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field from itertools import compress -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Protocol, cast import numpy as np import torch @@ -20,6 +20,12 @@ from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise.metadata import LayerPoolKey +class HiSparseSourceReady(Protocol): + """Return after Host writeback; raise if blocks are canceled or invalid.""" + + def wait_source_ready(self, req_id: str, source_block_ids: Sequence[int], end_token: int) -> None: ... + + def is_kv_save_role(kv_role: str, consumer_is_to_put: bool) -> bool: return kv_role in ("kv_producer", "kv_both") or consumer_is_to_put