[None][fix] Fix DeepSeek V4 KV cache warmup handling and serveral other issues

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=[

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)

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)

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)

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]

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(

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