import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch import numpy as np import torch from vllm.model_executor.layers.attention import MLAAttention from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor from vllm_ascend.core.kv_cache_interface import AscendMLAAttentionSpec from vllm_ascend.worker.model_runner_v1 import NPUModelRunner class TestNPUModelRunnerAcceptedTokens(unittest.TestCase): @patch("vllm_ascend.worker.model_runner_v1.mamba_utils.postprocess_mamba_align_gpu") def test_postprocess_writes_accepted_counts_to_independent_snapshot(self, mock_postprocess): runner = NPUModelRunner.__new__(NPUModelRunner) runner.use_async_scheduling = True runner.speculative_config = object() runner.model_config = SimpleNamespace(is_hybrid=True) runner.cache_config = SimpleNamespace(mamba_cache_mode="align") runner.num_accepted_tokens = SimpleNamespace( cpu=torch.zeros(2, dtype=torch.int32), gpu=torch.zeros(2, dtype=torch.int32), ) persistent_counts = torch.ones(2, dtype=torch.int32) runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu_tensor=persistent_counts) runner.kv_cache_config = object() runner.compilation_config = SimpleNamespace(static_forward_context={}) runner.model = SimpleNamespace(get_mamba_state_copy_func=lambda: ()) runner.num_accepted_tokens_event = MagicMock() runner._get_mamba_bufs = MagicMock() runner._update_states_after_model_execute( torch.tensor([[10, 11, -1], [20, -1, -1]]), MagicMock(), ) self.assertIs( mock_postprocess.call_args.kwargs["num_accepted_tokens_cpu_tensor"], runner.num_accepted_tokens.cpu, ) self.assertIsNot( mock_postprocess.call_args.kwargs["num_accepted_tokens_cpu_tensor"], persistent_counts, ) runner.num_accepted_tokens_event.record.assert_called_once_with() @patch( "vllm_ascend.worker.model_runner_v1.GPUModelRunner._update_states_after_model_execute", autospec=True, ) def test_non_async_postprocess_delegates_to_upstream(self, mock_postprocess): runner = NPUModelRunner.__new__(NPUModelRunner) runner.use_async_scheduling = False output_token_ids = torch.tensor([[10, -1]]) scheduler_output = MagicMock() runner._update_states_after_model_execute(output_token_ids, scheduler_output) mock_postprocess.assert_called_once_with(runner, output_token_ids, scheduler_output) def test_remap_uses_snapshot_after_persistent_row_is_overwritten(self): runner = NPUModelRunner.__new__(NPUModelRunner) previous_counts = np.ones(16, dtype=np.int32) previous_counts[4] = 3 previous_counts[11] = 4 persistent_counts = np.ones(16, dtype=np.int32) runner.num_accepted_tokens = SimpleNamespace(np=previous_counts) runner.prev_positions = SimpleNamespace(np=np.array([11, -1, 4] + [-1] * 13, dtype=np.int64)) runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu=persistent_counts) runner.use_async_scheduling = True runner._sync_num_accepted_tokens(num_reqs=3, has_prev_mapping=True) np.testing.assert_array_equal(previous_counts[:3], [4, 1, 3]) np.testing.assert_array_equal(persistent_counts[:3], [4, 1, 3]) def test_async_without_previous_mapping_initializes_current_rows(self): runner = NPUModelRunner.__new__(NPUModelRunner) snapshot = np.array([0, 4, 3, 9], dtype=np.int32) persistent_counts = np.array([7, 8, 6, 5], dtype=np.int32) runner.num_accepted_tokens = SimpleNamespace(np=snapshot) runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu=persistent_counts) runner.use_async_scheduling = True runner._sync_num_accepted_tokens(num_reqs=3, has_prev_mapping=False) np.testing.assert_array_equal(snapshot, [1, 1, 1, 9]) np.testing.assert_array_equal(persistent_counts, [1, 1, 1, 5]) def test_non_async_sync_uses_condensed_input_batch_rows(self): runner = NPUModelRunner.__new__(NPUModelRunner) snapshot = np.array([4, 2, 9], dtype=np.int32) persistent_counts = np.array([2, 1, 7], dtype=np.int32) runner.num_accepted_tokens = SimpleNamespace(np=snapshot) runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu=persistent_counts) runner.use_async_scheduling = False runner._sync_num_accepted_tokens(num_reqs=2, has_prev_mapping=False) np.testing.assert_array_equal(snapshot, [2, 1, 9]) np.testing.assert_array_equal(persistent_counts, [2, 1, 7]) class TestNPUModelRunnerKVCache(unittest.TestCase): def _build_runner(self): runner = NPUModelRunner.__new__(NPUModelRunner) runner.device = torch.device("cpu") runner.use_sparse = False runner.use_sparse_c8 = False runner.use_compress = False runner.use_hybrid_blocks = False runner.hybrid_with_attn_and_mamba = False runner.sfa_dcp_replicated_indexer_size = 1 runner.runner_only_attn_layers = set() runner.is_kv_consumer = False runner.vllm_config = MagicMock() runner.vllm_config.kv_transfer_config = None runner.model_config = MagicMock() runner.model_config.use_mla = True backend = MagicMock() backend.get_kv_cache_shape.side_effect = lambda num_blocks, block_size, num_kv_heads, head_size: ( 2, num_blocks, block_size, num_kv_heads, head_size, ) runner.attn_backend = backend return runner def test_allocate_kv_cache_uses_layer_spec_for_draft_gqa(self): runner = self._build_runner() kv_cache_spec = FullAttentionSpec( block_size=16, num_kv_heads=8, head_size=64, head_size_v=64, dtype=torch.float16, ) kv_cache_config = KVCacheConfig( num_blocks=2, kv_cache_tensors=[KVCacheTensor(size=kv_cache_spec.page_size_bytes * 2, shared_by=["draft_attn"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["draft_attn"], kv_cache_spec=kv_cache_spec)], ) kv_cache_raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config) k_cache_raw, v_cache_raw = kv_cache_raw_tensors["draft_attn"] self.assertEqual(k_cache_raw.numel(), kv_cache_spec.page_size_bytes) self.assertEqual(v_cache_raw.numel(), kv_cache_spec.page_size_bytes) def test_reshape_kv_cache_uses_layer_spec_for_draft_gqa(self): runner = self._build_runner() kv_cache_spec = FullAttentionSpec( block_size=16, num_kv_heads=8, head_size=64, head_size_v=64, dtype=torch.float16, ) kv_cache_config = KVCacheConfig( num_blocks=2, kv_cache_tensors=[KVCacheTensor(size=kv_cache_spec.page_size_bytes * 2, shared_by=["draft_attn"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["draft_attn"], kv_cache_spec=kv_cache_spec)], ) kv_cache_raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config) runner._kv_cache_spec_attn_group_iterator = lambda: [ SimpleNamespace( kv_cache_spec=kv_cache_spec, backend=runner.attn_backend, layer_names=["draft_attn"], ) ] kv_caches = runner._reshape_kv_cache_tensors(kv_cache_config, kv_cache_raw_tensors) k_cache, v_cache = kv_caches["draft_attn"] self.assertEqual(k_cache.shape, (2, 16, 8, 64)) self.assertEqual(v_cache.shape, (2, 16, 8, 64)) @patch("vllm_ascend.worker.model_runner_v1.has_ec_transfer", return_value=False) @patch("vllm_ascend.worker.model_runner_v1.get_layers_from_vllm_config") def test_sparse_layer_without_indexer_allocates_only_mla_kv_cache( self, mock_get_layers, _mock_has_ec_transfer, ): runner = self._build_runner() runner.use_sparse = True runner.block_size = 16 runner.sparse_head_dim = (512, 64, 128) runner.kv_cache_dtype = torch.bfloat16 runner.shared_kv_cache_layers = {} runner.ascend_config = MagicMock() runner.model_config.hf_text_config = SimpleNamespace( kv_lora_rank=512, qk_rope_head_dim=64, ) runner.vllm_config.cache_config.cache_dtype = "auto" attn_module = MLAAttention.__new__(MLAAttention) torch.nn.Module.__init__(attn_module) attn_module.impl = SimpleNamespace(has_indexer=False, use_sparse_c8=False) layer_name = "model.layers.1.self_attn.attn" mock_get_layers.return_value = {layer_name: attn_module} spec = runner.get_kv_cache_spec()[layer_name] self.assertEqual(spec.sparse_head_dim, (512, 64, 0)) kv_cache_config = KVCacheConfig( num_blocks=2, kv_cache_tensors=[ KVCacheTensor( size=spec.page_size_bytes * 2, shared_by=[layer_name], ) ], kv_cache_groups=[ KVCacheGroupSpec( layer_names=[layer_name], kv_cache_spec=spec, ) ], ) raw_caches = runner._allocate_kv_cache_tensors(kv_cache_config) raw_k_cache, raw_v_cache = raw_caches[layer_name] self.assertEqual(raw_k_cache.numel(), 2 * 16 * 512 * 2) self.assertEqual(raw_v_cache.numel(), 2 * 16 * 64 * 2) def test_sparse_c8_replicated_indexer_allocation_matches_page_size(self): runner = self._build_runner() runner.use_sparse = True runner.use_sparse_c8 = True runner.c8_k_cache_dtype = torch.int8 runner.c8_k_scale_cache_dtype = torch.float16 runner.model_config.hf_text_config = SimpleNamespace(index_head_dim=128) layer_name = "model.layers.0.self_attn.attn" num_blocks = 2 dcp_size = 2 spec = AscendMLAAttentionSpec( block_size=16, num_kv_heads=1, head_size=704, sparse_head_dim=(576, 0, 128), dtype=torch.bfloat16, cache_dtype_str="auto", cache_sparse_sfa_c8=True, cache_sparse_li_c8=True, c8_k_cache_dtype=torch.int8, c8_k_scale_cache_dtype=torch.float16, sfa_dcp_replicated_indexer_size=dcp_size, ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[ KVCacheTensor( size=spec.page_size_bytes * num_blocks, shared_by=[layer_name], ) ], kv_cache_groups=[ KVCacheGroupSpec( layer_names=[layer_name], kv_cache_spec=spec, ) ], ) raw_caches = runner._allocate_kv_cache_tensors(kv_cache_config) raw_k_cache, raw_indexer_cache, raw_indexer_scale_cache = raw_caches[layer_name] self.assertEqual(raw_k_cache.numel(), num_blocks * 16 * 576) self.assertEqual(raw_indexer_cache.numel(), num_blocks * dcp_size * 16 * 128) self.assertEqual(raw_indexer_scale_cache.numel(), num_blocks * dcp_size * 16 * 2) self.assertEqual( raw_k_cache.numel() + raw_indexer_cache.numel() + raw_indexer_scale_cache.numel(), spec.page_size_bytes * num_blocks, ) def test_sparse_replicated_indexer_page_size_uses_expanded_storage_once(self): block_size = 16 k_head_dim = 512 v_head_dim = 64 index_head_dim = 128 dcp_size = 4 expected_head_size = k_head_dim + v_head_dim + index_head_dim * dcp_size spec = AscendMLAAttentionSpec( block_size=block_size, num_kv_heads=1, head_size=k_head_dim + v_head_dim + index_head_dim, sparse_head_dim=(k_head_dim, v_head_dim, index_head_dim), dtype=torch.bfloat16, cache_dtype_str="auto", sfa_dcp_replicated_indexer_size=dcp_size, ) self.assertEqual(spec.page_size_bytes, block_size * expected_head_size * 2) self.assertEqual( spec.sparse_kv_cache_ratio, ( expected_head_size / k_head_dim, expected_head_size / v_head_dim, expected_head_size / (index_head_dim * dcp_size), None, ), ) def test_sparse_replicated_indexer_only_expands_indexer_cache(self): runner = self._build_runner() runner.use_sparse = True runner.model_config.hf_text_config = SimpleNamespace(index_head_dim=128) runner._get_attention_kv_cache_dims = lambda _layer_name, _spec: (512, 64) runner.attn_backend.get_kv_cache_shape.side_effect = lambda num_blocks, block_size, num_kv_heads, head_size: ( num_blocks, block_size, num_kv_heads, head_size, ) layer_name = "model.layers.0.self_attn.attn" num_blocks = 2 dcp_size = 4 spec = AscendMLAAttentionSpec( block_size=16, num_kv_heads=1, head_size=704, sparse_head_dim=(512, 64, 128), dtype=torch.bfloat16, cache_dtype_str="auto", sfa_dcp_replicated_indexer_size=dcp_size, ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[ KVCacheTensor( size=spec.page_size_bytes * num_blocks, shared_by=[layer_name], ) ], kv_cache_groups=[ KVCacheGroupSpec( layer_names=[layer_name], kv_cache_spec=spec, ) ], ) raw_caches = runner._allocate_kv_cache_tensors(kv_cache_config) raw_k_cache, raw_v_cache, raw_indexer_cache = raw_caches[layer_name] self.assertEqual(raw_k_cache.numel(), num_blocks * 16 * 512 * 2) self.assertEqual(raw_v_cache.numel(), num_blocks * 16 * 64 * 2) self.assertEqual(raw_indexer_cache.numel(), num_blocks * dcp_size * 16 * 128 * 2) self.assertEqual( raw_k_cache.numel() + raw_v_cache.numel() + raw_indexer_cache.numel(), spec.page_size_bytes * num_blocks, ) runner._kv_cache_spec_attn_group_iterator = lambda: [ SimpleNamespace( kv_cache_spec=spec, backend=runner.attn_backend, layer_names=[layer_name], ) ] k_cache, v_cache, indexer_cache = runner._reshape_kv_cache_tensors( kv_cache_config, raw_caches, )[layer_name] self.assertEqual(k_cache.shape, (num_blocks, 16, 1, 512)) self.assertEqual(v_cache.shape, (num_blocks, 16, 1, 64)) self.assertEqual(indexer_cache.shape, (num_blocks * dcp_size, 16, 1, 128)) class TestNPUModelRunnerOutputTokenIds(unittest.TestCase): def _build_runner(self): runner = NPUModelRunner.__new__(NPUModelRunner) runner.device = torch.device("cpu") runner.vllm_config = MagicMock() runner.model_config = MagicMock() runner.use_compress = False return runner @patch("vllm_ascend.worker.model_runner_v1.get_ascend_config") @patch("vllm_ascend.worker.model_runner_v1.lmhead_tp_enable") def test_sample_updates_output_token_ids_before_sampler(self, mock_lmhead_tp_enable, mock_get_ascend_config): """Verify output_token_ids are updated before sampler is called""" mock_lmhead_tp_enable.return_value = False mock_ascend_config = MagicMock() mock_ascend_config.enable_reduce_sample = False mock_get_ascend_config.return_value = mock_ascend_config # Build input batch with historical sampled tokens input_batch = MagicMock() input_batch.sampling_metadata.output_token_ids = [ [1, 2, 3, -1], [4, 5, -1], ] input_batch.sampling_metadata.top_k = None input_batch.num_reqs = 2 input_batch.top_k_cpu = None input_batch.prev_req_id_to_index = { "req0": 0, "req1": 1, } input_batch.sampled_token_ids_cpu = torch.tensor([6, 7]) input_batch.async_copy_ready_event = MagicMock() input_batch.async_copy_ready_event.synchronize = MagicMock() # Simulate the real behavior of InputBatch.update_async_output_token_ids def mock_update_output_token_ids(): output_token_ids = input_batch.sampling_metadata.output_token_ids sampled_ids = input_batch.sampled_token_ids_cpu.tolist() for index, req_id in enumerate(input_batch.prev_req_id_to_index): prev_index = input_batch.prev_req_id_to_index[req_id] req_output = output_token_ids[index] if req_output and req_output[-1] == -1: req_output[-1] = sampled_ids[prev_index] input_batch.update_async_output_token_ids.side_effect = mock_update_output_token_ids # Build runner and inject dependencies runner = self._build_runner() runner.input_batch = input_batch runner.sampler = MagicMock(return_value=MagicMock()) # Call sample method logits = torch.randn(2, 32000) runner._sample(logits=logits, spec_decode_metadata=None) # Verify sampler and update_async_output_token_ids were called runner.sampler.assert_called_once() input_batch.update_async_output_token_ids.assert_called_once() # Verify output_token_ids were updated before sampler is called call_kwargs = runner.sampler.call_args[1] actual_sampling_metadata = call_kwargs["sampling_metadata"] actual_output_token_ids = actual_sampling_metadata.output_token_ids self.assertEqual(actual_output_token_ids[0], [1, 2, 3, 6]) self.assertEqual(actual_output_token_ids[1], [4, 5, 7]) def test_placeholder_spec_tokens_are_sanitized_only_for_forward(self): runner = self._build_runner() runner.input_ids = SimpleNamespace( cpu=torch.tensor([11, -1, 33, -1], dtype=torch.int32), gpu=torch.tensor([11, -1, 33, -1], dtype=torch.int32), ) scheduler_output = SimpleNamespace( scheduled_spec_decode_tokens={"req0": [-1]}, ) runner._sanitize_placeholder_input_ids_for_forward( scheduler_output, num_forward_tokens=4, ) self.assertEqual(runner.input_ids.gpu.tolist(), [11, 0, 33, 0]) self.assertEqual(runner.input_ids.cpu.tolist(), [11, -1, 33, -1]) def test_placeholder_sanitization_is_scoped_to_current_forward(self): runner = self._build_runner() runner.input_ids = SimpleNamespace( cpu=torch.tensor([11, -1, 33, -1], dtype=torch.int32), gpu=torch.tensor([11, -1, 33, -1], dtype=torch.int32), ) scheduler_output = SimpleNamespace( scheduled_spec_decode_tokens={"req0": [-1]}, ) runner._sanitize_placeholder_input_ids_for_forward( scheduler_output, num_forward_tokens=2, ) self.assertEqual(runner.input_ids.gpu.tolist(), [11, 0, 33, -1]) def test_mtp3_placeholder_metadata_is_preserved_before_sanitizing_forward(self): runner = self._build_runner() runner.pcp_size = 1 runner.arange_np = np.arange(8, dtype=np.int32) runner._arange_scratch = np.empty(8, dtype=np.int32) runner.input_ids = SimpleNamespace( cpu=torch.tensor([11, -1, -1, -1], dtype=torch.int32), gpu=torch.tensor([11, -1, -1, -1], dtype=torch.int32), ) scheduler_output = SimpleNamespace( scheduled_spec_decode_tokens={"req0": [-1, -1, -1]}, ) spec_decode_metadata = runner._calc_spec_decode_metadata( num_draft_tokens=np.array([3], dtype=np.int32), cu_num_scheduled_tokens=np.array([4], dtype=np.int32), num_pcp_pads=None, ) runner._sanitize_placeholder_input_ids_for_forward( scheduler_output, num_forward_tokens=4, ) self.assertEqual(spec_decode_metadata.draft_token_ids.tolist(), [-1, -1, -1]) self.assertEqual(runner.input_ids.gpu.tolist(), [11, 0, 0, 0]) self.assertEqual(runner.input_ids.cpu.tolist(), [11, -1, -1, -1]) class TestNPUModelRunnerDebugger(unittest.TestCase): def _build_runner(self, debugger=None): runner = NPUModelRunner.__new__(NPUModelRunner) runner.debugger = debugger or MagicMock() runner.model = MagicMock() runner.model_config = MagicMock() runner.model_config.enforce_eager = False runner._debugger_started = True runner._debugger_step_dummy_data_before_execute = False runner.use_compress = False return runner def test_finalize_dump_data_stops_stop_capable_debugger(self): runner = self._build_runner() runner._finalize_dump_data() runner.debugger.stop.assert_called_once_with() runner.debugger.step.assert_called_once_with() self.assertFalse(runner._debugger_started) def test_finalize_dump_data_steps_graph_debugger_without_stop(self): debugger = MagicMock(spec=["start", "step"]) runner = self._build_runner(debugger) runner._finalize_dump_data() debugger.step.assert_called_once_with() self.assertTrue(runner._debugger_started) def test_start_dump_data_noop_when_already_started(self): runner = self._build_runner(MagicMock(spec=["start", "step"])) runner._start_dump_data() runner.debugger.start.assert_not_called() runner.debugger.step.assert_not_called() self.assertTrue(runner._debugger_started) class TestCorrectOptimisticSeqLensCpu(unittest.TestCase): """Regression tests for async spec-decode seq_lens correction. The helper must synchronize the device->host copy event *before* reading ``valid_sampled_token_count_cpu``. Reading it early consumes stale counts and corrupts the CPU seq_lens, which surfaced as an accuracy regression on DeepSeek-V4 (its compressed-KV slot mapping is built from these seq_lens). """ def _build_runner(self, optimistic, prev_positions, prev_drafts, counts_cpu): runner = NPUModelRunner.__new__(NPUModelRunner) runner.optimistic_seq_lens_cpu = optimistic runner.prev_positions = SimpleNamespace(np=prev_positions) runner.prev_num_draft_tokens = SimpleNamespace(np=prev_drafts) runner.valid_sampled_token_count_cpu = counts_cpu return runner def test_synchronizes_before_host_read(self): num_reqs = 3 # Optimistic (all drafts assumed accepted): # prev_computed=[100,200,50], prev_drafts=[2,3,1], sched=[3,4,2] # optimistic = prev_computed + (prev_drafts + 1) + sched optimistic = torch.tensor([106, 208, 54], dtype=torch.int64) prev_positions = np.array([0, 1, 2], dtype=np.int64) prev_drafts = np.array([2, 3, 1], dtype=np.int32) # CPU buffer initially holds STALE counts (== drafts + 1, i.e. "all # accepted"). If the helper reads before synchronizing, the correction # is a no-op and the assertion below fails. counts_cpu = torch.tensor([3, 4, 2], dtype=torch.int32) # The true counts that the async copy delivers on synchronize(). true_counts = np.array([2, 1, 2], dtype=np.int32) runner = self._build_runner(optimistic, prev_positions, prev_drafts, counts_cpu) event = MagicMock() event.synchronize.side_effect = lambda: counts_cpu.copy_(torch.from_numpy(true_counts)) runner.valid_sampled_token_count_event = event runner._correct_optimistic_seq_lens_cpu(num_reqs) event.synchronize.assert_called_once() # correction = (prev_drafts + 1 - true_counts) = [1, 3, 0] # corrected = optimistic - correction = [105, 205, 54] np.testing.assert_array_equal(optimistic.numpy(), np.array([105, 205, 54])) def test_asserts_event_present(self): runner = self._build_runner( torch.tensor([10], dtype=torch.int64), np.array([0], dtype=np.int64), np.array([1], dtype=np.int32), torch.tensor([1], dtype=torch.int32), ) runner.valid_sampled_token_count_event = None with self.assertRaises(AssertionError): runner._correct_optimistic_seq_lens_cpu(1) if __name__ == "__main__": unittest.main()