613 lines
25 KiB
Python
613 lines
25 KiB
Python
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()
|