diff --git a/.github/workflows/scripts/estimated_times.yaml b/.github/workflows/scripts/estimated_times.yaml index 6e0c1cdd3a7e..3cf82f7b6d2a 100644 --- a/.github/workflows/scripts/estimated_times.yaml +++ b/.github/workflows/scripts/estimated_times.yaml @@ -14,7 +14,7 @@ estimated_times: tests/e2e/pull_request/one_card/_310p/test_dense_model_310p.py: 1210 tests/e2e/pull_request/one_card/_310p/test_embedding_310p.py: 480 tests/e2e/pull_request/one_card/_310p/test_scoring_310p.py: 330 - tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py: 340 + tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py: 600 tests/e2e/pull_request/one_card/_310p/test_vl_model_310p.py: 370 tests/e2e/pull_request/one_card/aclgraph/test_aclgraph_batch_invariant.py: 550 tests/e2e/pull_request/one_card/compile/test_graphex_norm_quant_fusion.py: 170 diff --git a/tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py b/tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py index 3c86a83edd6f..db33098b5922 100644 --- a/tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py +++ b/tests/e2e/pull_request/one_card/_310p/test_spec_decode_mtp_310p.py @@ -15,13 +15,26 @@ # limitations under the License. # This file is a part of the vllm-ascend project. -from tests.e2e.conftest import VllmRunner +"""310P MTP e2e: MRv1 baseline + MRv2 eager smoke (1-card CI safe).""" + +import os +from unittest.mock import patch + +from tests.e2e.conftest import VllmRunner, wait_until_npu_memory_free + +# CI uses hub id; local verify can override (e.g. /home/weights/Qwen3.5-4B-W8A8). +QWEN35_MTP_MODEL = os.environ.get("QWEN35_MTP_MODEL", "Qwen/Qwen3.5-4B") +QWEN35_MTP_QUANTIZATION = os.environ.get("QWEN35_MTP_QUANTIZATION") # e.g. "ascend" + + +def _quant_kw(): + return {"quantization": QWEN35_MTP_QUANTIZATION} if QWEN35_MTP_QUANTIZATION else {} def test_qwen3_5_mtp_tp1_eager(): - example_prompts = ["Hello, my name is"] + """MRv1 baseline (no V2 runner env).""" with VllmRunner( - "Qwen/Qwen3.5-4B", + QWEN35_MTP_MODEL, tensor_parallel_size=1, enforce_eager=True, dtype="float16", @@ -31,5 +44,32 @@ def test_qwen3_5_mtp_tp1_eager(): "method": "qwen3_5_mtp", "num_speculative_tokens": 1, }, + **_quant_kw(), + ) as vllm_model: + vllm_model.generate_greedy(["Hello, my name is"], max_tokens=8) + + +@wait_until_npu_memory_free() +@patch.dict(os.environ, {"VLLM_USE_V2_MODEL_RUNNER": "1"}) +def test_qwen3_5_mtp_mrv2_tp1_eager(): + """MRv2 MTP eager smoke. + + FULL_DECODE_ONLY is covered by local/nightly runs: 1-card PR CI OOMs during + target+draft ACLGraph capture (Engine core init fails with empty Failed core + proc(s)). + """ + with VllmRunner( + QWEN35_MTP_MODEL, + tensor_parallel_size=1, + enforce_eager=True, + dtype="float16", + max_model_len=2048, + mamba_ssm_cache_dtype="float16", + speculative_config={ + "method": "mtp", + "num_speculative_tokens": 1, + }, + **_quant_kw(), ) as vllm_model: - vllm_model.generate_greedy(example_prompts, max_tokens=8) + assert vllm_model.model.llm_engine.vllm_config.use_v2_model_runner + vllm_model.generate_greedy(["Hello, my name is"], max_tokens=8) diff --git a/tests/ut/_310p/ops/test_gdn_310.py b/tests/ut/_310p/ops/test_gdn_310.py index 35990c0a27e2..06ec14e94110 100644 --- a/tests/ut/_310p/ops/test_gdn_310.py +++ b/tests/ut/_310p/ops/test_gdn_310.py @@ -18,7 +18,7 @@ from types import SimpleNamespace import torch -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm_ascend._310p.ops.fla.gdn_310 import ( AscendGatedDeltaNetAttention310, @@ -85,15 +85,17 @@ def test_builder310_pads_spec_decode_metadata_with_dummy_requests(): builder._pad_spec_decode_metadata(attn_metadata, graph_batch_size=4) + # 310P pads with PAD_SLOT_ID (-1), not NULL_BLOCK_ID (0), so FULL replay + # does not write into mamba block 0. Pad accepted tokens stay 1 (not 0). assert attn_metadata.spec_state_indices_tensor.tolist() == [ [3, 30], [4, 40], - [NULL_BLOCK_ID, NULL_BLOCK_ID], - [NULL_BLOCK_ID, NULL_BLOCK_ID], + [PAD_SLOT_ID, PAD_SLOT_ID], + [PAD_SLOT_ID, PAD_SLOT_ID], ] assert attn_metadata.spec_sequence_masks.tolist() == [True, True, False, False] assert attn_metadata.spec_query_start_loc.tolist() == [0, 4, 8, 8, 8] - assert attn_metadata.num_accepted_tokens.tolist() == [2, 3, 0, 0] + assert attn_metadata.num_accepted_tokens.tolist() == [2, 3, 1, 1] spec_meta = attn_metadata.spec_decode_metadata.spec_causal_conv1d assert spec_meta.query_start_loc.data_ptr() == attn_metadata.spec_query_start_loc.data_ptr() assert spec_meta.cache_indices.data_ptr() == attn_metadata.spec_state_indices_tensor.data_ptr() diff --git a/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py b/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py new file mode 100644 index 000000000000..d757e9f8a4fe --- /dev/null +++ b/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# Lean UTs for 310P MRv2 MTP (rejection offset, capture-safe draft step, RoPE flag). + +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from vllm.config.compilation import CUDAGraphMode + +from tests.ut.base import TestBase +from vllm_ascend._310p.ops.rotary_embedding import AscendRotaryEmbedding310 +from vllm_ascend._310p.worker.v2.spec_decode.aclgraph import AutoRegressiveAclGraphManager310 +from vllm_ascend._310p.worker.v2.spec_decode.mtp_speculator import AscendMTPSpeculator310 +from vllm_ascend._310p.worker.v2.spec_utils import ( + greedy_rejection_sample_cpu, + set_draft_step_host, + update_draft_inputs_cpu, +) + + +class TestMRv2Mtp310(TestBase): + def test_greedy_rejection_uses_logit_idx_plus_one(self): + # draft_sampled[logit_idx+1] must match argmax(logits[logit_idx]) to accept. + target_logits = torch.tensor( + [ + [0.1, 0.9, 0.0], # predicts token 1 + [0.0, 0.2, 0.8], # bonus → token 2 + ], + dtype=torch.float32, + ) + draft_sampled = torch.tensor([7, 1], dtype=torch.int32) + cu_num_logits = torch.tensor([0, 2], dtype=torch.int32) + + sampled, num_sampled = greedy_rejection_sample_cpu( + target_logits, draft_sampled, cu_num_logits, num_speculative_steps=1 + ) + + self.assertEqual(num_sampled.tolist(), [2]) + self.assertEqual(sampled[0, :2].tolist(), [1, 2]) + + def test_update_draft_inputs_uses_host_step_under_capture(self): + num_reqs = 2 + draft_tokens = torch.tensor([11, 22], dtype=torch.int32) + current_draft_step = torch.tensor([99], dtype=torch.int64) # must not .item() under capture + hidden_states = torch.randn(num_reqs, 4) + output_draft_tokens = torch.full((num_reqs, 2), -1, dtype=torch.int32) + next_input_hidden_states = torch.zeros(num_reqs, 4) + input_buffers = SimpleNamespace( + input_ids=torch.zeros(num_reqs, dtype=torch.int32), + positions=torch.tensor([3, 5], dtype=torch.int64), + seq_lens=torch.tensor([4, 6], dtype=torch.int32), + ) + set_draft_step_host(0) + + with patch("torch.npu.is_current_stream_capturing", return_value=True): + update_draft_inputs_cpu( + draft_tokens=draft_tokens, + current_draft_step=current_draft_step, + hidden_states=hidden_states, + output_draft_tokens=output_draft_tokens, + next_input_hidden_states=next_input_hidden_states, + input_buffers=input_buffers, + num_reqs=num_reqs, + max_model_len=128, + num_speculative_steps=2, + advance_draft_positions=True, + ) + + self.assertEqual(output_draft_tokens[:, 0].tolist(), [11, 22]) + self.assertEqual(input_buffers.positions.tolist(), [4, 6]) + + def test_run_model_sets_rope_flag(self): + flag_states: list[bool] = [] + + def mock_parent_run(self, *args, **kwargs): + del self, args, kwargs + flag_states.append(AscendRotaryEmbedding310._is_drafting_update_enabled) + return torch.zeros(1), torch.zeros(1) + + speculator = object.__new__(AscendMTPSpeculator310) + with patch( + "vllm_ascend.worker.v2.spec_decode.autoregressive.speculator.AscendAutoRegressiveSpeculator._run_model", + mock_parent_run, + ): + AscendMTPSpeculator310._run_model( + speculator, + num_tokens=1, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + + self.assertEqual(flag_states, [True]) + self.assertFalse(AscendRotaryEmbedding310._is_drafting_update_enabled) + + def test_decode_capture_routes_to_per_step(self): + manager = object.__new__(AutoRegressiveAclGraphManager310) + manager.is_draft_model_prefill = False + called = {"per_step": False} + + def fake_per_step(*args, **kwargs): + del args, kwargs + called["per_step"] = True + + with patch.object(AutoRegressiveAclGraphManager310, "_capture_decode_per_step", fake_per_step): + AutoRegressiveAclGraphManager310.capture( + manager, + forward_fn=lambda: None, + model_state=object(), + input_buffers=object(), + block_tables=object(), + attn_groups=[], + kv_cache_config=object(), + ) + + self.assertTrue(called["per_step"]) diff --git a/tests/ut/_310p/test_model_runner_v2_310p.py b/tests/ut/_310p/test_model_runner_v2_310p.py index 23df9ea27379..41975700d48d 100644 --- a/tests/ut/_310p/test_model_runner_v2_310p.py +++ b/tests/ut/_310p/test_model_runner_v2_310p.py @@ -340,10 +340,19 @@ def test_config_rejects_non_tp_parallelism(setting: str) -> None: NPUModelRunner310V2._validate_config(config) +def test_config_accepts_mtp_and_rejects_non_mtp() -> None: + """310P MRv2 allows method=mtp only.""" + NPUModelRunner310V2._validate_config( + _make_vllm_config(speculative_config=SimpleNamespace(method="mtp", num_speculative_tokens=1)) + ) + with pytest.raises(NotImplementedError, match="only supported via MTP"): + NPUModelRunner310V2._validate_config(_make_vllm_config(speculative_config=SimpleNamespace(method="eagle"))) + + @pytest.mark.parametrize( ("field", "value", "message"), [ - ("speculative_config", object(), "Speculative decoding"), + ("speculative_config", object(), "only supported via MTP"), ("kv_transfer_config", object(), "KV cache transfer"), ("lora_config", object(), "LoRA"), ], @@ -353,6 +362,26 @@ def test_config_rejects_out_of_scope_features(field, value, message) -> None: NPUModelRunner310V2._validate_config(_make_vllm_config(**{field: value})) +def test_copy_kv_cache_blocks_flattens_mamba_lists() -> None: + """Prefix-cache CoW must flatten list[Tensor] mamba layers for upstream copy.""" + runner = object.__new__(NPUModelRunner310V2) + runner._attn_kv_copy_params = [] + t0 = torch.zeros(4, 2) + t1 = torch.zeros(4, 2) + runner.kv_caches = [[t0, t1], torch.zeros(2)] # hybrid: mamba list + other + runner.kv_cache_config = SimpleNamespace(num_blocks=4) + copies = [SimpleNamespace(src_block_id=0, dst_block_id=1)] + + with patch.object(model_runner_module, "copy_kv_cache_blocks_inplace") as mock_copy: + NPUModelRunner310V2._copy_kv_cache_blocks_310p(runner, copies) + + mock_copy.assert_called_once() + tensors_arg, num_blocks, copies_arg = mock_copy.call_args[0] + assert tensors_arg == [t0, t1] + assert num_blocks == 4 + assert copies_arg is copies + + def test_sampler_rejects_random_sampling_parameters() -> None: sampler = Ascend310PSampler() sampler.add_request(0, 4, SamplingParams(temperature=0)) diff --git a/vllm_ascend/_310p/attention/attention_v1.py b/vllm_ascend/_310p/attention/attention_v1.py index 4dff07b91deb..c791781c13e4 100644 --- a/vllm_ascend/_310p/attention/attention_v1.py +++ b/vllm_ascend/_310p/attention/attention_v1.py @@ -184,6 +184,12 @@ def forward_paged_attention( Any: The result of the attention operation. """ if attn_metadata.seq_lens.device != query.device: + # Pageable H2D is illegal under NPU GLOBAL ACLGraph capture. + if torch.npu.is_current_stream_capturing(): + raise RuntimeError( + "310P paged attention: seq_lens must already be on-device before " + "ACLGraph capture; move it outside torch.npu.graph()." + ) attn_metadata.seq_lens = attn_metadata.seq_lens.to( device=query.device, non_blocking=True, @@ -265,6 +271,8 @@ def forward_chunked_prefill_310(self, query, attn_metadata, output): block_table = attn_metadata.block_tables if attn_metadata.seq_lens.device != query.device: + if torch.npu.is_current_stream_capturing(): + raise RuntimeError("310P splitfuse: seq_lens must already be on-device before ACLGraph capture.") attn_metadata.seq_lens = attn_metadata.seq_lens.to( device=query.device, non_blocking=True, diff --git a/vllm_ascend/_310p/ops/fla/fused_recurrent_gated_delta_rule.py b/vllm_ascend/_310p/ops/fla/fused_recurrent_gated_delta_rule.py index 9487e2efc26f..521a72b883df 100644 --- a/vllm_ascend/_310p/ops/fla/fused_recurrent_gated_delta_rule.py +++ b/vllm_ascend/_310p/ops/fla/fused_recurrent_gated_delta_rule.py @@ -96,10 +96,12 @@ def _run_recurrent_gated_delta_rule( if seq_len <= 0: continue + # Match NPU recurrent_gated_delta_rule_v310: num_accepted_tokens only + # selects the resume state slot (accepted-1). Do NOT truncate the + # current query (MTP verify still has 1+K tokens when prior accept=1). accepted = None if num_accepted_tokens is not None: accepted = int(num_accepted_tokens[seq_idx].item()) - seq_len = min(seq_len, accepted) if seq_len <= 0: continue diff --git a/vllm_ascend/_310p/ops/fla/gdn_310.py b/vllm_ascend/_310p/ops/fla/gdn_310.py index 48e450186927..77a40c8704e0 100644 --- a/vllm_ascend/_310p/ops/fla/gdn_310.py +++ b/vllm_ascend/_310p/ops/fla/gdn_310.py @@ -64,16 +64,17 @@ def _flatten_state_indices( return ssm_state_indices[:total_tokens].to(torch.int32).contiguous() num_seqs = (cu_seqlens[1:] - cu_seqlens[:-1]).shape[0] - seq_lens = cu_seqlens[1 : num_seqs + 1] - cu_seqlens[:num_seqs] - ssm_state_indices = ssm_state_indices[:num_seqs] + q_per_seq = ssm_state_indices.shape[1] # Uniform spec-decode ACL graph uses fixed q_len per request; reshape avoids - # NPU masked_select which breaks stream capture (aclnnMaskedSelect / 107027). - if _EXTRA_CTX.capturing or (seq_lens.numel() > 0 and torch.all(seq_lens == seq_lens[0])): - q_per_seq = ssm_state_indices.shape[1] - flat = ssm_state_indices[:, :q_per_seq].reshape(-1) + # NPU masked_select and seq_lens scalar reads which break stream capture. + if _EXTRA_CTX.capturing or total_tokens == num_seqs * q_per_seq: + flat = ssm_state_indices[:num_seqs, :q_per_seq].reshape(-1) return flat[:total_tokens].to(torch.int32).contiguous() + seq_lens = cu_seqlens[1 : num_seqs + 1] - cu_seqlens[:num_seqs] + ssm_state_indices = ssm_state_indices[:num_seqs] + # Eager mixed batches with variable seq_lens: compact on CPU, copy back async. ssm_cpu = ssm_state_indices.cpu() seq_lens_cpu = seq_lens.cpu() @@ -119,10 +120,8 @@ def npu_recurrent_gated_delta_rule_310( total_tokens = v.shape[1] flat_state_indices = _flatten_state_indices(ssm_state_indices, cu_seqlens, total_tokens) actual_seq_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int32).contiguous() - flat_state_indices = torch.clamp_min( - flat_state_indices, - 0, - ).contiguous() + # Do not clamp PAD_SLOT_ID (-1) to 0 — that would write mamba block 0. + # Callers must slice away padded requests before invoking this helper. accepted_tokens = None if num_accepted_tokens is not None: accepted_tokens = _mask_padded_recurrent_accepted_tokens( @@ -262,15 +261,20 @@ def _forward_core( spec_valid_tokens, token_dim=0, ) + # Always slice to real spec rows (GPU GDN pattern). FULL-graph pad + # tails must not reach npu_causal_conv1d_310 / recurrent kernels — + # even with PAD_SLOT_ID, cu_seqlens / accepted desync under MTP + # has been observed to poison hybrid state and emit wrong tokens. + num_spec = attn_metadata.num_spec_decodes mixed_qkv_spec = torch.ops._C_ascend.npu_causal_conv1d_310( mixed_qkv_spec, conv_weights, bias=self.conv1d.bias, conv_states=conv_state, - query_start_loc=spec_query_start_loc_device, - cache_indices=spec_causal_conv1d_meta.cache_indices, + query_start_loc=spec_query_start_loc_device[: num_spec + 1], + cache_indices=spec_causal_conv1d_meta.cache_indices[:num_spec], initial_state_mode=None, - num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens, + num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens[:num_spec], activation_mode=activation_num, pad_slot_id=PAD_SLOT_ID, run_mode=1, @@ -335,6 +339,7 @@ def _forward_core( # 2.1: Process the multi-query part if spec_sequence_masks is not None: + num_spec = attn_metadata.num_spec_decodes core_attn_out_spec = npu_recurrent_gated_delta_rule_310( q=query_spec, k=key_spec, @@ -342,9 +347,9 @@ def _forward_core( g=g_spec, beta=beta_spec, state=ssm_state, - cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1], - ssm_state_indices=spec_state_indices_tensor[: attn_metadata.num_spec_decodes], - num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens, + cu_seqlens=spec_query_start_loc[: num_spec + 1], + ssm_state_indices=spec_state_indices_tensor[:num_spec], + num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens[:num_spec], use_qk_l2norm_in_kernel=True, ) else: diff --git a/vllm_ascend/_310p/ops/gdn_attn_builder_310.py b/vllm_ascend/_310p/ops/gdn_attn_builder_310.py index cf24eebec84e..5334bed151ce 100644 --- a/vllm_ascend/_310p/ops/gdn_attn_builder_310.py +++ b/vllm_ascend/_310p/ops/gdn_attn_builder_310.py @@ -23,7 +23,7 @@ import torch from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID from vllm_ascend._310p.ops.fla.cumpute_causal_conv1d_metadata_310 import ( compute_causal_conv1d_metadata, @@ -128,7 +128,9 @@ def _pad_spec_decode_metadata( non_blocking=True, ) attn_metadata.spec_state_indices_tensor = self.spec_state_indices_tensor[:graph_batch_size] - attn_metadata.spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID) + # Match PAD_SLOT_ID expected by npu_causal_conv1d_310; NULL_BLOCK_ID(=0) + # would write padding rows into mamba block 0 under FULL graph replay. + attn_metadata.spec_state_indices_tensor[num_spec_decodes:].fill_(PAD_SLOT_ID) self.spec_sequence_masks[:num_spec_decodes].copy_( spec_sequence_masks[:num_spec_decodes], @@ -169,7 +171,9 @@ def _pad_spec_decode_metadata( non_blocking=True, ) attn_metadata.num_accepted_tokens = self.num_accepted_tokens[:graph_batch_size] - attn_metadata.num_accepted_tokens[num_spec_decodes:].fill_(0) + # Match Ascend mainline / Mamba neutral value (1), not 0. Padding with 0 + # makes accepted-1=-1 and corrupts GDN recurrent init for real rows. + attn_metadata.num_accepted_tokens[num_spec_decodes:].fill_(1) self._attach_spec_decode_metadata(attn_metadata) def _pad_decode_metadata( @@ -231,11 +235,15 @@ def build( # type: ignore[override] return attn_metadata graph_batch_size = common_attn_metadata.num_reqs + # Spec metadata is request-granular (see vLLM GDN comment). Do not gate + # on ``num_spec_decode_tokens``: concurrent MTP has tokens = reqs*(1+K), + # which can equal/exceed ``decode_cudagraph_max_bs`` (often set from + # max capture *token* size). Skipping pad then leaves ephemeral tensors + # that FULL replay cannot refresh — concurrent SpecDecoding poison. if ( attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0 and attn_metadata.num_spec_decodes <= self.decode_cudagraph_max_bs - and attn_metadata.num_spec_decode_tokens <= self.decode_cudagraph_max_bs ): self._pad_spec_decode_metadata(attn_metadata, graph_batch_size) elif ( diff --git a/vllm_ascend/_310p/worker/v2/aclgraph.py b/vllm_ascend/_310p/worker/v2/aclgraph.py new file mode 100644 index 000000000000..94ca811b18a2 --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/aclgraph.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""310P ACLGraph helpers for MRv2 MTP FULL_DECODE_ONLY. + +Target verify on 310P maps MTP uniform batches (q_len = 1+K) to SpecDecoding +(splitfuse), not DecodeOnly (PA). Upstream ``AscendInputBatch.make_dummy`` +always tags DecodeOnly, so FULL capture would record the wrong attention path +and replay would diverge from runtime SpecDecoding. Wrap capture only on 310P. + +MRv1 concurrent SpecDecoding FULL relies on buffer-address refresh (no FIA +``graph_task`` on 310P). Mirror that: sync before replay so H2D into capture- +stable seq_lens / slot_mapping / GDN pad buffers is visible to the graph. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn as nn +from vllm.sequence import IntermediateTensors +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.utils import AttentionGroup + +from vllm_ascend.worker.v2.aclgraph_utils import ModelAclGraphManager +from vllm_ascend.worker.v2.input_batch import AscendInputBatch + + +class ModelAclGraphManager310(ModelAclGraphManager): + """310P target ACLGraph manager: MTP capture uses SpecDecoding metadata.""" + + def run_fullgraph(self, desc: BatchExecutionDescriptor) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + # MRv1 ``_model_forward`` synchronizes before speculative FULL replay so + # CPU→NPU refreshes of capture-stable buffers land before the graph. + if self.vllm_config.speculative_config is not None: + torch.npu.current_stream().synchronize() + return super().run_fullgraph(desc) + + def capture( + self, + model: nn.Module, + model_state: ModelState, + input_buffers: InputBuffers, + intermediate_tensors: IntermediateTensors | None, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + has_lora: bool = False, + use_aux_hidden_state_outputs: bool = False, + lora_capture_hook: Callable[[int, int, int], None] | None = None, + progress_bar_desc: str = "Capturing CUDA graphs", + pcp_manager: Any = None, + ) -> None: + speculative = self.vllm_config.speculative_config is not None + if not speculative or not self.cudagraph_mode.has_full_cudagraphs(): + return super().capture( + model, + model_state, + input_buffers, + intermediate_tensors, + block_tables, + attn_groups, + kv_cache_config, + has_lora=has_lora, + use_aux_hidden_state_outputs=use_aux_hidden_state_outputs, + lora_capture_hook=lora_capture_hook, + progress_bar_desc=progress_bar_desc, + pcp_manager=pcp_manager, + ) + + # Lazy import: attention_v1 pulls DeviceOperator and circularizes at module load. + from vllm_ascend.attention.attention_v1 import AscendAttentionState + + orig_make_dummy = AscendInputBatch.make_dummy + + @classmethod + def make_dummy_mtp( + cls, + num_reqs: int, + num_tokens: int, + input_buffers_arg: Any, + max_query_len: int | None = None, + ) -> AscendInputBatch: + kwargs: dict[str, Any] = {} + if max_query_len is not None: + kwargs["max_query_len"] = max_query_len + batch = orig_make_dummy(num_reqs, num_tokens, input_buffers_arg, **kwargs) + # Uniform MTP verify: q_len = 1+K (>1). Decode-only graphs stay PA. + if num_reqs > 0 and (num_tokens // num_reqs) > 1: + batch.attn_state = AscendAttentionState.SpecDecoding + return batch + + AscendInputBatch.make_dummy = make_dummy_mtp # type: ignore[method-assign] + try: + return super().capture( + model, + model_state, + input_buffers, + intermediate_tensors, + block_tables, + attn_groups, + kv_cache_config, + has_lora=has_lora, + use_aux_hidden_state_outputs=use_aux_hidden_state_outputs, + lora_capture_hook=lora_capture_hook, + progress_bar_desc=progress_bar_desc, + pcp_manager=pcp_manager, + ) + finally: + AscendInputBatch.make_dummy = orig_make_dummy # type: ignore[method-assign] diff --git a/vllm_ascend/_310p/worker/v2/model_runner.py b/vllm_ascend/_310p/worker/v2/model_runner.py index 2d3734b41837..604e36cd5382 100644 --- a/vllm_ascend/_310p/worker/v2/model_runner.py +++ b/vllm_ascend/_310p/worker/v2/model_runner.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors from __future__ import annotations +from collections.abc import Sequence from copy import deepcopy from typing import Any @@ -16,6 +18,7 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import get_dtype_size from vllm.v1 import kv_cache_interface +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -33,16 +36,21 @@ from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.kv_connector import get_kv_connector from vllm.v1.worker.gpu.model_runner import BatchReqState, sort_batch_req_ids -from vllm.v1.worker.utils import bind_kv_cache +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.utils import bind_kv_cache, copy_kv_cache_blocks_inplace from vllm_ascend._310p.attention.attention_v1 import AscendAttentionBackend310 +from vllm_ascend._310p.worker.v2.aclgraph import ModelAclGraphManager310 from vllm_ascend._310p.worker.v2.block_table import Ascend310PBlockTables from vllm_ascend._310p.worker.v2.kv_block_zeroer import AscendKVBlockZeroer310V2 +from vllm_ascend._310p.worker.v2.spec_utils import ( + combine_sampled_and_draft_tokens_cpu, + expand_idx_mapping_cpu, +) from vllm_ascend._310p.worker.v2.states import Ascend310PRequestState from vllm_ascend.core.kv_cache_interface import get_storage_block_size from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, get_kv_cache_tensor_layers, vllm_version_is -from vllm_ascend.worker.v2.aclgraph_utils import ModelAclGraphManager from vllm_ascend.worker.v2.attn_utils import build_attn_state from vllm_ascend.worker.v2.input_batch import AscendInputBatch from vllm_ascend.worker.v2.model_runner import NPUModelRunner @@ -79,6 +87,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # (same as MRv1 `_determine_batch_execution_and_padding`). FULL_DECODE_ONLY # already keeps those batches eager via mixed_mode=NONE. self._force_eager_pc_batch = False + self._force_eager_spec_batch = False + self._spec_dummy_capture = False + # Populated in initialize_kv_cache; also set here so UT can call + # ``_allocate_kv_cache_tensors`` without going through that path. + self._attn_kv_copy_params: list[tuple[torch.Tensor, torch.Tensor, int]] = [] + self._attn_kv_storage_ptrs: set[int] = set() @staticmethod def _validate_config(vllm_config: VllmConfig) -> None: @@ -106,9 +120,14 @@ def _validate_config(vllm_config: VllmConfig) -> None: ) if getattr(parallel_config, "enable_expert_parallel", False): raise NotImplementedError("Expert parallelism is not supported by model runner v2 on 310P.") - # TODO: Support speculative decoding in the next 310P MRV2 iteration. if vllm_config.speculative_config is not None: - raise NotImplementedError("Speculative decoding is not supported by model runner v2 on 310P.") + spec = vllm_config.speculative_config + # Bare objects (UT fixtures) and non-MTP methods are out of scope. + method = getattr(spec, "method", None) + if method != "mtp": + raise NotImplementedError( + f"Speculative decoding is only supported via MTP on 310P model runner v2, got {method!r}." + ) if vllm_config.kv_transfer_config is not None: raise NotImplementedError("KV cache transfer is not supported by model runner v2 on 310P.") # Prefix caching is supported: 310P MRv2 reuses CPU Ascend310PBlockTables / @@ -124,11 +143,6 @@ def _prepare_inputs_310p( ) -> AscendInputBatch: # TODO: Refactor this Triton-free input preparation through Triton # Dispatcher after vLLM RFC #45133 lands. - # ``super().execute_model`` has already run finish/add/update_requests and - # ``apply_staged_writes``; sync GPU counts now so mamba preprocess matches - # the CPU/np values used for positions and slot mappings. - self._sync_num_computed_tokens_gpu_from_np() - num_tokens = scheduler_output.total_num_scheduled_tokens num_tokens_after_padding = batch_desc.num_tokens assert num_tokens > 0 @@ -140,19 +154,36 @@ def _prepare_inputs_310p( scheduler_output.scheduled_spec_decode_tokens, self.decode_query_len, ) + # MTP: prior step async-D2Hs GPU→``num_computed_tokens_cpu``. Refresh np + # from that buffer BEFORE mirroring np→GPU, otherwise hybrid GDN + # preprocess_state sees stale counts and poisons SpecDecoding state. self._update_seq_lens_cpu(scheduler_output, req_ids) + # ``super().execute_model`` already ran finish/add/update_requests and + # ``apply_staged_writes``; sync GPU counts so mamba preprocess matches + # the CPU/np values used for positions and slot mappings. + self._sync_num_computed_tokens_gpu_from_np() num_scheduled_tokens = np.fromiter( map(num_tokens_per_req.get, req_ids), dtype=np.int32, count=num_reqs, ) + num_valid_tokens = num_scheduled_tokens + draft_tokens_map = scheduler_output.scheduled_spec_decode_tokens + if draft_tokens_map: + num_valid_tokens = np.array( + [ + num_tokens - len(draft_tokens_map.get(req_id, ())) + for num_tokens, req_id in zip(num_scheduled_tokens, req_ids) + ], + dtype=np.int32, + ) attn_state = build_attn_state( self.vllm_config, self.input_buffers.seq_lens_np, num_reqs, num_scheduled_tokens, - num_scheduled_tokens, + num_valid_tokens, ) idx_mapping_np = np.fromiter( map(self.req_states.req_id_to_index.get, req_ids), @@ -161,6 +192,33 @@ def _prepare_inputs_310p( ) idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + num_draft_tokens_per_req = None + if not draft_tokens_map: + total_num_draft_tokens = 0 + cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) + cu_num_logits = torch.arange(num_reqs + 1, device=self.device, dtype=torch.int32) + expanded_idx_mapping = idx_mapping + expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=self.device) + else: + num_draft_tokens_per_req = np.fromiter( + (len(draft_tokens_map.get(req_id, ())) for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step + total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) + num_logits_per_req = num_draft_tokens_per_req + num_bonus_tokens + cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) + cu_num_logits_np[0] = 0 + np.cumsum(num_logits_per_req, out=cu_num_logits_np[1:]) + cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) + total_num_logits = int(cu_num_logits_np[-1]) + expanded_idx_mapping, expanded_local_pos = expand_idx_mapping_cpu( + idx_mapping, + total_num_logits, + cu_num_logits_np, + ) + num_reqs_padded = batch_desc.num_reqs or num_reqs query_start_loc_np = np.empty(self.max_num_reqs + 2, dtype=np.int32) query_start_loc_np[0] = 0 @@ -208,9 +266,10 @@ def _prepare_inputs_310p( num_scheduled_tokens=num_scheduled_tokens, ) seq_lens = self.input_buffers.seq_lens[:num_reqs_padded] + # Pad rows must not carry stale host seq_lens into GDN/attention metadata. + self.input_buffers.seq_lens_np[num_reqs:num_reqs_padded] = 0 self.input_buffers.seq_lens_np[num_reqs_padded:] = 0 - cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) - cu_num_logits = torch.arange(num_reqs + 1, device=self.device, dtype=torch.int32) + total_num_logits = num_reqs if not draft_tokens_map else int(cu_num_logits_np[-1]) logits_indices = self._combine_sampled_and_draft_tokens( self.input_buffers.input_ids, idx_mapping, @@ -220,7 +279,7 @@ def _prepare_inputs_310p( self.req_states.prefill_len.gpu, self.req_states.draft_tokens, cu_num_logits, - num_reqs, + total_num_logits, self.model_state.num_new_sampled_tokens_per_step, idx_mapping_np=idx_mapping_np, query_start_loc_np=query_start_loc_np, @@ -240,13 +299,13 @@ def _prepare_inputs_310p( num_reqs_after_padding=num_reqs_padded, idx_mapping=idx_mapping, idx_mapping_np=idx_mapping_np, - expanded_idx_mapping=idx_mapping, - expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=self.device), + expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, num_scheduled_tokens=num_scheduled_tokens, num_tokens=num_tokens, num_tokens_after_padding=num_tokens_after_padding, - num_draft_tokens=0, - num_draft_tokens_per_req=None, + num_draft_tokens=total_num_draft_tokens if draft_tokens_map else 0, + num_draft_tokens_per_req=num_draft_tokens_per_req, query_start_loc=query_start_loc, query_start_loc_np=query_start_loc_np, seq_lens=seq_lens, @@ -327,6 +386,85 @@ def _scheduler_output_needs_pc_eager(self, scheduler_output: SchedulerOutput) -> # Avoid importing AscendAttentionState at module top (heavy attention_v1). return attn_state.name in ("PrefillCacheHit", "ChunkedPrefill") + def _scheduler_output_needs_spec_eager(self, scheduler_output: SchedulerOutput) -> bool: + """Force eager when MTP verify batch is not uniform SpecDecoding. + + Step 2 (FULL_DECODE_ONLY): mirror MRv1 ``_determine_batch_execution_and_padding``. + Uniform decode with ``q_len == decode_query_len`` (1+K) may replay SpecDecoding + FULL graphs; mixed / prefill / non-uniform MTP schedules stay eager. + """ + if self.speculative_config is None: + return False + cudagraph_mode = self.compilation_config.cudagraph_mode + if not cudagraph_mode.has_full_cudagraphs(): + return False + + num_tokens_per_req = scheduler_output.num_scheduled_tokens + num_reqs = len(num_tokens_per_req) + if num_reqs == 0: + return False + + # Prefer stable req order (same as prepare_inputs) for draft counts. + req_ids = sort_batch_req_ids( + num_tokens_per_req, + scheduler_output.scheduled_spec_decode_tokens, + self.decode_query_len, + ) + num_scheduled = np.fromiter( + (num_tokens_per_req[req_id] for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + if not np.all(num_scheduled == self.decode_query_len): + return True + if scheduler_output.total_num_scheduled_tokens != int(num_scheduled.sum()): + return True + + computed_by_req: dict[str, int] = {} + for req in scheduler_output.scheduled_new_reqs: + computed_by_req[req.req_id] = int(req.num_computed_tokens) + cached = scheduler_output.scheduled_cached_reqs + if cached is not None: + for req_id, num_computed in zip(cached.req_ids, cached.num_computed_tokens): + computed_by_req[req_id] = int(num_computed) + for req_id in req_ids: + if req_id in computed_by_req: + continue + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is not None: + computed_by_req[req_id] = int(self.req_states.num_computed_tokens_np[req_idx]) + + if any(computed_by_req.get(req_id, 0) == 0 for req_id in req_ids): + return True + + draft_tokens_map = scheduler_output.scheduled_spec_decode_tokens or {} + # SpecDecoding requires bonus/valid token count == 1 (rest are drafts). + num_valid_tokens = np.fromiter( + (int(num_tokens_per_req[req_id]) - len(draft_tokens_map.get(req_id, ())) for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + if not np.all(num_valid_tokens == 1): + return True + + # Allow concurrent uniform SpecDecoding FULL (MRv1 contract). Accuracy + # regressions are caught by E2E; do not blanket-eager on num_reqs. + seq_lens = np.fromiter( + (computed_by_req[req_id] + num_tokens_per_req[req_id] for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + attn_state = build_attn_state( + self.vllm_config, + seq_lens, + num_reqs, + num_scheduled, + num_valid_tokens, + ) + from vllm_ascend.attention.attention_v1 import AscendAttentionState + + return attn_state != AscendAttentionState.SpecDecoding + def _install_pc_eager_cudagraph_dispatch(self) -> None: """Wrap ACLGraph dispatch so PrefillCacheHit cannot replay FULL mixed graphs.""" manager = self.cudagraph_manager @@ -342,7 +480,7 @@ def dispatch( num_active_loras: int, max_query_len: int | None = None, ) -> BatchExecutionDescriptor: - if runner._force_eager_pc_batch: + if runner._force_eager_pc_batch or runner._force_eager_spec_batch: return BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, @@ -371,10 +509,14 @@ def _sync_num_computed_tokens_gpu_from_np(self) -> None: early in ``execute_model`` leaves stale counts and corrupts recurrent state. """ np_vals = self.req_states.num_computed_tokens_np + # Copy host buffer so non-blocking H2D cannot race later np mutations. + host = torch.tensor(np_vals, dtype=torch.int32) gpu = self.req_states.num_computed_tokens.gpu - gpu.copy_(torch.from_numpy(np_vals).to(device=gpu.device, dtype=gpu.dtype)) - self.req_states.num_computed_tokens_cpu.copy_(torch.from_numpy(np_vals)) - self.req_states.num_computed_tokens.cpu.copy_(torch.from_numpy(np_vals)) + if host.dtype != gpu.dtype: + host = host.to(dtype=gpu.dtype) + gpu.copy_(host, non_blocking=True) + self.req_states.num_computed_tokens_cpu.copy_(host) + self.req_states.num_computed_tokens.cpu.copy_(host) def _advance_num_computed_tokens(self, valid_indices: torch.Tensor, query_lens: torch.Tensor) -> None: """Advance per-request computed counts on both CPU mirror and GPU tensor.""" @@ -401,8 +543,10 @@ def execute_model( valid_dummy_state_slots: bool = False, ): self._force_eager_pc_batch = False + self._force_eager_spec_batch = False if not dummy_run: self._force_eager_pc_batch = self._scheduler_output_needs_pc_eager(scheduler_output) + self._force_eager_spec_batch = self._scheduler_output_needs_spec_eager(scheduler_output) try: return super().execute_model( scheduler_output, @@ -415,6 +559,7 @@ def execute_model( ) finally: self._force_eager_pc_batch = False + self._force_eager_spec_batch = False def prepare_inputs( # type: ignore[misc, override] self, @@ -434,6 +579,68 @@ def finish_requests(self, scheduler_output: SchedulerOutput) -> None: # does not need this because it does not use that CPU gather path. torch.npu.current_stream().synchronize() + @staticmethod + def _dedupe_kv_cache_block_copies( + kv_cache_block_copies: Sequence[KVCacheBlockCopy], + ) -> list[KVCacheBlockCopy]: + """Drop duplicate CoW pairs from hybrid multi-manager prefix-cache hits.""" + seen: set[tuple[int, int]] = set() + deduped: list[KVCacheBlockCopy] = [] + for copy in kv_cache_block_copies: + key = (copy.src_block_id, copy.dst_block_id) + if key in seen: + continue + seen.add(key) + deduped.append(copy) + return deduped + + def _copy_kv_cache_blocks_310p(self, kv_cache_block_copies: Sequence[KVCacheBlockCopy]) -> None: + """Copy-on-write for hybrid prefix cache on 310P NZ attention + ND Mamba.""" + if not kv_cache_block_copies: + return + + indices_np = np.array( + [[copy.src_block_id, copy.dst_block_id] for copy in kv_cache_block_copies], + dtype=np.int64, + ) + seen_attn_storage: set[int] = set() + for k_cache, v_cache, blocks_per_kv_block in self._attn_kv_copy_params: + storage_ptr = k_cache.untyped_storage().data_ptr() + if storage_ptr in seen_attn_storage: + continue + seen_attn_storage.add(storage_ptr) + for src_block_id, dst_block_id in indices_np: + src_start = int(src_block_id) * blocks_per_kv_block + src_end = src_start + blocks_per_kv_block + dst_start = int(dst_block_id) * blocks_per_kv_block + dst_end = dst_start + blocks_per_kv_block + k_cache[dst_start:dst_end].copy_(k_cache[src_start:src_end]) + v_cache[dst_start:dst_end].copy_(v_cache[src_start:src_end]) + + # Mamba layers store a list[Tensor] per layer (conv/ssm views). Upstream + # copy_kv_cache_blocks_inplace expects Iterable[Tensor], not nested lists. + mamba_tensors: list[torch.Tensor] = [] + for entry in self.kv_caches: + if isinstance(entry, list): + mamba_tensors.extend(t for t in entry if isinstance(t, torch.Tensor)) + if mamba_tensors: + copy_kv_cache_blocks_inplace( + mamba_tensors, + self.kv_cache_config.num_blocks, + kv_cache_block_copies, + ) + + def update_requests(self, scheduler_output: SchedulerOutput) -> None: + copies = scheduler_output.kv_cache_block_copies + pending_copies: list[KVCacheBlockCopy] | None = None + if copies: + pending_copies = self._dedupe_kv_cache_block_copies(copies) + # Skip upstream copy (mishandles 310P NZ attention storages). + scheduler_output.kv_cache_block_copies = None + super().update_requests(scheduler_output) + if pending_copies: + self._copy_kv_cache_blocks_310p(pending_copies) + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: """Restore linear-attention specs omitted by some upstream V2 versions.""" kv_cache_spec = super().get_kv_cache_spec() @@ -503,7 +710,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: kv_cache_config=kv_cache_config, max_num_reqs=self.max_num_reqs, ) - self.cudagraph_manager = ModelAclGraphManager( + self.cudagraph_manager = ModelAclGraphManager310( self.vllm_config, self.device, cudagraph_mode, @@ -512,8 +719,20 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: lora_capture_cases=self.lora_capture_cases, ) check_attention_cp_compatibility(self.vllm_config) + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.set_attn( + self.model_state, + self.kv_cache_config, + self.block_tables, + self.input_buffers, + self.attn_groups, + ) + if self.speculator is not None: + self.speculator.init_cudagraph_manager(cudagraph_mode) shared_layers = get_shared_kv_cache_layers(self.vllm_config) + self._attn_kv_copy_params = [] + self._attn_kv_storage_ptrs = set() kv_caches_dict = self._allocate_kv_cache_tensors(kv_cache_config, shared_layers) self.kv_caches: list[Any] = [] bind_kv_cache( @@ -568,6 +787,11 @@ def _allocate_kv_cache_tensors( shared_layers: dict[str, str], ) -> dict[str, Any]: """Allocate attention caches as NZ and hybrid Mamba state as ND.""" + # UT fixtures may construct via ``object.__new__`` without ``__init__``. + if not hasattr(self, "_attn_kv_copy_params"): + self._attn_kv_copy_params = [] + if not hasattr(self, "_attn_kv_storage_ptrs"): + self._attn_kv_storage_ptrs = set() layer_specs: dict[str, KVCacheSpec] = {} layer_group_ids: dict[str, int] = {} for group_id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups): @@ -643,41 +867,47 @@ def _allocate_kv_cache_tensors( # Symmetric NZ only: K/V share the 4D view ``kv_cache_shape[1:]``. kv_view_shape = kv_cache_shape[1:] if legacy_shared_by: - cache: Any = ( - torch_npu.empty_with_format( - size=kv_view_shape, - dtype=kv_cache_spec.dtype, - device=self.device, - acl_format=ACL_FORMAT_FRACTAL_NZ, - ), - torch_npu.empty_with_format( - size=kv_view_shape, - dtype=kv_cache_spec.dtype, - device=self.device, - acl_format=ACL_FORMAT_FRACTAL_NZ, - ), + k_cache = torch_npu.empty_with_format( + size=kv_view_shape, + dtype=kv_cache_spec.dtype, + device=self.device, + acl_format=ACL_FORMAT_FRACTAL_NZ, + ) + v_cache = torch_npu.empty_with_format( + size=kv_view_shape, + dtype=kv_cache_spec.dtype, + device=self.device, + acl_format=ACL_FORMAT_FRACTAL_NZ, ) + cache: Any = (k_cache, v_cache) for name in cache_layer_names: kv_caches[name] = cache + storage_ptr = k_cache.untyped_storage().data_ptr() + if storage_ptr not in self._attn_kv_storage_ptrs: + self._attn_kv_storage_ptrs.add(storage_ptr) + self._attn_kv_copy_params.append((k_cache, v_cache, blocks_per_kv_block)) else: # Standardized descriptors list distinct layer # regions. Allocate one private NZ K/V pair per layer; # only explicit shared_layers below may alias. for name in cache_layer_names: - kv_caches[name] = ( - torch_npu.empty_with_format( - size=kv_view_shape, - dtype=kv_cache_spec.dtype, - device=self.device, - acl_format=ACL_FORMAT_FRACTAL_NZ, - ), - torch_npu.empty_with_format( - size=kv_view_shape, - dtype=kv_cache_spec.dtype, - device=self.device, - acl_format=ACL_FORMAT_FRACTAL_NZ, - ), + k_cache = torch_npu.empty_with_format( + size=kv_view_shape, + dtype=kv_cache_spec.dtype, + device=self.device, + acl_format=ACL_FORMAT_FRACTAL_NZ, + ) + v_cache = torch_npu.empty_with_format( + size=kv_view_shape, + dtype=kv_cache_spec.dtype, + device=self.device, + acl_format=ACL_FORMAT_FRACTAL_NZ, ) + kv_caches[name] = (k_cache, v_cache) + storage_ptr = k_cache.untyped_storage().data_ptr() + if storage_ptr not in self._attn_kv_storage_ptrs: + self._attn_kv_storage_ptrs.add(storage_ptr) + self._attn_kv_copy_params.append((k_cache, v_cache, blocks_per_kv_block)) elif isinstance(kv_cache_spec, MambaSpec): # Hybrid recurrent state stays ND (int8 raw plus views). # Main's descriptor.size is the entire virtual backing; @@ -813,20 +1043,20 @@ def _combine_sampled_and_draft_tokens( seq_lens_np: np.ndarray, prefill_len_np: np.ndarray, ) -> torch.Tensor: - # TODO: Refactor this CPU fallback to use Triton Dispatcher after vLLM - # RFC #45133 lands. - del idx_mapping, query_start_loc, seq_lens, prefill_len - del draft_tokens, cu_num_logits, num_bonus_tokens - if num_logits != len(idx_mapping_np): - # TODO: Support draft tokens in the next 310P MRV2 iteration. - raise NotImplementedError("310P MRV2 does not support draft tokens.") - logits_indices_np = np.empty(num_logits, dtype=np.int64) - for batch_idx, req_idx in enumerate(idx_mapping_np): - query_end = int(query_start_loc_np[batch_idx + 1]) - logits_indices_np[batch_idx] = query_end - 1 - if seq_lens_np[batch_idx] > prefill_len_np[batch_idx]: - input_ids[query_end - 1 : query_end].copy_(last_sampled_tokens[req_idx]) - return async_copy_to_gpu(logits_indices_np, device=self.device) + cu_num_logits_np = cu_num_logits.detach().cpu().numpy().astype(np.int32, copy=False) + return combine_sampled_and_draft_tokens_cpu( + input_ids, + idx_mapping_np, + last_sampled_tokens, + query_start_loc_np, + seq_lens_np, + prefill_len_np, + draft_tokens, + cu_num_logits_np, + num_logits, + num_bonus_tokens, + device=self.device, + ) def prepare_attn( self, @@ -869,13 +1099,22 @@ def sample( # TODO: Restore MRV1 structured output support in the next 310P MRV2 iteration. raise NotImplementedError("Structured output is not supported by model runner v2 on 310P.") logits = self.model.compute_logits(hidden_states[input_batch.logits_indices]) - sampler_output = self.sampler(logits, input_batch) - can_sample_np = input_batch.seq_lens_np[: input_batch.num_reqs] >= input_batch.prefill_len_np - num_sampled = async_copy_to_gpu(can_sample_np.astype(np.int32), device=self.device) - num_rejected = torch.zeros_like(num_sampled) - sampler_output.num_sampled = num_sampled - sampler_output.num_rejected = num_rejected - return sampler_output, num_sampled, num_rejected + if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: + sampler_output = self.sampler(logits, input_batch) + can_sample_np = input_batch.seq_lens_np[: input_batch.num_reqs] >= input_batch.prefill_len_np + num_sampled = async_copy_to_gpu(can_sample_np.astype(np.int32), device=self.device) + num_rejected = torch.zeros_like(num_sampled) + sampler_output.num_sampled = num_sampled + sampler_output.num_rejected = num_rejected + return sampler_output, num_sampled, num_rejected + + assert self.speculator is not None + sampler_output = self.rejection_sampler( + logits, + input_batch, + self.speculator.draft_logits, + ) + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected def postprocess_sampled( self, @@ -885,35 +1124,109 @@ def postprocess_sampled( num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, ) -> None: - # TODO: Refactor this 310P state update to use Triton Dispatcher after - # vLLM RFC #45133 lands. - del num_rejected num_entries = min(idx_mapping.shape[0], sampled_tokens.shape[0], num_sampled.shape[0]) idx_mapping = idx_mapping[:num_entries] sampled_tokens = sampled_tokens[:num_entries] num_sampled = num_sampled[:num_entries] valid_mask = idx_mapping >= 0 valid_indices = idx_mapping.masked_select(valid_mask) - sampled = sampled_tokens[:, 0].masked_select(valid_mask).to(self.req_states.last_sampled_tokens.dtype) valid_num_sampled = num_sampled.masked_select(valid_mask) has_sample = valid_num_sampled > 0 - token_positions = self.req_states.total_len.gpu[valid_indices].to(torch.int64) - old_tokens = self.req_states.all_token_ids.gpu[valid_indices, token_positions] - stored_tokens = torch.where(has_sample, sampled.to(torch.int32), old_tokens) - self.req_states.all_token_ids.gpu.index_put_((valid_indices, token_positions), stored_tokens) - old_last = self.req_states.last_sampled_tokens[valid_indices, 0] - self.req_states.last_sampled_tokens.index_copy_( - 0, - valid_indices, - torch.where(has_sample, sampled, old_last).unsqueeze(-1), - ) - self.req_states.total_len.gpu.index_add_(0, valid_indices, valid_num_sampled) + if self.speculator is not None and sampled_tokens.ndim == 2: + # MTP writeback: contiguous per-req slices only (avoid UVA scatter / + # index_copy_ corruption under concurrency on 310P). + valid_batch = torch.nonzero(valid_mask, as_tuple=False).flatten() + if valid_batch.numel() > 0: + req_idx_t = idx_mapping[valid_batch].to(torch.long) + count_t = num_sampled[valid_batch].to(torch.long) + start_t = self.req_states.total_len.gpu[req_idx_t].to(torch.long) + counts_host = count_t.detach().cpu().tolist() + reqs_host = req_idx_t.detach().cpu().tolist() + starts_host = start_t.detach().cpu().tolist() + batch_host = valid_batch.detach().cpu().tolist() + for i, count in enumerate(counts_host): + count = int(count) + if count <= 0: + continue + req_idx = int(reqs_host[i]) + start_pos = int(starts_host[i]) + row = sampled_tokens[int(batch_host[i]), :count].to(dtype=self.req_states.all_token_ids.gpu.dtype) + self.req_states.all_token_ids.gpu[req_idx, start_pos : start_pos + count] = row + self.req_states.last_sampled_tokens[req_idx, 0] = row[-1].to( + dtype=self.req_states.last_sampled_tokens.dtype + ) + self.req_states.total_len.gpu[req_idx] = start_pos + count + else: + sampled = sampled_tokens[:, 0].masked_select(valid_mask).to(self.req_states.last_sampled_tokens.dtype) + token_positions = self.req_states.total_len.gpu[valid_indices].to(torch.int64) + old_tokens = self.req_states.all_token_ids.gpu[valid_indices, token_positions] + stored_tokens = torch.where(has_sample, sampled.to(torch.int32), old_tokens) + self.req_states.all_token_ids.gpu.index_put_((valid_indices, token_positions), stored_tokens) + old_last = self.req_states.last_sampled_tokens[valid_indices, 0] + self.req_states.last_sampled_tokens.index_copy_( + 0, + valid_indices, + torch.where(has_sample, sampled, old_last).unsqueeze(-1), + ) + self.req_states.total_len.gpu.index_add_(0, valid_indices, valid_num_sampled) if query_start_loc is not None: query_lens = self._get_valid_query_lens(idx_mapping, query_start_loc) - self._advance_num_computed_tokens(valid_indices, query_lens) - self.model_state.postprocess_state(idx_mapping, num_sampled) + if self.speculator is not None: + num_rejected_valid = num_rejected.masked_select(valid_mask) + advance_lens = query_lens - num_rejected_valid.to(query_lens.dtype) + self._advance_num_computed_tokens(valid_indices, advance_lens) + else: + self._advance_num_computed_tokens(valid_indices, query_lens) + + self.model_state.postprocess_state( + idx_mapping, + num_sampled, + self.req_states.num_computed_tokens.gpu, + ) + + if self.speculator is not None: + self._copy_num_computed_tokens_to_cpu() + + def _copy_num_computed_tokens_to_cpu(self) -> None: + default_stream = torch.npu.current_stream() + assert self.num_computed_tokens_stream is not None + assert self.num_computed_tokens_cpu is not None + with torch.npu.stream(self.num_computed_tokens_stream): + self.num_computed_tokens_stream.wait_stream(default_stream) + self.num_computed_tokens_cpu.copy_( + self.req_states.num_computed_tokens.gpu, + non_blocking=True, + ) + self.num_computed_tokens_event.record() + + def _update_seq_lens_cpu( + self, + scheduler_output: SchedulerOutput, + req_ids: list[str], + ) -> None: + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + + if self.speculator is not None: + self.num_computed_tokens_event.synchronize() + cached_ids = scheduler_output.scheduled_cached_reqs.req_ids + if cached_ids: + req_indices = [self.req_states.req_id_to_index[req_id] for req_id in cached_ids] + vals = self.num_computed_tokens_cpu[req_indices].tolist() + for req_index, val in zip(req_indices, vals): + self.req_states.num_computed_tokens_cpu[req_index] = val + self.req_states.num_computed_tokens_np[req_index] = int(val) + else: + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + req_index = self.req_states.req_id_to_index[req_id] + self.req_states.num_computed_tokens_cpu[req_index] = self.req_states.num_computed_tokens_np[req_index] + + for i, req_id in enumerate(req_ids): + req_index = self.req_states.req_id_to_index[req_id] + num_computed_tokens = self.req_states.num_computed_tokens_cpu[req_index] + self.input_buffers.seq_lens_cpu[i] = num_computed_tokens + num_scheduled_tokens[req_id] + self.input_buffers.seq_lens_np[i] = self.input_buffers.seq_lens_cpu[i] @staticmethod def _get_valid_query_lens( diff --git a/vllm_ascend/_310p/worker/v2/model_state.py b/vllm_ascend/_310p/worker/v2/model_state.py index 23b7445c6171..661207a4be4e 100644 --- a/vllm_ascend/_310p/worker/v2/model_state.py +++ b/vllm_ascend/_310p/worker/v2/model_state.py @@ -1,25 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors """MRV2 model state for Ascend 310P (dense/VL + hybrid/GDN).""" from typing import Any +import numpy as np import torch import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridAttnMetadata from vllm.v1.worker.utils import AttentionGroup from vllm_ascend._310p.ops.rotary_embedding import prepare_mrope_cos_sin_slices_from_runner from vllm_ascend._310p.worker.v2.rope import Ascend310PRopeState, get_310p_rope_state from vllm_ascend.device.device_op import DeviceOperator +from vllm_ascend.worker.v2.attn_utils import build_attn_metadata from vllm_ascend.worker.v2.input_batch import AscendInputBatch from vllm_ascend.worker.v2.model_states.default import AscendModelState from vllm_ascend.worker.v2.model_states.mamba_hybrid import AscendMambaHybridModelState +from .rejection_sampler import RejectionSampler310V2 from .sampler import Ascend310PSampler @@ -121,7 +126,19 @@ def prepare_inputs(self, input_batch: AscendInputBatch, req_states): def custom_sampler(self, sampler): del sampler - return Ascend310PSampler(), None + # MTP propose/_dummy_run reads sampler.sampling_states.temperature/seeds. + # ``object.__new__`` UT fixtures may omit attrs set in real ``__init__``. + max_num_reqs = int(getattr(self, "max_num_reqs", 1) or 1) + device = getattr(self, "device", torch.device("cpu")) + base_sampler = Ascend310PSampler(max_num_reqs, device) + vllm_config = getattr(self, "vllm_config", None) + spec_config = None if vllm_config is None else vllm_config.speculative_config + if spec_config is None: + return base_sampler, None + method = getattr(spec_config, "method", None) + if method != "mtp": + raise NotImplementedError(f"310P MRv2 only supports MTP speculative decoding, got {method!r}.") + return base_sampler, RejectionSampler310V2(base_sampler, spec_config, device) class Ascend310PModelState(_Ascend310PModelStateMixin, AscendModelState): @@ -175,6 +192,109 @@ def __init__( self._capture_seq_lens_by_ptr = {} self._replace_310p_rope_state(encoder_cache) + def prepare_attn( + self, + input_batch: AscendInputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + ubatch_idx: int = 0, + ) -> dict[str, Any]: + """310P hybrid FULL: correct actual/padded token counts + uniform SpecDecoding pads. + + Upstream ``AscendMambaHybridModelState.prepare_attn`` passes only + ``num_tokens`` (padded under FULL), so GDN treats pad tokens as actual. + Pad rows also get ``draft_tokens=-1``, which turns a uniform SpecDecoding + FULL batch into mixed Spec+Prefill — diverging from the captured uniform + graph. Mirror AscendModelState's actual/padded split and keep pad rows on + the SpecDecoding path (draft=K, accepted=1), matching MRv1 pad accepted=1. + """ + assert ubatch_idx == 0, "DBO is not supported on Ascend" + if for_capture: + self._record_capture_seq_lens(input_batch.seq_lens) + elif cudagraph_mode == CUDAGraphMode.FULL: + self._refresh_capture_seq_lens(input_batch.seq_lens) + + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_input_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_input_tokens = input_batch.num_tokens + num_actual_reqs = input_batch.num_reqs + num_actual_tokens = input_batch.num_tokens + + is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") + is_prefilling[:num_actual_reqs] = torch.from_numpy(input_batch.is_prefilling_np) + + num_accepted_tokens = None + num_decode_draft_tokens_cpu = None + if not for_capture and self.vllm_config.num_speculative_tokens > 0: + num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) + num_accepted_tokens[:num_actual_reqs] = self.num_accepted_tokens_gpu[input_batch.idx_mapping] + + num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) + num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req + if num_draft_tokens_per_req is not None: + is_decode = input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1 + spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode + num_decode_draft_tokens_np[:num_actual_reqs] = np.where( + spec_decode_mask, + num_draft_tokens_per_req, + -1, + ) + # Align with upstream #15707: only promote pad rows to Spec when + # every real request is SpecDecoding and pad query lens == 1+K. + # Also keep pad rows Spec when attn_state is already SpecDecoding + # (310P target FULL capture / concurrent pad), matching MRv1. + if cudagraph_mode == CUDAGraphMode.FULL and num_reqs > num_actual_reqs: + expected_query_len = int(self.vllm_config.num_speculative_tokens) + 1 + padded_query_lens = np.diff(input_batch.query_start_loc_np[: num_reqs + 1])[num_actual_reqs:] + attn_state = input_batch.attn_state + is_spec = attn_state is not None and getattr(attn_state, "name", "") == "SpecDecoding" + if (spec_decode_mask.all() or is_spec) and np.all(padded_query_lens == expected_query_len): + num_decode_draft_tokens_np[num_actual_reqs:] = padded_query_lens - 1 + num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + + # Host seq_lens for pad rows must be 0 (GPU already zeroed in prepare_inputs). + seq_lens_np = input_batch.seq_lens_np + if seq_lens_np is not None and num_reqs > num_actual_reqs: + seq_lens_np = seq_lens_np.copy() + seq_lens_np[num_actual_reqs:num_reqs] = 0 + + model_specific_metadata = MambaHybridAttnMetadata( + is_prefilling=is_prefilling, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + self.attn_metadata = build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_actual_reqs=num_actual_reqs, + num_tokens=num_input_tokens, + num_actual_tokens=num_actual_tokens, + num_input_tokens=num_input_tokens, + is_prefilling=is_prefilling, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np), + max_query_len=input_batch.num_scheduled_tokens.max().item(), + seq_lens=input_batch.seq_lens, + max_seq_len=self.max_model_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + seq_lens_np=seq_lens_np, + positions=input_batch.positions, + attn_state=input_batch.attn_state, + model_specific_attn_metadata=model_specific_metadata, + for_cudagraph_capture=for_capture, + ) + return self.attn_metadata + def preprocess_state( self, input_batch: AscendInputBatch, diff --git a/vllm_ascend/_310p/worker/v2/rejection_sampler.py b/vllm_ascend/_310p/worker/v2/rejection_sampler.py new file mode 100644 index 000000000000..65d470f6e773 --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/rejection_sampler.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""Greedy rejection sampler for 310P MRv2 MTP (no Triton).""" + +from __future__ import annotations + +import torch +from vllm.config import SpeculativeConfig +from vllm.v1.worker.gpu.sample.output import SamplerOutput + +from vllm_ascend._310p.worker.v2.spec_utils import ( + get_num_sampled_and_rejected_cpu, + greedy_rejection_sample_cpu, +) +from vllm_ascend.worker.v2.input_batch import AscendInputBatch + + +class RejectionSampler310V2: + """Greedy MTP rejection sampler for 310P MRv2.""" + + def __init__( + self, + sampler, + spec_config: SpeculativeConfig, + device: torch.device, + ) -> None: + del device + self.sampler = sampler + self.num_speculative_steps = spec_config.num_speculative_tokens + + def __call__( + self, + logits: torch.Tensor, + input_batch: AscendInputBatch, + draft_logits: torch.Tensor | None, + ) -> SamplerOutput: + del draft_logits + draft_sampled = input_batch.input_ids[input_batch.logits_indices] + sampled, num_sampled = greedy_rejection_sample_cpu( + logits, + draft_sampled, + input_batch.cu_num_logits, + self.num_speculative_steps, + ) + num_sampled, num_rejected = get_num_sampled_and_rejected_cpu( + num_sampled, + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping_np, + input_batch.prefill_len_np, + ) + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + num_rejected=num_rejected, + ) diff --git a/vllm_ascend/_310p/worker/v2/sampler.py b/vllm_ascend/_310p/worker/v2/sampler.py index 91008beb5397..babe37d63a5d 100644 --- a/vllm_ascend/_310p/worker/v2/sampler.py +++ b/vllm_ascend/_310p/worker/v2/sampler.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors from types import SimpleNamespace @@ -9,13 +10,35 @@ class Ascend310PSampler: - """Triton-free sampler for 310P MRV2.""" + """Triton-free sampler for 310P MRV2. + + Exposes a minimal ``sampling_states`` surface so MTP draft ``propose()`` + (and ``_dummy_run``) can read ``temperature.gpu`` / ``seeds.gpu`` without + pulling in UVA-backed SamplingStates. + """ # TODO: Refactor this sampler to register 310P implementations through # Triton Dispatcher after vLLM RFC #45133 lands. - def __init__(self) -> None: + def __init__( + self, + max_num_reqs: int = 1, + device: torch.device | str | None = None, + ) -> None: self.penalties_state = SimpleNamespace(output_bin_counts=None) + self.max_num_reqs = max_num_reqs + if device is None: + device = torch.device("cpu") + elif not isinstance(device, torch.device): + device = torch.device(device) + self.device = device + # Greedy-only: temperature stays 0; seeds unused but must exist for MTP. + temperature_gpu = torch.zeros(max_num_reqs, dtype=torch.float32, device=device) + seeds_gpu = torch.zeros(max_num_reqs, dtype=torch.int64, device=device) + self.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(gpu=temperature_gpu), + seeds=SimpleNamespace(gpu=seeds_gpu), + ) def add_request( self, @@ -23,7 +46,7 @@ def add_request( prompt_len: int, sampling_params: SamplingParams, ) -> None: - del req_idx, prompt_len + del prompt_len unsupported = [] if sampling_params.temperature != 0: unsupported.append("temperature") @@ -50,6 +73,10 @@ def add_request( raise NotImplementedError( f"Unsupported sampling parameters on model runner v2 for 310P: {', '.join(unsupported)}." ) + if 0 <= req_idx < self.max_num_reqs: + self.sampling_states.temperature.gpu[req_idx] = 0.0 + seed = getattr(sampling_params, "seed", None) + self.sampling_states.seeds.gpu[req_idx] = 0 if seed is None else int(seed) def apply_staged_writes(self) -> None: pass diff --git a/vllm_ascend/_310p/worker/v2/spec_decode/__init__.py b/vllm_ascend/_310p/worker/v2/spec_decode/__init__.py new file mode 100644 index 000000000000..65680dc3fc48 --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/spec_decode/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. + +"""310P MRv2 speculative-decoding helpers.""" diff --git a/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py b/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py new file mode 100644 index 000000000000..1e232987b9e7 --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""310P draft ACLGraph manager: SpecDecoding capture + per-step decode FULL. + +K=1 draft-prefill FULL reuses target batch shape with ``q_len = 1+K = 2``. +Eager draft uses target SpecDecoding (splitfuse) metadata; if capture tags +``DecodeOnly`` via ``AscendInputBatch.make_dummy``, the graph records PA and +replay yields bad drafts (accept ~55% with intact final accuracy after +rejection). Mirror target ``ModelAclGraphManager310``: force SpecDecoding when +``num_tokens // num_reqs > 1`` during draft-prefill capture. + +K>1 draft-decode cannot record CPU slot_mapping D2H/H2D inside an NPU GLOBAL +graph. Align with MRv1: capture a **single** decode step, recompute slots on +the host between steps, then ``run_fullgraph`` per step. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import logger +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, + prepare_inputs_to_capture, +) +from vllm.v1.worker.gpu.input_batch import InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.utils import AttentionGroup + +from vllm_ascend._310p.worker.v2.spec_utils import set_draft_step_host +from vllm_ascend.worker.v2.aclgraph_utils import model_capture_wrapper +from vllm_ascend.worker.v2.input_batch import AscendInputBatch +from vllm_ascend.worker.v2.spec_decode.autoregressive.aclgraph import ( + AutoRegressiveAclGraphManager, +) +from vllm_ascend.worker.v2.utils import communicator_switch + + +class AutoRegressiveAclGraphManager310(AutoRegressiveAclGraphManager): + """310P draft FULL: SpecDecoding prefill + per-step decode graphs.""" + + def capture( + self, + forward_fn: Callable, + model_state: ModelState, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + progress_bar_desc: str = "Capturing CUDA graphs", + ) -> None: + if not self.is_draft_model_prefill: + self._capture_decode_per_step( + model_state, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + progress_bar_desc=progress_bar_desc, + ) + return + + if not self.cudagraph_mode.has_full_cudagraphs(): + return super().capture( + forward_fn, + model_state, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + progress_bar_desc=progress_bar_desc, + ) + + from vllm_ascend.attention.attention_v1 import AscendAttentionState + + orig_make_dummy = AscendInputBatch.make_dummy + + @classmethod + def make_dummy_draft_prefill( + cls, + num_reqs: int, + num_tokens: int, + input_buffers_arg: Any, + max_query_len: int | None = None, + ) -> AscendInputBatch: + kwargs: dict[str, Any] = {} + if max_query_len is not None: + kwargs["max_query_len"] = max_query_len + batch = orig_make_dummy(num_reqs, num_tokens, input_buffers_arg, **kwargs) + if num_reqs > 0 and (num_tokens // num_reqs) > 1: + batch.attn_state = AscendAttentionState.SpecDecoding + return batch + + AscendInputBatch.make_dummy = make_dummy_draft_prefill # type: ignore[method-assign] + try: + logger.info( + "Capturing 310P draft-prefill FULL with SpecDecoding make_dummy " + "(q_len>1 → splitfuse, not DecodeOnly/PA)." + ) + return super().capture( + forward_fn, + model_state, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + progress_bar_desc=progress_bar_desc, + ) + finally: + AscendInputBatch.make_dummy = orig_make_dummy # type: ignore[method-assign] + + def _capture_decode_per_step( + self, + model_state: ModelState, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + progress_bar_desc: str, + ) -> None: + """Capture one draft-decode step (q_len=1); runtime loops with host slots.""" + logger.info( + "Capturing 310P draft-decode FULL as per-step graphs " + "(CPU slot_mapping updated between steps, outside capture)." + ) + + with communicator_switch(), model_capture_wrapper(self.speculator, False): + + def create_forward_fn(desc: BatchExecutionDescriptor, warmup: bool): + del warmup + num_tokens = desc.num_tokens + num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) + num_tokens_across_dp = ( + torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu") + if self.dp_size > 1 + else None + ) + # Runs outside torch.cuda.graph / acl capture — D2H/H2D allowed. + prepare_inputs_to_capture( + num_reqs, + num_tokens, + model_state, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + full_cudagraph=(desc.cg_mode == CUDAGraphMode.FULL), + ) + seq_ub = input_buffers.seq_lens_cpu[:num_reqs] + # Stable device slot buffer; content refreshed before each replay. + slot_tensor = self.speculator.block_tables.get_dummy_slot_mappings(num_tokens) + slot_by_layer = build_slot_mappings_by_layer(slot_tensor, self.speculator.kv_cache_config) + attn_metadata = self.speculator._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=num_reqs, + num_tokens_padded=num_tokens, + seq_lens_cpu_upper_bound=seq_ub, + step=1, + ) + # Move capture-stable attn tensors to NPU before graph begin + # (pageable H2D inside capture is banned). + device = self.speculator.device + if attn_metadata is not None: + for meta in attn_metadata.values(): + if meta is None: + continue + seq_lens = getattr(meta, "seq_lens", None) + if seq_lens is not None and seq_lens.device != device: + meta.seq_lens = seq_lens.to(device=device, non_blocking=False) + torch.npu.current_stream().synchronize() + self.speculator.current_draft_step.fill_(1) + set_draft_step_host(1) + + def run(cg_mode: CUDAGraphMode) -> None: + del cg_mode + # Record model forward + sample + device-only draft write. + # Do NOT call update_draft_inputs_cpu / attn metadata H2D + # (pageable memcpy is banned under NPU GLOBAL capture). + self.speculator._prepare_eplb_forward(num_reqs) + last_hidden_states, _hidden_states = self.speculator._run_model( + num_tokens, + attn_metadata, + slot_by_layer, + num_tokens_across_dp, + CUDAGraphMode.NONE, + ) + last_hidden_states = last_hidden_states[:num_reqs] + positions = self.speculator.input_buffers.positions[:num_reqs] + idx_mapping = self.speculator.idx_mapping[:num_reqs] + draft_tokens = self.speculator.sample_draft( + last_hidden_states, + positions, + idx_mapping, + self.speculator.temperature, + self.speculator.seeds, + self.speculator.current_draft_step, + self.speculator.draft_logits, + ) + # Persist sampled drafts into the capture-stable buffer. + self.speculator.draft_tokens[:num_reqs, 1].copy_(draft_tokens[:num_reqs]) + + return run + + CudaGraphManager.capture(self, create_forward_fn, progress_bar_desc=progress_bar_desc) + + def run_fullgraph(self, desc: BatchExecutionDescriptor) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + num_tokens = desc.num_tokens + if self.is_draft_model_prefill: + logger.info_once( + "AutoRegressiveAclGraphManager310: draft prefill run_fullgraph with num_tokens=%s", + num_tokens, + ) + else: + logger.info_once( + "AutoRegressiveAclGraphManager310: draft decode per-step run_fullgraph with num_tokens=%s", + num_tokens, + ) + + # Ensure H2D into capture-stable buffers is visible before replay. + torch.npu.current_stream().synchronize() + ms = self.speculator.model_state + runtime_seq_lens = self.speculator.input_buffers.seq_lens + refresh = getattr(ms, "_refresh_capture_seq_lens", None) + if callable(refresh): + refresh(runtime_seq_lens) + + if self.is_draft_model_prefill: + return super().run_fullgraph(desc) + + # Per-step decode: parent builds multi-step FIA metadatas; 310P has no + # FIA graph_task for that path. Replay the captured single-step graph + # after host-side slot/seq updates (done by the speculator loop). + pending_attn = getattr(self.speculator, "_pending_draft_attn_metadata", None) + if pending_attn is not None: + ms.attn_metadata = pending_attn + return CudaGraphManager.run_fullgraph(self, desc) diff --git a/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py b/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py new file mode 100644 index 000000000000..b5a9968c262d --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""310P MTP speculator: CPU block-table slot mappings + RoPE flag + draft quant. + +Target path aligns with MRv1 concurrent uniform SpecDecoding FULL (hybrid +``prepare_attn`` actual/pad split). Draft-prefill FULL (K=1) uses +``AutoRegressiveAclGraphManager310`` with SpecDecoding capture (splitfuse). +K>1 draft-decode FULL uses per-step graphs: host slot_mapping between steps. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from vllm.config import VllmConfig, replace +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import logger +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +from vllm.v1.worker.gpu.spec_decode.mtp.speculator import MTPSpeculator + +from vllm_ascend._310p.ops.rotary_embedding import AscendRotaryEmbedding310 +from vllm_ascend._310p.worker.v2.spec_utils import set_draft_step_host +from vllm_ascend.worker.v2.spec_decode.autoregressive.speculator import ( + AscendAutoRegressiveSpeculator, +) + + +class AscendMTPSpeculator310(AscendAutoRegressiveSpeculator, MTPSpeculator): + """Ascend MTP speculator for 310P MRv2 (Triton-free draft loop).""" + + def _create_draft_vllm_config(self) -> VllmConfig: + draft_model_config = self.speculative_config.draft_model_config + if draft_model_config.hf_overrides is None: + draft_model_config.hf_overrides = {} + + parallel_config = replace( + self.vllm_config.parallel_config, + pipeline_parallel_size=1, + ) + draft_vllm_config = replace( + self.vllm_config, + model_config=draft_model_config, + parallel_config=parallel_config, + ) + + target_path = os.path.realpath(self.vllm_config.model_config.model) + draft_path = os.path.realpath(draft_model_config.model) + if target_path == draft_path and self.vllm_config.quant_config is not None: + draft_vllm_config = replace( + draft_vllm_config, + quant_config=self.vllm_config.quant_config, + ) + return draft_vllm_config + + def load_draft_model( + self, + target_model: nn.Module, + target_attn_layer_names: set[str], + ) -> nn.Module: + draft_model = load_eagle_model(target_model, self.draft_vllm_config) + spec_config = self.vllm_config.speculative_config + draft_hf_config = spec_config.draft_model_config.hf_config if spec_config is not None else None + self.share_mtp_topk_indices = ( + getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + and hasattr(draft_model.model, "set_skip_topk") + and hasattr(draft_model.model, "compact_topk_indices") + ) + return draft_model + + def _as_numpy_host(self, value: torch.Tensor | np.ndarray) -> np.ndarray: + if isinstance(value, np.ndarray): + return value.astype(np.int64, copy=False) + if value.device.type == "cpu": + return value.detach().numpy().astype(np.int64, copy=False) + # Sync D2H is illegal while an NPU stream is capturing (GLOBAL mode). + if torch.npu.is_current_stream_capturing(): + raise RuntimeError( + "310P draft slot_mapping cannot D2H while the NPU stream is capturing; " + "prepare host mirrors outside ACLGraph capture." + ) + return value.detach().cpu().numpy().astype(np.int64, copy=False) + + def _compute_draft_slot_mappings( + self, + idx_mapping: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + num_tokens_padded: int, + ) -> dict[str, torch.Tensor]: + idx_mapping_np = self._as_numpy_host(idx_mapping) + query_start_loc_np = self._as_numpy_host(query_start_loc) + positions_np = self._as_numpy_host(positions) + slot_mappings = self.block_tables.compute_slot_mappings( + idx_mapping_np, # type: ignore[arg-type] + query_start_loc_np, # type: ignore[arg-type] + positions_np, # type: ignore[arg-type] + num_tokens_padded=num_tokens_padded, + ) + return build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config) + + @contextmanager + def _rope_position_flag_310p(self): + AscendRotaryEmbedding310.set_rope_position_flag_310p(True) + try: + yield + finally: + AscendRotaryEmbedding310.set_rope_position_flag_310p(False) + + def capture(self) -> None: + """Capture draft-prefill FULL + per-step draft-decode FULL on 310P.""" + self.last_token_indices.zero_() + logger.info( + "Capturing 310P MTP draft ACLGraph (draft-prefill FULL + SpecDecoding; draft-decode per-step FULL)." + ) + super().capture() + + @torch.inference_mode() + def _run_model( + self, + num_tokens: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + with self._rope_position_flag_310p(): + return super()._run_model( + num_tokens, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + mm_inputs, + ) + + def _calc_next_seq_lens_cpu(self, seq_lens_cpu, num_reqs, num_reqs_padded, step): + """Match prepare_decode_inputs: seq = target - rejected + step.""" + next_seqs_cpu = seq_lens_cpu[:num_reqs_padded].clone() + rejected = getattr(self, "_last_num_rejected_cpu", None) + if rejected is not None and rejected.numel() >= num_reqs: + next_seqs_cpu[:num_reqs] = next_seqs_cpu[:num_reqs] - rejected[:num_reqs].to(next_seqs_cpu.dtype) + next_seqs_cpu = torch.clamp(next_seqs_cpu + step, max=self.max_model_len) + next_seqs_cpu[num_reqs:].fill_(0) + return next_seqs_cpu + + @torch.inference_mode() + def propose(self, *args, **kwargs): # type: ignore[no-untyped-def] + num_rejected = kwargs.get("num_rejected") + if num_rejected is None and len(args) >= 7: + num_rejected = args[6] + if isinstance(num_rejected, torch.Tensor): + self._last_num_rejected_cpu = num_rejected.detach().to("cpu") + return super().propose(*args, **kwargs) + + def _generate_draft( + self, + num_reqs: int, + num_tokens_padded: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> None: + """Skip Ascend metadata H2D under ACLGraph capture (pageable memcpy ban).""" + # Call GPU AR generate_draft (sample + update_draft_inputs) without the + # Ascend post-step ``seq_lens_cpu.copy_`` which is illegal while capturing. + from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( + AutoRegressiveSpeculator, + ) + + AutoRegressiveSpeculator._generate_draft( + self, + num_reqs, + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + if attn_metadata is None or torch.npu.is_current_stream_capturing(): + return + self._update_decode_attn_metadata(attn_metadata, 1, num_reqs) + + def _set_draft_step(self, step: int) -> None: + self.current_draft_step.fill_(step) + set_draft_step_host(step) + + def _multi_step_decode( + self, + num_reqs: int, + skip_attn: bool, + batch_desc: BatchExecutionDescriptor, + num_tokens_across_dp: torch.Tensor | None, + seq_lens_cpu_upper_bound: torch.Tensor | None = None, + ) -> None: + """K>1 draft decode: per-step FULL replay or eager CPU slot mappings.""" + assert seq_lens_cpu_upper_bound is not None + positions = self.input_buffers.positions[:num_reqs] + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] + idx_mapping = self.idx_mapping[:num_reqs] + seq_ub = seq_lens_cpu_upper_bound + rejected = getattr(self, "_last_num_rejected_cpu", None) + if rejected is not None and rejected.numel() >= num_reqs: + seq_ub = seq_lens_cpu_upper_bound.clone() + seq_ub[:num_reqs] = seq_ub[:num_reqs] - rejected[:num_reqs].to(seq_ub.dtype) + + use_full = batch_desc.cg_mode == CUDAGraphMode.FULL + attn_metadata = None + slot_mappings_by_layer = None + for step in range(1, self.num_speculative_steps): + if not skip_attn and (self.advance_draft_positions or step == 1): + # Host slot_mapping + attn metadata must run outside capture. + slot_mappings_by_layer = self._compute_draft_slot_mappings( + idx_mapping, + query_start_loc, + positions, + batch_desc.num_tokens, + ) + attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=batch_desc.num_reqs or num_reqs, + num_tokens_padded=batch_desc.num_tokens, + seq_lens_cpu_upper_bound=seq_ub, + step=step, + ) + if attn_metadata is not None: + for meta in attn_metadata.values(): + if meta is None: + continue + seq_lens = getattr(meta, "seq_lens", None) + if seq_lens is not None and seq_lens.device != self.device: + meta.seq_lens = seq_lens.to(device=self.device, non_blocking=False) + + self._set_draft_step(step) + if use_full: + assert self.decode_cudagraph_manager is not None + self._pending_draft_attn_metadata = attn_metadata + self.decode_cudagraph_manager.run_fullgraph(batch_desc) + if attn_metadata is not None: + self._update_decode_attn_metadata(attn_metadata, 1, num_reqs) + else: + self._generate_draft( + num_reqs, + batch_desc.num_tokens, + attn_metadata, + slot_mappings_by_layer, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) diff --git a/vllm_ascend/_310p/worker/v2/spec_utils.py b/vllm_ascend/_310p/worker/v2/spec_utils.py new file mode 100644 index 000000000000..b70d961dd80f --- /dev/null +++ b/vllm_ascend/_310p/worker/v2/spec_utils.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""CPU fallbacks for MRv2 spec-decode helpers (310P has no Triton).""" + +from __future__ import annotations + +import numpy as np +import torch +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers + + +def expand_idx_mapping_cpu( + idx_mapping: torch.Tensor, + total_num_logits: int, + cu_num_logits_np: np.ndarray, +) -> tuple[torch.Tensor, torch.Tensor]: + device = idx_mapping.device + expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty(total_num_logits, dtype=torch.int32, device=device) + for req_idx in range(cu_num_logits_np.shape[0] - 1): + start = int(cu_num_logits_np[req_idx]) + end = int(cu_num_logits_np[req_idx + 1]) + num_tokens = end - start + if num_tokens <= 0: + continue + expanded_idx_mapping[start:end] = idx_mapping[req_idx] + expanded_local_pos[start:end] = torch.arange(num_tokens, dtype=torch.int32, device=device) + return expanded_idx_mapping, expanded_local_pos + + +def combine_sampled_and_draft_tokens_cpu( + input_ids: torch.Tensor, + idx_mapping_np: np.ndarray, + last_sampled_tokens: torch.Tensor, + query_start_loc_np: np.ndarray, + seq_lens_np: np.ndarray, + prefill_len_np: np.ndarray, + draft_tokens: torch.Tensor, + cu_num_logits_np: np.ndarray, + num_logits: int, + num_new_sampled_tokens: int, + *, + device: torch.device, +) -> torch.Tensor: + del device + assert num_new_sampled_tokens in (0, 1) + num_reqs = idx_mapping_np.shape[0] + logits_indices = torch.empty(num_logits, dtype=torch.int64, device=input_ids.device) + # One small D2H for sampled/draft token tables (not the full input_ids buffer). + last_sampled_cpu = last_sampled_tokens.detach().cpu() + draft_cpu = draft_tokens.detach().cpu() + # Host staging only for rows we write; copy back via indexed slices. + writes: list[tuple[int, int]] = [] + host_vals: list[int] = [] + + for batch_idx in range(num_reqs): + req_state_idx = int(idx_mapping_np[batch_idx]) + cu_start = int(cu_num_logits_np[batch_idx]) + cu_end = int(cu_num_logits_np[batch_idx + 1]) + num_req_logits = cu_end - cu_start + num_draft_tokens = num_req_logits - num_new_sampled_tokens + + query_end = int(query_start_loc_np[batch_idx + 1]) + logits_start = query_end - num_req_logits + for offset in range(num_req_logits): + logits_indices[cu_start + offset] = logits_start + offset + + seq_len = int(seq_lens_np[batch_idx]) + prefill_len = int(prefill_len_np[batch_idx]) + if seq_len <= prefill_len: + continue + + first_logit_seq_pos = seq_len - num_req_logits + if num_new_sampled_tokens > 0 and first_logit_seq_pos >= prefill_len: + last_token_id = int(last_sampled_cpu[req_state_idx].item()) + writes.append((logits_start, logits_start + 1)) + host_vals.append(last_token_id) + + if num_draft_tokens > 0: + draft_row = draft_cpu[req_state_idx, :num_draft_tokens].tolist() + draft_start = query_end - num_draft_tokens + for i, tok in enumerate(draft_row): + writes.append((draft_start + i, draft_start + i + 1)) + host_vals.append(int(tok)) + + if host_vals: + # Avoid NPU index_copy_/IndexPutV2 (unreliable on some 310P layouts). + # Direct indexing matches gdn_310._merge_spec_and_non_spec_outputs_310. + idx = torch.tensor([s for s, _ in writes], dtype=torch.long, device=input_ids.device) + vals = torch.tensor(host_vals, dtype=input_ids.dtype, device=input_ids.device) + input_ids[idx] = vals + return logits_indices + + +def get_num_sampled_and_rejected_cpu( + num_sampled: torch.Tensor, + seq_lens: torch.Tensor, + cu_num_logits: torch.Tensor, + idx_mapping_np: np.ndarray, + prefill_len_np: np.ndarray, +) -> tuple[torch.Tensor, torch.Tensor]: + num_reqs = idx_mapping_np.shape[0] + num_rejected = torch.empty_like(num_sampled) + cu_np = cu_num_logits.detach().cpu().numpy() + seq_np = seq_lens.detach().cpu().numpy() + sampled_np = num_sampled.detach().cpu().numpy().copy() + + for batch_idx in range(num_reqs): + seq_len = int(seq_np[batch_idx]) + prefill_len_i = int(prefill_len_np[batch_idx]) + is_chunked_prefilling = seq_len < prefill_len_i + if is_chunked_prefilling: + sampled_np[batch_idx] = 0 + num_rejected[batch_idx] = 0 + continue + num_logits = int(cu_np[batch_idx + 1] - cu_np[batch_idx]) + num_rejected[batch_idx] = num_logits - int(sampled_np[batch_idx]) + + num_sampled_out = torch.from_numpy(sampled_np).to(device=num_sampled.device, dtype=num_sampled.dtype) + return num_sampled_out, num_rejected.to(device=num_sampled.device) + + +def greedy_rejection_sample_cpu( + target_logits: torch.Tensor, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + num_speculative_steps: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Greedy (temperature=0) rejection sampling for MTP verify. + + Argmax runs on-device so we only D2H token ids (not the full vocab logits). + Acceptance bookkeeping stays on CPU to avoid per-element NPU syncs. + """ + cu_np = cu_num_logits.detach().cpu().numpy() + num_reqs = cu_np.shape[0] - 1 + max_tokens = num_speculative_steps + 1 + sampled_cpu = torch.full((num_reqs, max_tokens), -1, dtype=torch.int32) + num_sampled_cpu = torch.zeros(num_reqs, dtype=torch.int32) + # Avoid full-vocab logits D2H (dominant cost on 310P MTP verify). + target_argmax_cpu = target_logits.argmax(dim=-1).to(dtype=torch.int32).detach().cpu().numpy() + draft_cpu = draft_sampled.detach().cpu().to(dtype=torch.int32).numpy() + + for req_idx in range(num_reqs): + start = int(cu_np[req_idx]) + end = int(cu_np[req_idx + 1]) + num_logits = end - start + if num_logits <= 0: + continue + accepted = 0 + # draft_sampled = input_ids[logits_indices] = [last_sampled, draft_0, ...] + # logits[i] predicts token i+1, which is draft_sampled[i+1] (upstream + # rejection_sampler_utils loads draft_sampled_ptr + logit_idx + 1). + for logit_idx in range(start, end): + target_token = int(target_argmax_cpu[logit_idx]) + is_bonus = logit_idx >= end - 1 + if accepted < num_speculative_steps and not is_bonus: + draft_token = int(draft_cpu[logit_idx + 1]) + if draft_token == target_token: + sampled_cpu[req_idx, accepted] = target_token + accepted += 1 + continue + sampled_cpu[req_idx, accepted] = target_token + accepted += 1 + break + if accepted == 0: + sampled_cpu[req_idx, 0] = int(target_argmax_cpu[start]) + accepted = 1 + num_sampled_cpu[req_idx] = accepted + + return ( + sampled_cpu.to(device=target_logits.device, non_blocking=True), + num_sampled_cpu.to(device=target_logits.device, non_blocking=True), + ) + + +def prepare_prefill_inputs_cpu( + last_token_indices: torch.Tensor, + current_draft_step: torch.Tensor, + input_buffers: InputBuffers, + input_batch: InputBatch, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + max_num_reqs: int, +) -> torch.Tensor: + """Build draft-prefill inputs with minimal host/device sync. + + Prefer on-device token/position copies and host np mirrors for metadata. + Full-buffer ``.cpu()`` of target/draft tensors was a major 310P MTP cost. + """ + del max_num_reqs + num_reqs = input_batch.num_reqs + query_start_loc_np = input_batch.query_start_loc_np + idx_mapping_np = input_batch.idx_mapping_np + seq_lens_np = getattr(input_batch, "seq_lens_np", None) + if seq_lens_np is None: + seq_lens_np = input_batch.seq_lens.detach().cpu().numpy() + + # Small metadata only (one sync each); avoid D2H of full token buffers. + num_sampled_np = num_sampled[:num_reqs].detach().cpu().numpy() + num_rejected_np = num_rejected[:num_reqs].detach().cpu().numpy() + last_sampled_np = last_sampled.detach().cpu().numpy() + next_prefill_np = next_prefill_tokens.detach().cpu().numpy() + + target_input_ids = input_batch.input_ids + target_positions = input_batch.positions + draft_input_ids = input_buffers.input_ids + draft_positions = input_buffers.positions + device = draft_input_ids.device + + last_token_indices_host = np.zeros(last_token_indices.shape[0], dtype=np.int64) + draft_qsl_host = np.zeros(input_buffers.query_start_loc.shape[0], dtype=np.int32) + draft_seq_host = np.zeros(input_buffers.seq_lens.shape[0], dtype=np.int32) + next_tokens_host = np.empty(num_reqs, dtype=np.int32) + + for req_idx in range(num_reqs): + req_state_idx = int(idx_mapping_np[req_idx]) + query_start = int(query_start_loc_np[req_idx]) + query_end = int(query_start_loc_np[req_idx + 1]) + query_len = query_end - query_start + seq_len = int(seq_lens_np[req_idx]) + query_len -= int(num_rejected_np[req_idx]) + + if int(num_sampled_np[req_idx]) > 0: + next_tokens_host[req_idx] = int(np.asarray(last_sampled_np[req_state_idx]).reshape(-1)[0]) + else: + next_tokens_host[req_idx] = int(np.asarray(next_prefill_np[req_state_idx]).reshape(-1)[0]) + + # After subtracting rejected tokens, only the kept prefix is valid + # (match upstream Triton prepare_prefill_inputs). + kept_end = query_start + query_len + if query_len > 1: + # 310P: NPU slice-assign can corrupt the destination tail element + # (see llm_base_proposer_310.set_inputs_first_pass). Save/restore + # draft_input_ids[kept_end-1] around the shift copy. + tail_idx = kept_end - 1 + tail_save = draft_input_ids[tail_idx].clone() + draft_input_ids[query_start : kept_end - 1].copy_( + target_input_ids[query_start + 1 : kept_end], + non_blocking=True, + ) + draft_input_ids[tail_idx] = tail_save + last_token_index = kept_end - 1 + last_token_indices_host[req_idx] = last_token_index + draft_positions[query_start:kept_end].copy_( + target_positions[query_start:kept_end], + non_blocking=True, + ) + draft_qsl_host[req_idx] = query_start + draft_seq_host[req_idx] = seq_len + + current_draft_step.fill_(0) + if num_reqs > 0: + query_end = int(query_start_loc_np[num_reqs]) + draft_qsl_host[num_reqs:] = query_end + draft_seq_host[num_reqs:] = 0 + last_token_indices_host[num_reqs:] = 0 + # Scatter next tokens without index_copy_ (310P IndexPutV2 issues). + next_tokens = torch.from_numpy(next_tokens_host).to(device=device, non_blocking=True) + last_idx = torch.from_numpy(last_token_indices_host[:num_reqs]).to( + device=device, dtype=torch.long, non_blocking=True + ) + draft_input_ids[last_idx] = next_tokens.to(dtype=draft_input_ids.dtype) + + input_buffers.query_start_loc.copy_( + torch.from_numpy(draft_qsl_host).to(device=device, non_blocking=True), + non_blocking=True, + ) + input_buffers.seq_lens.copy_( + torch.from_numpy(draft_seq_host).to(device=device, dtype=input_buffers.seq_lens.dtype, non_blocking=True), + non_blocking=True, + ) + last_token_indices.copy_( + torch.from_numpy(last_token_indices_host).to(device=device, dtype=last_token_indices.dtype, non_blocking=True), + non_blocking=True, + ) + return last_token_indices + + +def prepare_decode_inputs_cpu( + draft_tokens: torch.Tensor, + target_seq_lens: torch.Tensor, + num_rejected: torch.Tensor, + input_buffers: InputBuffers, + max_model_len: int, + max_num_reqs: int, + advance_draft_positions: bool = True, +) -> None: + """Prepare draft decode inputs with small host syncs (K>1 path).""" + del max_num_reqs + num_reqs = draft_tokens.shape[0] + device = input_buffers.input_ids.device + draft_np = draft_tokens[:num_reqs].detach().cpu().numpy() + target_seq_np = target_seq_lens[:num_reqs].detach().cpu().numpy() + rejected_np = num_rejected[:num_reqs].detach().cpu().numpy() + + input_ids_host = draft_np.astype(np.int64, copy=False) + seq_host = np.zeros(input_buffers.seq_lens.shape[0], dtype=np.int64) + qsl_host = np.arange(num_reqs + 1, dtype=np.int64) + if qsl_host.shape[0] < input_buffers.query_start_loc.shape[0]: + qsl_full = np.zeros(input_buffers.query_start_loc.shape[0], dtype=np.int64) + qsl_full[: qsl_host.shape[0]] = qsl_host + qsl_full[qsl_host.shape[0] :] = num_reqs + qsl_host = qsl_full + + for req_idx in range(num_reqs): + seq_len = int(target_seq_np[req_idx]) - int(rejected_np[req_idx]) + if advance_draft_positions: + seq_len = min(seq_len + 1, max_model_len) + seq_host[req_idx] = seq_len + + input_buffers.input_ids[:num_reqs].copy_( + torch.from_numpy(input_ids_host).to(device=device, dtype=input_buffers.input_ids.dtype, non_blocking=True), + non_blocking=True, + ) + if advance_draft_positions: + # positions += 1 on-device for the active rows (avoid full-buffer D2H). + pos = input_buffers.positions[:num_reqs] + pos.add_(1) + pos.clamp_max_(max_model_len - 1) + input_buffers.query_start_loc.copy_( + torch.from_numpy(qsl_host).to(device=device, dtype=input_buffers.query_start_loc.dtype, non_blocking=True), + non_blocking=True, + ) + input_buffers.seq_lens.copy_( + torch.from_numpy(seq_host).to(device=device, dtype=input_buffers.seq_lens.dtype, non_blocking=True), + non_blocking=True, + ) + + +# Host mirror for draft step. ``current_draft_step.item()`` is a sync D2H and +# is illegal under NPU GLOBAL ACLGraph capture; callers set this before fill_. +_DRAFT_STEP_HOST: int = 0 + + +def set_draft_step_host(step: int) -> None: + global _DRAFT_STEP_HOST + _DRAFT_STEP_HOST = int(step) + + +def update_draft_inputs_cpu( + draft_tokens: torch.Tensor, + current_draft_step: torch.Tensor, + hidden_states: torch.Tensor, + output_draft_tokens: torch.Tensor, + next_input_hidden_states: torch.Tensor, + input_buffers: InputBuffers, + num_reqs: int, + max_model_len: int, + num_speculative_steps: int, + advance_draft_positions: bool = True, +) -> None: + """Update draft buffers for the next step using on-device ops where possible.""" + # ``.item()`` is a sync D2H and is illegal under NPU GLOBAL ACLGraph capture. + if torch.npu.is_current_stream_capturing(): + step = _DRAFT_STEP_HOST + else: + step = int(current_draft_step.item()) + set_draft_step_host(step) + tokens = draft_tokens[:num_reqs] + output_draft_tokens[:num_reqs, step].copy_(tokens) + if step >= num_speculative_steps - 1: + return + input_buffers.input_ids[:num_reqs].copy_(tokens) + next_input_hidden_states[:num_reqs].copy_(hidden_states[:num_reqs]) + if advance_draft_positions: + pos = input_buffers.positions[:num_reqs] + pos.add_(1) + pos.clamp_max_(max_model_len - 1) + seq = input_buffers.seq_lens[:num_reqs] + seq.add_(1) + seq.clamp_max_(max_model_len) diff --git a/vllm_ascend/patch/worker/__init__.py b/vllm_ascend/patch/worker/__init__.py index 74461f70ef1a..787096ba0971 100644 --- a/vllm_ascend/patch/worker/__init__.py +++ b/vllm_ascend/patch/worker/__init__.py @@ -36,6 +36,7 @@ import vllm_ascend.patch.worker.patch_qwen3vl # noqa else: import vllm_ascend.patch.worker.patch_idex_310 # noqa + import vllm_ascend.patch.worker.patch_v2.patch_spec_decode_310 # noqa import vllm_ascend.patch.worker.patch_rejection_sampler # noqa import vllm_ascend.patch.worker.patch_kimi_k25 # noqa @@ -61,5 +62,16 @@ import vllm_ascend.patch.worker.patch_v2.patch_dflash_speculator # noqa import vllm_ascend.patch.worker.patch_v2.patch_dspark # noqa +# 310P: draft FULL must use AutoRegressiveAclGraphManager310 (no FIA graph_task). +# patch_eagle_speculator above installs the 910 manager; re-override here. +if not get_current_hardware_profile().supports(HardwareCapability.STANDARD_WORKER_PATCHES): + from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as _ar_spec + + from vllm_ascend._310p.worker.v2.spec_decode.aclgraph import ( + AutoRegressiveAclGraphManager310, + ) + + _ar_spec.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager310 + # only patch routed experts capture in main2main. import vllm_ascend.patch.worker.patch_routed_experts_capture # noqa diff --git a/vllm_ascend/patch/worker/patch_v2/patch_spec_decode_310.py b/vllm_ascend/patch/worker/patch_v2/patch_spec_decode_310.py new file mode 100644 index 000000000000..f64bd9881b31 --- /dev/null +++ b/vllm_ascend/patch/worker/patch_v2/patch_spec_decode_310.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# mypy: ignore-errors + +"""310P MRv2 patches: CPU spec-decode helpers + MTP routing (no mainline edits).""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as ar_speculator + +from vllm_ascend._310p.worker.v2.spec_utils import ( + prepare_decode_inputs_cpu, + prepare_prefill_inputs_cpu, + update_draft_inputs_cpu, +) +from vllm_ascend.worker.v2 import aclgraph_utils as aclgraph_utils_mod +from vllm_ascend.worker.v2 import spec_decode as ascend_spec_decode + +ar_speculator.prepare_prefill_inputs = prepare_prefill_inputs_cpu +ar_speculator.prepare_decode_inputs = prepare_decode_inputs_cpu +ar_speculator.update_draft_inputs = update_draft_inputs_cpu +# Draft ACLGraph manager override is applied in patch/worker/__init__.py +# *after* patch_eagle_speculator (which would otherwise win). + + +def _embed_input_ids(self: nn.Module, input_ids: torch.Tensor, **kwargs): + return self.original_model.embed_input_ids(input_ids, **kwargs) # type: ignore[operator] + + +# Draft ModelWithContext must forward embed_input_ids for Qwen MTP. +if not hasattr(aclgraph_utils_mod.ModelWithContext, "embed_input_ids"): + aclgraph_utils_mod.ModelWithContext.embed_input_ids = _embed_input_ids # type: ignore[attr-defined] + +_orig_init_speculator = ascend_spec_decode.init_speculator + + +def _init_speculator_310p(vllm_config, device): + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + if ( + speculative_config.method == "mtp" + and not speculative_config.use_gemma4_mtp() + and not speculative_config.use_step3p5_mtp() + ): + from vllm_ascend._310p.worker.v2.spec_decode.mtp_speculator import ( + AscendMTPSpeculator310, + ) + + return AscendMTPSpeculator310(vllm_config, device) + return _orig_init_speculator(vllm_config, device) + + +ascend_spec_decode.init_speculator = _init_speculator_310p + +# Rebind if model_runner already imported ``init_speculator`` by name. +try: + from vllm_ascend.worker.v2 import model_runner as _mr + + _mr.init_speculator = _init_speculator_310p +except Exception: + pass