diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index c7ca1891b642..eb03e9d73335 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -888,7 +888,9 @@ def _add_layer( # the typical step. An all-generation typical_step over-provisions the # compressed-cache pool at the expense of the SWA pool, starving the # SWA pool and artificially capping the achievable batch size. - ctx_capacity = max_num_tokens if max_num_tokens is not None else typical_seq_len + ctx_capacity = ( + max_num_tokens if max_num_tokens is not None else typical_seq_len + ) + self.num_extra_kv_tokens generation_history_length = max(0, typical_seq_len - max_draft_len - 1) typical_step = BatchDesc( kv_caches=[ diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..a76a6ac90706 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -39,6 +39,27 @@ KeyType: TypeAlias = Tuple[int, int, bool, bool, bool] +def _save_spec_decode_capture_state( + attn_metadata: Any, enable_spec_decode: bool) -> Optional[torch.Tensor]: + if not enable_spec_decode or not hasattr(attn_metadata, 'kv_lens_cuda'): + return None + return attn_metadata.kv_lens_cuda[:attn_metadata.num_seqs].clone() + + +def _restore_spec_decode_capture_state( + attn_metadata: Any, saved_kv_lens_cuda: Optional[torch.Tensor]) -> None: + if saved_kv_lens_cuda is None: + return + # Speculative decoding updates kv_lens_cuda in-place during every forward. + # CUDA graph warmup reuses one dummy request for multiple eager forwards, so + # letting those updates accumulate would make later warmups/capture advertise + # more KV tokens than the dummy request actually allocated. Restore the + # single-step input state outside the graph after each forward instead. + batch_size = saved_kv_lens_cuda.shape[0] + attn_metadata.kv_lens_cuda[:batch_size].copy_(saved_kv_lens_cuda) + attn_metadata.on_update_kv_lens() + + @dataclass class CUDAGraphRunnerConfig: """Configuration for the CUDAGraphRunner, passed from the ModelEngine.""" @@ -407,9 +428,12 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) + attn_metadata = capture_inputs["attn_metadata"] + saved_kv_lens_cuda = _save_spec_decode_capture_state( + attn_metadata, enable_spec_decode) self.graph_metadata[key] = { - "attn_metadata": initial_inputs["attn_metadata"], + "attn_metadata": attn_metadata, "spec_metadata": initial_inputs.get("spec_metadata", None), } @@ -433,12 +457,16 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs) if postprocess_fn is not None: postprocess_fn(capture_inputs) + _restore_spec_decode_capture_state(attn_metadata, + saved_kv_lens_cuda) with torch.cuda.graph(graph, pool=self.memory_pool): output = _setup_spec_decoding_and_forward( key, forward_fn, capture_inputs) if postprocess_fn is not None: postprocess_fn(capture_inputs) + _restore_spec_decode_capture_state(attn_metadata, + saved_kv_lens_cuda) self.graphs[key] = graph self.graph_outputs[key] = make_weak_ref(output) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f633a590f4c8..9c7ef2f16c75 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -109,6 +109,7 @@ l0_b200: - unittest/_torch/compilation - unittest/_torch/debugger - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp + - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d5d316518f22..520943fd97ab 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -171,6 +171,8 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) + - examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8] TIMEOUT (60) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 460fdde0ab82..3fcfc8bafb78 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -84,6 +84,35 @@ def get_num_attention_layers(self) -> int: assert cost.intercept == 0 +def test_mtp_extra_tokens_are_in_context_capacity(): + cache_manager = object.__new__(DeepseekV4CacheManager) + cache_manager.pp_layers = [0] + cache_manager._compress_ratios = [1] + cache_manager._get_attn_bytes_per_block = lambda _attn_type, _layer_idx: 1 + cache_manager._get_window_size = lambda _compress_ratio, _attn_type: 128 + cache_manager.max_batch_size = 1 + cache_manager.max_seq_len = 264 + cache_manager._max_num_tokens = 256 + cache_manager._max_draft_len = 3 + cache_manager.num_extra_kv_tokens = 2 + cache_manager.enable_stats = False + cache_manager.enable_swa_scratch_reuse = False + cache_manager.block_reuse_policy = BlockReusePolicy.ALL_REUSABLE + + config = cache_manager._build_cache_config( + KvCacheConfig(), + tokens_per_block=128, + vocab_size=129280, + cache_tiers=[GpuCacheTierConfig(quota=1024)], + ) + + assert config.typical_step is not None + assert config.typical_step.kv_caches[0].capacity == 258 + assert config.typical_step.kv_caches[0].history_length == 0 + assert config.constraints[1].kv_caches[0].capacity == 258 + assert config.constraints[1].kv_caches[0].history_length == 0 + + def test_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(DeepseekV4CacheManager) manager.pp_layers = [0, 1] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a1583b33b756..51c36071274c 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -11,6 +11,8 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + _restore_spec_decode_capture_state, _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( PyTorchModelEngine, _build_request_multimodal_input) @@ -167,6 +169,25 @@ def test_build_request_multimodal_input_skips_when_cache_disabled( self.assertIsNone( _build_request_multimodal_input(request, cache_enabled=False)) + def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None: + attn_metadata = Mock() + attn_metadata.num_seqs = 1 + attn_metadata.kv_lens_cuda = torch.tensor([4095], dtype=torch.int32) + + saved_kv_lens_cuda = _save_spec_decode_capture_state( + attn_metadata, enable_spec_decode=True) + + # CUDA graph capture performs two eager warmup forwards. A speculative + # draft loop may advance the static attention metadata during each + # forward, but the next warmup must start from the original input. + for _ in range(2): + attn_metadata.kv_lens_cuda.add_(1) + _restore_spec_decode_capture_state(attn_metadata, + saved_kv_lens_cuda) + self.assertEqual(attn_metadata.kv_lens_cuda.tolist(), [4095]) + + self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2) + def test_pad_generation_requests(self) -> None: model_engine, kv_cache_manager = create_model_engine_and_kvcache() resource_manager = ResourceManager( diff --git a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py index 0fbe520ee067..0486f387a749 100644 --- a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py +++ b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py @@ -53,7 +53,6 @@ MAX_BATCH_SIZE = 16 VOCAB_SIZE = 129280 NUM_KV_HEADS = 1 -INDEXER_QUANT_BLOCK_SIZE = 128 # DeepSeek-V4 specific ratios (mirrors module constants) @@ -407,32 +406,31 @@ def _expected_valid_blocks( def _split_blockwise_buffer( buffer: torch.Tensor, - index_head_dim: int = INDEX_HEAD_DIM, - quant_block_size: int = INDEXER_QUANT_BLOCK_SIZE, + data_size: int, + scale_size: int, + scale_dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Split a blockwise FP8 quantized buffer into value and scale buffers. + """Split a blockwise quantized buffer into value and scale buffers. Args: buffer: shape [num_blocks, tokens_per_block, bytes_per_token] Returns: - (values_buffer, scales_buffer) where values are uint8 and scales are float32 + (values_buffer, scales_buffer), preserving the manager's scale dtype. """ num_blocks, tokens_per_block, bytes_per_token = buffer.shape bytes_per_block = bytes_per_token * tokens_per_block # Value buffer - value_shape = (num_blocks, tokens_per_block, index_head_dim) - value_stride = (bytes_per_block, index_head_dim, 1) + value_shape = (num_blocks, tokens_per_block, data_size) + value_stride = (bytes_per_block, data_size, 1) value_buffer = buffer.as_strided(value_shape, value_stride, 0).view(torch.uint8) # Scale buffer - scale_dim = index_head_dim // quant_block_size - scale_bytes = scale_dim * 4 # float32 = 4 bytes - scale_shape = (num_blocks, tokens_per_block, scale_bytes) - scale_stride = (bytes_per_block, scale_bytes, 1) - scale_offset = index_head_dim * tokens_per_block - scale_buffer = buffer.as_strided(scale_shape, scale_stride, scale_offset).view(torch.float32) + scale_shape = (num_blocks, tokens_per_block, scale_size) + scale_stride = (bytes_per_block, scale_size, 1) + scale_offset = data_size * tokens_per_block + scale_buffer = buffer.as_strided(scale_shape, scale_stride, scale_offset).view(scale_dtype) return value_buffer, scale_buffer @@ -458,7 +456,13 @@ def _read_cache_data( return torch.tensor([]), None if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: - values_buf, scales_buf = _split_blockwise_buffer(buffer) + scale_dtype = torch.float32 if mgr._indexer_k_dtype == "fp8" else torch.uint8 + values_buf, scales_buf = _split_blockwise_buffer( + buffer, + data_size=mgr._indexer_data_size, + scale_size=mgr._indexer_scale_size, + scale_dtype=scale_dtype, + ) return values_buf[indices], scales_buf[indices] return buffer[indices], None