diff --git a/.github/workflows/scripts/estimated_times.yaml b/.github/workflows/scripts/estimated_times.yaml index a7ab525c7927..344f072df7fc 100644 --- a/.github/workflows/scripts/estimated_times.yaml +++ b/.github/workflows/scripts/estimated_times.yaml @@ -15,6 +15,7 @@ estimated_times: 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: 600 + tests/e2e/pull_request/one_card/_310p/test_sampling_temp_310p.py: 400 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_sampling_temp_310p.py b/tests/e2e/pull_request/one_card/_310p/test_sampling_temp_310p.py new file mode 100644 index 000000000000..0bba27c1df57 --- /dev/null +++ b/tests/e2e/pull_request/one_card/_310p/test_sampling_temp_310p.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# 310P MRV2 e2e smoke for temperature / top-k / top-p post-process. +# Style mirrors MRV1 ``tests/e2e/pull_request/one_card/test_sampler.py`` +# (single generate with SamplingParams; no distributional asserts). + +from __future__ import annotations + +import os +from unittest.mock import patch + +from vllm import SamplingParams + +from tests.e2e.conftest import VllmRunner, wait_until_npu_memory_free + +# CI uses hub id; local verify can override (e.g. /home/weights/Qwen3-0.6B). +QWEN3_06B_MODEL = os.environ.get("QWEN3_06B_MODEL", "Qwen/Qwen3-0.6B") + + +@wait_until_npu_memory_free(0.7) +@patch.dict(os.environ, {"VLLM_USE_V2_MODEL_RUNNER": "1"}) +def test_qwen3_mrv2_310p_topk() -> None: + """MRV1 ``test_qwen3_topk`` counterpart on 310P MRV2. + + MRV1 uses ``temperature=0.0, top_k=50, top_p=0.9`` (vLLM clears top_k/top_p + when temperature==0). Here temperature>0 so Ascend310PSampler's + temperature / top-k / top-p path is actually exercised under ACLGraph. + """ + example_prompts = ["Hello, my name is"] + sampling_params = SamplingParams(max_tokens=5, temperature=0.8, top_k=50, top_p=0.9) + with VllmRunner( + QWEN3_06B_MODEL, + tensor_parallel_size=1, + enforce_eager=False, + dtype="float16", + max_model_len=2048, + max_num_seqs=8, + gpu_memory_utilization=0.8, + enable_prefix_caching=True, + additional_config={"ascend_compilation_config": {"fuse_norm_quant": False, "enable_npugraph_ex": False}}, + compilation_config={ + "cudagraph_mode": "FULL_DECODE_ONLY", + "cudagraph_capture_sizes": [1, 8], + }, + ) as vllm_model: + assert vllm_model.model.llm_engine.vllm_config.use_v2_model_runner + vllm_model.generate(example_prompts, sampling_params) diff --git a/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py b/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py index d757e9f8a4fe..cdb3d5c50be5 100644 --- a/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py +++ b/tests/ut/_310p/spec_decode/test_mtp_mrv2_310.py @@ -6,6 +6,7 @@ from types import SimpleNamespace from unittest.mock import patch +import numpy as np import torch from vllm.config.compilation import CUDAGraphMode @@ -15,6 +16,7 @@ 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, + probabilistic_rejection_sample_cpu, set_draft_step_host, update_draft_inputs_cpu, ) @@ -40,6 +42,103 @@ def test_greedy_rejection_uses_logit_idx_plus_one(self): self.assertEqual(num_sampled.tolist(), [2]) self.assertEqual(sampled[0, :2].tolist(), [1, 2]) + def test_probabilistic_rejection_accepts_when_u_below_p_draft(self): + """MRV1 IS_NGRAM: accept draft iff u < p(draft); forced u=0 always accepts.""" + # After softmax, token1 dominates first row → p(draft=1) high. + target_logits = torch.tensor( + [ + [0.0, 5.0, 0.0], + [0.0, 0.0, 5.0], + ], + dtype=torch.float32, + ) + draft_sampled = torch.tensor([9, 1], dtype=torch.int32) # draft at +1 is 1 + cu_num_logits = torch.tensor([0, 2], dtype=torch.int32) + temperature_np = np.array([0.8], dtype=np.float32) + idx_mapping_np = np.array([0], dtype=np.int32) + + with patch( + "vllm_ascend._310p.worker.v2.spec_utils._draw_uniform_cpu", + side_effect=[0.0, 0.99], # accept draft; bonus near end of CDF → token 2 + ): + sampled, num_sampled = probabilistic_rejection_sample_cpu( + target_logits, + draft_sampled, + cu_num_logits, + num_speculative_steps=1, + temperature_np=temperature_np, + idx_mapping_np=idx_mapping_np, + source_generators={}, + ) + + self.assertEqual(int(num_sampled[0].item()), 2) + self.assertEqual(int(sampled[0, 0].item()), 1) # accepted draft + self.assertEqual(int(sampled[0, 1].item()), 2) # bonus from last logit row + + def test_probabilistic_rejection_recovers_when_u_rejects_draft(self): + """u >= p(draft) → recovered token from residual (draft mass zeroed).""" + target_logits = torch.tensor( + [ + [0.0, 5.0, 4.0], # draft=1; residual keeps token 2 + [1.0, 0.0, 0.0], + ], + dtype=torch.float32, + ) + draft_sampled = torch.tensor([0, 1], dtype=torch.int32) + cu_num_logits = torch.tensor([0, 2], dtype=torch.int32) + temperature_np = np.array([1.0], dtype=np.float32) + idx_mapping_np = np.array([0], dtype=np.int32) + + with patch( + "vllm_ascend._310p.worker.v2.spec_utils._draw_uniform_cpu", + return_value=0.999, # always reject / pick high CDF + ): + sampled, num_sampled = probabilistic_rejection_sample_cpu( + target_logits, + draft_sampled, + cu_num_logits, + num_speculative_steps=1, + temperature_np=temperature_np, + idx_mapping_np=idx_mapping_np, + source_generators={}, + ) + + self.assertEqual(int(num_sampled[0].item()), 1) + # Recovered must not be the rejected draft token 1. + self.assertNotEqual(int(sampled[0, 0].item()), 1) + + def test_prepare_decode_inputs_advances_sample_src_positions(self): + """Upstream K>1 path passes sample_src_positions; CPU fallback must accept it.""" + from vllm_ascend._310p.worker.v2.spec_utils import prepare_decode_inputs_cpu + + num_reqs = 2 + draft_tokens = torch.tensor([7, 8], dtype=torch.int32) + target_seq_lens = torch.tensor([10, 12], dtype=torch.int32) + num_rejected = torch.tensor([0, 1], dtype=torch.int32) + sample_src_positions = torch.tensor([5, 6], dtype=torch.int64) + input_buffers = SimpleNamespace( + input_ids=torch.zeros(num_reqs, dtype=torch.int32), + positions=torch.tensor([4, 5], dtype=torch.int64), + query_start_loc=torch.zeros(num_reqs + 1, dtype=torch.int32), + seq_lens=torch.zeros(num_reqs, dtype=torch.int32), + ) + + prepare_decode_inputs_cpu( + draft_tokens, + target_seq_lens, + num_rejected, + input_buffers, + sample_src_positions, + max_model_len=128, + max_num_reqs=num_reqs, + advance_draft_positions=True, + ) + + self.assertEqual(input_buffers.input_ids.tolist(), [7, 8]) + self.assertEqual(sample_src_positions.tolist(), [6, 7]) + self.assertEqual(input_buffers.positions.tolist(), [5, 6]) + self.assertEqual(input_buffers.seq_lens.tolist(), [11, 12]) + def test_update_draft_inputs_uses_host_step_under_capture(self): num_reqs = 2 draft_tokens = torch.tensor([11, 22], dtype=torch.int32) @@ -47,6 +146,7 @@ def test_update_draft_inputs_uses_host_step_under_capture(self): 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) + sample_src_positions = torch.tensor([9, 10], dtype=torch.int64) input_buffers = SimpleNamespace( input_ids=torch.zeros(num_reqs, dtype=torch.int32), positions=torch.tensor([3, 5], dtype=torch.int64), @@ -62,6 +162,7 @@ def test_update_draft_inputs_uses_host_step_under_capture(self): output_draft_tokens=output_draft_tokens, next_input_hidden_states=next_input_hidden_states, input_buffers=input_buffers, + sample_src_positions=sample_src_positions, num_reqs=num_reqs, max_model_len=128, num_speculative_steps=2, @@ -70,6 +171,7 @@ def test_update_draft_inputs_uses_host_step_under_capture(self): self.assertEqual(output_draft_tokens[:, 0].tolist(), [11, 22]) self.assertEqual(input_buffers.positions.tolist(), [4, 6]) + self.assertEqual(sample_src_positions.tolist(), [10, 11]) def test_run_model_sets_rope_flag(self): flag_states: list[bool] = [] diff --git a/tests/ut/_310p/test_model_runner_v2_310p.py b/tests/ut/_310p/test_model_runner_v2_310p.py index 41975700d48d..56ffa4fa4aaf 100644 --- a/tests/ut/_310p/test_model_runner_v2_310p.py +++ b/tests/ut/_310p/test_model_runner_v2_310p.py @@ -382,11 +382,152 @@ def test_copy_kv_cache_blocks_flattens_mamba_lists() -> None: assert copies_arg is copies -def test_sampler_rejects_random_sampling_parameters() -> None: - sampler = Ascend310PSampler() +def test_sampler_accepts_temperature_and_rejects_penalties() -> None: + sampler = Ascend310PSampler(max_num_reqs=4, device="cpu", vocab_size=16) sampler.add_request(0, 4, SamplingParams(temperature=0)) + sampler.add_request(1, 4, SamplingParams(temperature=0.8, top_p=0.9, top_k=8, seed=7)) + assert sampler.sampling_states.temperature.gpu[1].item() == pytest.approx(0.8) + assert sampler.sampling_states.top_p.gpu[1].item() == pytest.approx(0.9) + assert int(sampler.sampling_states.top_k.gpu[1].item()) == 8 with pytest.raises(NotImplementedError, match="Unsupported sampling parameters"): - sampler.add_request(1, 4, SamplingParams(temperature=1)) + sampler.add_request(2, 4, SamplingParams(temperature=0, frequency_penalty=0.5)) + + +def test_sampler_temperature_scales_logits_before_argmax() -> None: + """Non-1 temperature must change relative logits before greedy/top paths.""" + from vllm_ascend._310p.worker.v2.sampler import _apply_temperature_pytorch + + logits = torch.tensor([[2.0, 4.0, 0.0], [1.0, 1.0, 1.0]], dtype=torch.float32) + expanded = torch.tensor([0, 1], dtype=torch.int64) + temperature = torch.tensor([0.5, 1.0], dtype=torch.float32) + _apply_temperature_pytorch(logits, expanded, temperature) + torch.testing.assert_close(logits[0], torch.tensor([4.0, 8.0, 0.0])) + torch.testing.assert_close(logits[1], torch.tensor([1.0, 1.0, 1.0])) + + +def test_sampler_greedy_call_returns_argmax() -> None: + sampler = Ascend310PSampler(max_num_reqs=2, device="cpu", vocab_size=4) + sampler.add_request(0, 2, SamplingParams(temperature=0)) + logits = torch.tensor([[0.1, 3.0, 0.2, 0.0]], dtype=torch.float32) + input_batch = SimpleNamespace( + expanded_idx_mapping=torch.tensor([0], dtype=torch.int32), + idx_mapping_np=np.array([0], dtype=np.int32), + num_reqs=1, + seq_lens=torch.ones(1, dtype=torch.int32), + ) + out = sampler(logits, input_batch) + assert int(out.sampled_token_ids.view(-1)[0].item()) == 1 + + +def _make_sampler_batch(req_idx: int = 0) -> SimpleNamespace: + return SimpleNamespace( + expanded_idx_mapping=torch.tensor([req_idx], dtype=torch.int32), + idx_mapping_np=np.array([req_idx], dtype=np.int32), + num_reqs=1, + seq_lens=torch.ones(1, dtype=torch.int32), + ) + + +def test_310p_mrv2_apply_top_k_top_p_masks_logits() -> None: + """Align with MRV1 ``tests/ut/sample/test_sampler.py`` mask checks. + + 310P MRV2 calls the same Triton-free ``apply_top_k_top_p`` path; assert + discarded logits become ``-inf`` and the kept set size matches top-k / top-p. + """ + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p + + # token0 > token1 >> rest + logits = torch.tensor([[5.0, 4.0, 1.0, 0.0, -1.0]], dtype=torch.float32) + k = torch.tensor([2], dtype=torch.int32) + p = torch.tensor([0.5], dtype=torch.float32) + + top_k_only = apply_top_k_top_p(logits.clone(), k, None) + finite_k = (top_k_only[0] > float("-inf")).nonzero(as_tuple=False).view(-1) + assert finite_k.tolist() == [0, 1] + assert top_k_only[0, 0] == pytest.approx(5.0) + assert top_k_only[0, 1] == pytest.approx(4.0) + assert not torch.isfinite(top_k_only[0, 2:]).any() + + # Mirror MRV1: combined filter keeps finite values and original shape. + filtered = apply_top_k_top_p(logits.clone(), k, p) + assert filtered.shape == logits.shape + assert torch.isfinite(filtered).any() + assert 0 in (filtered[0] > float("-inf")).nonzero(as_tuple=False).view(-1).tolist() + + top_p_only = apply_top_k_top_p(logits.clone(), None, p) + finite_p = (top_p_only[0] > float("-inf")).nonzero(as_tuple=False).view(-1) + assert 0 in finite_p.tolist() + assert finite_p.numel() >= 1 + assert finite_p.numel() < logits.shape[-1] + + +def test_sampler_top_k_restricts_softmax_mass() -> None: + """With temp>0, top-k must zero-out discarded tokens before inverse-CDF. + + vLLM clears top_k/top_p when temperature==0, so exercise the random path + with ``_random_sample_310p`` mocked (CPU-safe; mirrors MRV1 CDF unit style). + """ + import vllm_ascend._310p.worker.v2.sampler as sampler_mod + + sampler = Ascend310PSampler(max_num_reqs=1, device="cpu", vocab_size=5) + sampler.add_request(0, 2, SamplingParams(temperature=0.8, top_k=2, top_p=1.0, seed=7)) + # max at idx1, second at idx2; top_k=2 must keep only {1,2} + logits = torch.tensor([[1.0, 5.0, 4.0, 0.0, -2.0]], dtype=torch.float32) + captured: dict[str, torch.Tensor] = {} + + def _capture_and_argmax(probs: torch.Tensor, generators): + del generators + captured["probs"] = probs.detach().cpu().clone() + return probs.argmax(dim=-1) + + with patch.object(sampler_mod, "_random_sample_310p", side_effect=_capture_and_argmax): + out = sampler(logits, _make_sampler_batch()) + + probs = captured["probs"][0] + assert probs[3:].sum().item() == pytest.approx(0.0, abs=1e-6) + assert probs[0].item() == pytest.approx(0.0, abs=1e-6) + assert probs[1:3].sum().item() == pytest.approx(1.0, abs=1e-5) + assert int(out.sampled_token_ids.view(-1)[0].item()) == 1 + + +def test_sampler_top_k_one_matches_argmax_with_temperature() -> None: + """top_k=1 leaves one candidate; sampled token equals global argmax.""" + import vllm_ascend._310p.worker.v2.sampler as sampler_mod + + sampler = Ascend310PSampler(max_num_reqs=1, device="cpu", vocab_size=5) + sampler.add_request(0, 2, SamplingParams(temperature=0.9, top_k=1, top_p=1.0, seed=3)) + logits = torch.tensor([[0.2, 0.1, 3.0, 1.5, -1.0]], dtype=torch.float32) + + def _argmax_sample(probs: torch.Tensor, generators): + del generators + return probs.argmax(dim=-1) + + with patch.object(sampler_mod, "_random_sample_310p", side_effect=_argmax_sample): + out = sampler(logits, _make_sampler_batch()) + assert int(out.sampled_token_ids.view(-1)[0].item()) == int(logits.argmax(dim=-1).item()) + + +def test_sampler_top_p_restricts_softmax_mass() -> None: + """top_p must drop the long tail before sampling (MRV1-style mask contract).""" + import vllm_ascend._310p.worker.v2.sampler as sampler_mod + + sampler = Ascend310PSampler(max_num_reqs=1, device="cpu", vocab_size=5) + sampler.add_request(0, 2, SamplingParams(temperature=0.8, top_k=-1, top_p=0.5, seed=11)) + logits = torch.tensor([[5.0, 4.0, 1.0, 0.0, -1.0]], dtype=torch.float32) + captured: dict[str, torch.Tensor] = {} + + def _capture_and_argmax(probs: torch.Tensor, generators): + del generators + captured["probs"] = probs.detach().cpu().clone() + return probs.argmax(dim=-1) + + with patch.object(sampler_mod, "_random_sample_310p", side_effect=_capture_and_argmax): + out = sampler(logits, _make_sampler_batch()) + + probs = captured["probs"][0] + assert probs[0].item() == pytest.approx(1.0, abs=1e-5) or probs[:2].sum().item() == pytest.approx(1.0, abs=1e-5) + assert probs[2:].sum().item() == pytest.approx(0.0, abs=1e-5) + assert int(out.sampled_token_ids.view(-1)[0].item()) == 0 def test_block_tables_use_cpu_metadata_for_gather_and_slot_mapping() -> None: diff --git a/vllm_ascend/_310p/worker/v2/model_state.py b/vllm_ascend/_310p/worker/v2/model_state.py index 661207a4be4e..3060c49d8ca3 100644 --- a/vllm_ascend/_310p/worker/v2/model_state.py +++ b/vllm_ascend/_310p/worker/v2/model_state.py @@ -125,12 +125,16 @@ def prepare_inputs(self, input_batch: AscendInputBatch, req_states): return {"positions": positions} def custom_sampler(self, sampler): - del sampler # 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) + vocab_size = getattr(getattr(sampler, "sampling_states", None), "vocab_size", None) + if vocab_size is None: + model_config = getattr(self, "model_config", None) + if model_config is not None and hasattr(model_config, "get_vocab_size"): + vocab_size = model_config.get_vocab_size() + base_sampler = Ascend310PSampler(max_num_reqs, device, vocab_size=vocab_size) 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: diff --git a/vllm_ascend/_310p/worker/v2/rejection_sampler.py b/vllm_ascend/_310p/worker/v2/rejection_sampler.py index 65d470f6e773..9d20229535b2 100644 --- a/vllm_ascend/_310p/worker/v2/rejection_sampler.py +++ b/vllm_ascend/_310p/worker/v2/rejection_sampler.py @@ -2,10 +2,11 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. # mypy: ignore-errors -"""Greedy rejection sampler for 310P MRv2 MTP (no Triton).""" +"""MTP rejection sampler for 310P MRv2 (greedy + temperature, no Triton).""" from __future__ import annotations +import numpy as np import torch from vllm.config import SpeculativeConfig from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -13,12 +14,21 @@ from vllm_ascend._310p.worker.v2.spec_utils import ( get_num_sampled_and_rejected_cpu, greedy_rejection_sample_cpu, + probabilistic_rejection_sample_cpu, ) from vllm_ascend.worker.v2.input_batch import AscendInputBatch +_SAMPLING_EPS = 1e-5 + class RejectionSampler310V2: - """Greedy MTP rejection sampler for 310P MRv2.""" + """MTP rejection sampler for 310P MRv2. + + - ``temperature≈0``: greedy argmax verify (``greedy_rejection_sample_cpu``). + - ``temperature>0``: Leviathan / IS_NGRAM path aligned with MRV1 + ``AscendRejectionSampler310`` (accept iff ``u < p(draft)``, else recovered + token from residual; bonus via inverse-CDF). + """ def __init__( self, @@ -38,12 +48,39 @@ def __call__( ) -> SamplerOutput: del draft_logits draft_sampled = input_batch.input_ids[input_batch.logits_indices] - sampled, num_sampled = greedy_rejection_sample_cpu( + idx_mapping_np = input_batch.idx_mapping_np + expanded_idx_mapping = input_batch.expanded_idx_mapping + + processed = self.sampler.apply_sampling_params( logits, - draft_sampled, - input_batch.cu_num_logits, - self.num_speculative_steps, + expanded_idx_mapping, + idx_mapping_np, ) + + temperature_np = getattr(self.sampler, "_temperature_np", None) + if temperature_np is None: + temperature_np = np.zeros(int(idx_mapping_np.max(initial=0)) + 1, dtype=np.float32) + + any_random = bool(np.any(temperature_np[idx_mapping_np] >= _SAMPLING_EPS)) + if not any_random: + sampled, num_sampled = greedy_rejection_sample_cpu( + processed, + draft_sampled, + input_batch.cu_num_logits, + self.num_speculative_steps, + ) + else: + source_generators = getattr(self.sampler, "_source_generators", {}) + sampled, num_sampled = probabilistic_rejection_sample_cpu( + processed, + draft_sampled, + input_batch.cu_num_logits, + self.num_speculative_steps, + temperature_np, + idx_mapping_np, + source_generators, + ) + num_sampled, num_rejected = get_num_sampled_and_rejected_cpu( num_sampled, input_batch.seq_lens, diff --git a/vllm_ascend/_310p/worker/v2/sampler.py b/vllm_ascend/_310p/worker/v2/sampler.py index babe37d63a5d..fc2263c20b93 100644 --- a/vllm_ascend/_310p/worker/v2/sampler.py +++ b/vllm_ascend/_310p/worker/v2/sampler.py @@ -2,19 +2,54 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. # mypy: ignore-errors +"""Triton-free MRV2 sampler for Ascend 310P. + +Aligns temperature / top-k / top-p post-processing with MRV1's +``AscendSampler310`` (inverse-CDF random sampling) while keeping the MRV2 +``sampling_states`` surface required by MTP draft ``propose()``. +""" + +from __future__ import annotations + from types import SimpleNamespace +import numpy as np import torch from vllm.sampling_params import SamplingParams +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm_ascend._310p.sample.sampler import ( + _prepare_cpu_generators_310p, + _random_sample_310p, +) + +_NP_INT64_MIN = np.iinfo(np.int64).min +_NP_INT64_MAX = np.iinfo(np.int64).max +_SAMPLING_EPS = 1e-5 + + +def _apply_temperature_pytorch( + logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + temperature: torch.Tensor, +) -> None: + """In-place temperature scaling (Triton-free; skips temp in {0, 1}).""" + temps = temperature[expanded_idx_mapping].to(dtype=torch.float32) + need_scale = (temps != 0.0) & (temps != 1.0) + if not bool(need_scale.any().item()): + return + # Work in FP32 for numerical stability, write back to logits dtype. + scaled = logits.to(torch.float32) / temps.unsqueeze(-1) + logits.copy_(torch.where(need_scale.unsqueeze(-1), scaled.to(logits.dtype), logits)) + class Ascend310PSampler: - """Triton-free sampler for 310P MRV2. + """Triton-free sampler for 310P MRV2 with MRV1-aligned temperature sampling. - 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. + Exposes ``sampling_states.temperature/seeds/top_k/top_p`` so MTP draft + ``propose()`` (and ``_dummy_run``) can read device buffers without UVA + SamplingStates. """ # TODO: Refactor this sampler to register 310P implementations through @@ -24,21 +59,39 @@ def __init__( self, max_num_reqs: int = 1, device: torch.device | str | None = None, + vocab_size: int | None = None, ) -> None: self.penalties_state = SimpleNamespace(output_bin_counts=None) self.max_num_reqs = max_num_reqs + self.vocab_size = int(vocab_size) if vocab_size is not None else 0 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) + top_k_gpu = torch.full( + (max_num_reqs,), + fill_value=max(self.vocab_size, 1), + dtype=torch.int32, + device=device, + ) + top_p_gpu = torch.ones(max_num_reqs, dtype=torch.float32, device=device) self.sampling_states = SimpleNamespace( temperature=SimpleNamespace(gpu=temperature_gpu), seeds=SimpleNamespace(gpu=seeds_gpu), + top_k=SimpleNamespace(gpu=top_k_gpu), + top_p=SimpleNamespace(gpu=top_p_gpu), + vocab_size=self.vocab_size, ) + self._temperature_np = np.zeros(max_num_reqs, dtype=np.float32) + self._top_k_np = np.full(max_num_reqs, fill_value=max(self.vocab_size, 1), dtype=np.int32) + self._top_p_np = np.ones(max_num_reqs, dtype=np.float32) + self._seeds_set = np.zeros(max_num_reqs, dtype=bool) + # Source generators keyed by req_idx (MRV1 CPU RNG cache semantics). + self._source_generators: dict[int, torch.Generator] = {} def add_request( self, @@ -48,12 +101,6 @@ def add_request( ) -> None: del prompt_len unsupported = [] - if sampling_params.temperature != 0: - unsupported.append("temperature") - if sampling_params.top_p != 1.0: - unsupported.append("top_p") - if sampling_params.top_k not in (-1, 0): - unsupported.append("top_k") if sampling_params.min_p != 0.0: unsupported.append("min_p") if sampling_params.repetition_penalty != 1.0: @@ -69,21 +116,160 @@ def add_request( ): unsupported.append("logits processors") if unsupported: - # TODO: Support additional sampling features in the next 310P MRV2 iteration. + # TODO: Support penalties / logprobs / min_p in a later 310P MRV2 iteration. 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) + + if not (0 <= req_idx < self.max_num_reqs): + return + + temperature = float(sampling_params.temperature) + top_p = float(sampling_params.top_p) + top_k = int(sampling_params.top_k) + if self.vocab_size <= 0: + # Late-bind from logits on first sample if caller omitted vocab_size. + vocab_size = max(top_k, 1) + else: + vocab_size = self.vocab_size + if top_k <= 0 or top_k > vocab_size: + top_k = vocab_size + + self._temperature_np[req_idx] = temperature + self._top_p_np[req_idx] = top_p + self._top_k_np[req_idx] = top_k + self.sampling_states.temperature.gpu[req_idx] = temperature + self.sampling_states.top_p.gpu[req_idx] = top_p + self.sampling_states.top_k.gpu[req_idx] = top_k + + seed = getattr(sampling_params, "seed", None) + self._seeds_set[req_idx] = seed is not None + if seed is None: + seed = int(np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)) + self.sampling_states.seeds.gpu[req_idx] = int(seed) + + if self._seeds_set[req_idx]: + source = torch.Generator(device="cpu") + source.manual_seed(int(seed)) + self._source_generators[req_idx] = source + else: + self._source_generators.pop(req_idx, None) def apply_staged_writes(self) -> None: + # Params are written directly to device buffers in ``add_request``. pass + def _maybe_bind_vocab_size(self, vocab_size: int) -> None: + if self.vocab_size > 0 or vocab_size <= 0: + return + self.vocab_size = int(vocab_size) + self.sampling_states.vocab_size = self.vocab_size + # Replace unset top_k defaults (initialized to 1 when vocab was unknown). + unset = self._top_k_np <= 1 + self._top_k_np[unset] = self.vocab_size + self.sampling_states.top_k.gpu[unset] = self.vocab_size + + def _batch_needs_logits_processing(self, idx_mapping_np: np.ndarray) -> bool: + temp = self._temperature_np[idx_mapping_np] + return bool( + np.any((temp != 0.0) & (temp != 1.0)) + or np.any(self._top_k_np[idx_mapping_np] != max(self.vocab_size, 1)) + or np.any(self._top_p_np[idx_mapping_np] != 1.0) + ) + + def _any_random(self, idx_mapping_np: np.ndarray) -> bool: + return bool(np.any(self._temperature_np[idx_mapping_np] >= _SAMPLING_EPS)) + + def apply_sampling_params( + self, + logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + ) -> torch.Tensor: + """FP32 copy + temperature + top-k/top-p for MTP rejection verify. + + Mirrors upstream ``Sampler.apply_sampling_params`` (without penalties / + logit bias) so ``RejectionSampler310V2`` matches MRV1 tempered verify. + """ + self._maybe_bind_vocab_size(logits.shape[-1]) + processed = torch.empty_like(logits, dtype=torch.float32).copy_(logits) + if not self._batch_needs_logits_processing(idx_mapping_np): + return processed + _apply_temperature_pytorch( + processed, + expanded_idx_mapping, + self.sampling_states.temperature.gpu, + ) + do_top_k = np.any(self._top_k_np[idx_mapping_np] != max(self.vocab_size, 1)) + do_top_p = np.any(self._top_p_np[idx_mapping_np] != 1.0) + top_k = self.sampling_states.top_k.gpu[expanded_idx_mapping] if do_top_k else None + top_p = self.sampling_states.top_p.gpu[expanded_idx_mapping] if do_top_p else None + return apply_top_k_top_p(processed, top_k, top_p) + + def _build_row_generators( + self, + expanded_idx_mapping_np: np.ndarray, + ) -> dict[int, torch.Generator]: + """Map logit-row index → CPU generator for seeded random requests.""" + sources: dict[int, torch.Generator] = {} + for row_idx, req_idx in enumerate(expanded_idx_mapping_np.tolist()): + if req_idx < 0: + continue + if self._temperature_np[req_idx] < _SAMPLING_EPS: + continue + source = self._source_generators.get(int(req_idx)) + if source is not None: + sources[row_idx] = source + return _prepare_cpu_generators_310p(sources) + def __call__(self, logits: torch.Tensor, input_batch) -> SamplerOutput: - sampled = logits.argmax(dim=-1).to(torch.int32) - num_sampled = input_batch.seq_lens.new_ones(input_batch.num_reqs) + self._maybe_bind_vocab_size(logits.shape[-1]) + + expanded_idx_mapping = input_batch.expanded_idx_mapping + idx_mapping_np = input_batch.idx_mapping_np + num_reqs = input_batch.num_reqs + + # Fast path: all greedy and no top-k/top-p → argmax (previous behavior). + if not self._any_random(idx_mapping_np) and not self._batch_needs_logits_processing(idx_mapping_np): + sampled = logits.argmax(dim=-1).to(torch.int32) + num_sampled = input_batch.seq_lens.new_ones(num_reqs) + return SamplerOutput( + sampled_token_ids=sampled.view(-1, 1), + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + num_rejected=torch.zeros_like(num_sampled), + ) + + # Copy to FP32 before in-place temperature / top-k / top-p, matching + # upstream MRV2 Sampler.apply_sampling_params. + processed = torch.empty_like(logits, dtype=torch.float32).copy_(logits) + + if self._batch_needs_logits_processing(idx_mapping_np): + _apply_temperature_pytorch( + processed, + expanded_idx_mapping, + self.sampling_states.temperature.gpu, + ) + do_top_k = np.any(self._top_k_np[idx_mapping_np] != max(self.vocab_size, 1)) + do_top_p = np.any(self._top_p_np[idx_mapping_np] != 1.0) + top_k = self.sampling_states.top_k.gpu[expanded_idx_mapping] if do_top_k else None + top_p = self.sampling_states.top_p.gpu[expanded_idx_mapping] if do_top_p else None + processed = apply_top_k_top_p(processed, top_k, top_p) + + if not self._any_random(idx_mapping_np): + sampled = processed.argmax(dim=-1).to(torch.int32) + else: + # Greedy rows keep argmax; random rows use MRV1 inverse-CDF sampling. + greedy = processed.argmax(dim=-1) + probs = processed.softmax(dim=-1, dtype=torch.float32) + expanded_np = expanded_idx_mapping.detach().to("cpu").numpy() + generators = self._build_row_generators(expanded_np) + random_sampled = _random_sample_310p(probs, generators) + is_random = self.sampling_states.temperature.gpu[expanded_idx_mapping] >= _SAMPLING_EPS + sampled = torch.where(is_random, random_sampled, greedy).to(torch.int32) + + num_sampled = input_batch.seq_lens.new_ones(num_reqs) return SamplerOutput( sampled_token_ids=sampled.view(-1, 1), logprobs_tensors=None, diff --git a/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py b/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py index 1e232987b9e7..afb82c0e070f 100644 --- a/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py +++ b/vllm_ascend/_310p/worker/v2/spec_decode/aclgraph.py @@ -194,11 +194,11 @@ def run(cg_mode: CUDAGraphMode) -> None: CUDAGraphMode.NONE, ) last_hidden_states = last_hidden_states[:num_reqs] - positions = self.speculator.input_buffers.positions[:num_reqs] + sample_src_positions = self.speculator.sample_src_positions[:num_reqs] idx_mapping = self.speculator.idx_mapping[:num_reqs] draft_tokens = self.speculator.sample_draft( last_hidden_states, - positions, + sample_src_positions, idx_mapping, self.speculator.temperature, self.speculator.seeds, diff --git a/vllm_ascend/_310p/worker/v2/spec_utils.py b/vllm_ascend/_310p/worker/v2/spec_utils.py index b70d961dd80f..e2c2342102a1 100644 --- a/vllm_ascend/_310p/worker/v2/spec_utils.py +++ b/vllm_ascend/_310p/worker/v2/spec_utils.py @@ -122,6 +122,10 @@ def get_num_sampled_and_rejected_cpu( return num_sampled_out, num_rejected.to(device=num_sampled.device) +_SAMPLING_EPS = 1e-5 +_UNIFORM_TINY = float(torch.finfo(torch.float32).tiny) + + def greedy_rejection_sample_cpu( target_logits: torch.Tensor, draft_sampled: torch.Tensor, @@ -175,6 +179,123 @@ def greedy_rejection_sample_cpu( ) +def _draw_uniform_cpu(generator: torch.Generator | None) -> float: + if generator is None: + u = torch.rand((), dtype=torch.float32) + else: + u = torch.rand((), dtype=torch.float32, generator=generator) + return float(u.clamp_min_(_UNIFORM_TINY).item()) + + +def _sample_from_probs_row_cpu(probs_row: torch.Tensor, generator: torch.Generator | None) -> int: + """Inverse-CDF sample one token from a 1-D probability row (CPU).""" + row = probs_row.detach().to(dtype=torch.float32, device="cpu") + total = float(row.sum().item()) + if total <= 0.0: + return int(row.argmax().item()) + u = _draw_uniform_cpu(generator) * total + cdf = row.cumsum(dim=-1) + return int(torch.searchsorted(cdf, torch.tensor([u], dtype=torch.float32), right=True).item()) + + +def probabilistic_rejection_sample_cpu( + target_logits: torch.Tensor, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + num_speculative_steps: int, + temperature_np: np.ndarray, + idx_mapping_np: np.ndarray, + source_generators: dict[int, torch.Generator], +) -> tuple[torch.Tensor, torch.Tensor]: + """MTP rejection with temperature (Leviathan, draft one-hot / IS_NGRAM). + + Aligns with MRV1 ``rejection_random_sample_pytorch`` when ``draft_probs`` is + None: accept draft iff ``u < p_target(draft)``; on reject sample a recovered + token from the residual distribution; bonus token sampled from the last + logit row. Greedy requests (``temperature≈0``) keep argmax semantics. + + ``target_logits`` must already include temperature / top-k / top-p processing. + """ + from vllm_ascend._310p.sample.sampler import _prepare_cpu_generators_310p + + 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) + + # Softmax once on-device; per-step gathers stay cheap for small K. + probs = torch.softmax(target_logits, dim=-1, dtype=torch.float32) + 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() + + # Prepare CPU RNG keyed by request-state index (MRV1 generator cache). + sources = {int(req): gen for req, gen in source_generators.items()} + prepared = _prepare_cpu_generators_310p(sources) if sources else {} + + for batch_idx in range(num_reqs): + start = int(cu_np[batch_idx]) + end = int(cu_np[batch_idx + 1]) + if end <= start: + continue + req_state_idx = int(idx_mapping_np[batch_idx]) + is_greedy = float(temperature_np[req_state_idx]) < _SAMPLING_EPS + gen = prepared.get(req_state_idx) + + accepted = 0 + for logit_idx in range(start, end): + is_bonus = logit_idx >= end - 1 + if is_greedy: + target_token = int(target_argmax_cpu[logit_idx]) + if accepted < num_speculative_steps and not is_bonus: + draft_token = int(draft_cpu[logit_idx + 1]) + if draft_token == target_token: + sampled_cpu[batch_idx, accepted] = target_token + accepted += 1 + continue + sampled_cpu[batch_idx, accepted] = target_token + accepted += 1 + break + + # Probabilistic (MRV1 IS_NGRAM / MTP): u < p(draft). + if accepted < num_speculative_steps and not is_bonus: + draft_token = int(draft_cpu[logit_idx + 1]) + if draft_token < 0: + # Invalid pad draft → force recover/bonus path. + row = probs[logit_idx].detach().cpu() + sampled_cpu[batch_idx, accepted] = _sample_from_probs_row_cpu(row, gen) + accepted += 1 + break + p_draft = float(probs[logit_idx, draft_token].item()) + u = _draw_uniform_cpu(gen) + if u < p_draft: + sampled_cpu[batch_idx, accepted] = draft_token + accepted += 1 + continue + # Reject: sample recovered token from residual (zero draft mass). + row = probs[logit_idx].detach().cpu().clone() + row[draft_token] = 0.0 + sampled_cpu[batch_idx, accepted] = _sample_from_probs_row_cpu(row, gen) + accepted += 1 + break + + # Bonus token (all drafts accepted) or final logit. + row = probs[logit_idx].detach().cpu() + sampled_cpu[batch_idx, accepted] = _sample_from_probs_row_cpu(row, gen) + accepted += 1 + break + + if accepted == 0: + sampled_cpu[batch_idx, 0] = int(target_argmax_cpu[start]) + accepted = 1 + num_sampled_cpu[batch_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, @@ -285,11 +406,16 @@ def prepare_decode_inputs_cpu( target_seq_lens: torch.Tensor, num_rejected: torch.Tensor, input_buffers: InputBuffers, + sample_src_positions: torch.Tensor, 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).""" + """Prepare draft decode inputs with small host syncs (K>1 path). + + Signature matches upstream ``prepare_decode_inputs`` (incl. + ``sample_src_positions``) so patched call sites stay compatible. + """ del max_num_reqs num_reqs = draft_tokens.shape[0] device = input_buffers.input_ids.device @@ -316,6 +442,8 @@ def prepare_decode_inputs_cpu( torch.from_numpy(input_ids_host).to(device=device, dtype=input_buffers.input_ids.dtype, non_blocking=True), non_blocking=True, ) + # Align with Triton: always advance the draft sampling key for decode steps. + sample_src_positions[:num_reqs].add_(1) if advance_draft_positions: # positions += 1 on-device for the active rows (avoid full-buffer D2H). pos = input_buffers.positions[:num_reqs] @@ -348,12 +476,17 @@ def update_draft_inputs_cpu( output_draft_tokens: torch.Tensor, next_input_hidden_states: torch.Tensor, input_buffers: InputBuffers, + sample_src_positions: torch.Tensor, 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.""" + """Update draft buffers for the next step using on-device ops where possible. + + Signature matches upstream ``update_draft_inputs`` (incl. + ``sample_src_positions``). + """ # ``.item()`` is a sync D2H and is illegal under NPU GLOBAL ACLGraph capture. if torch.npu.is_current_stream_capturing(): step = _DRAFT_STEP_HOST @@ -364,6 +497,8 @@ def update_draft_inputs_cpu( output_draft_tokens[:num_reqs, step].copy_(tokens) if step >= num_speculative_steps - 1: return + # Align with Triton: advance sampling key before preparing next draft inputs. + sample_src_positions[:num_reqs].add_(1) input_buffers.input_ids[:num_reqs].copy_(tokens) next_input_hidden_states[:num_reqs].copy_(hidden_states[:num_reqs]) if advance_draft_positions: