[None][fix] Warmup: no generation shapes on a context-only KV cache

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py

@@ -22,7 +22,11 @@
 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata
 from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs
 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager
-from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
+from tensorrt_llm._torch.pyexecutor.resource_manager import (
+    ResourceManager,
+    ResourceManagerType,
+    is_context_only_kv_cache,
+)
 from tensorrt_llm._torch.pyexecutor.runtime_hooks import runtime_hooks_of
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
     warmup_sampling_module,
@@ -296,6 +300,8 @@ def _get_max_shape_warmup_requests(
     ) -> list[tuple[int, int]]:
         """
         Returns warmup configs covering the maximum context and generation shapes.
+
+        A context-only KV cache cannot hold generation rows, so it gets the context shape only.
         """
 
         kv_cache_manager = resource_manager.get_resource_manager(
@@ -316,10 +322,10 @@ def _get_max_shape_warmup_requests(
             // self._config.max_beam_width,
         )
 
-        warmup_requests_configs = [
-            (curr_max_num_tokens, 0),  # max_num_tokens, pure context
-            (max_batch_size, max_batch_size),  # max_batch_size, pure generation
-        ]
+        warmup_requests_configs = [(curr_max_num_tokens, 0)]  # max_num_tokens, pure context
+        if not is_context_only_kv_cache(kv_cache_manager):
+            # max_batch_size, pure generation
+            warmup_requests_configs.append((max_batch_size, max_batch_size))
 
         return warmup_requests_configs
 
@@ -333,8 +339,12 @@ def _get_full_general_warmup_requests(
         and small-context (2-token path) cases.
         """
         max_configs = self._get_max_shape_warmup_requests(resource_manager)
-        # Specialize for 1 token pure ctx and pure gen
-        one_token_configs = [(1, 0), (1, 1)]
+        # Specialize for 1 token pure ctx and pure gen (no gen on a context-only cache)
+        one_token_configs = [(1, 0)]
+        if not is_context_only_kv_cache(
+            resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER)
+        ):
+            one_token_configs.append((1, 1))
         # Small ctx specialization
         small_ctx_configs = [(2, 0)]
 
@@ -881,8 +891,13 @@ def _run_attention_warmup(
         logger.info("Running TRTLLM-Gen FMHA JIT warmup")
 
         warmup_requests_configs = []
-        if self.guided_decoder is None:
-            # doesn't support guided decoding
+        # A context-only cache cannot hold generation rows. The property is the same on every
+        # rank of a role, so its shapes are left out here rather than skipped by agreement.
+        supports_generation = not is_context_only_kv_cache(
+            resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER)
+        )
+        if self.guided_decoder is None and supports_generation:
+            # doesn't support guided decoding or a context-only cache
             warmup_requests_configs.append(
                 (1 + self._config.max_total_draft_tokens, 1)
             )  # one generation request
@@ -898,7 +913,7 @@ def _run_attention_warmup(
         # here: they always run on a ``MambaHybridCacheManager``, so
         # ``can_run_general_warmup`` is False and the configs above never run.
 
-        if self.guided_decoder is None and can_run_general_warmup:
+        if self.guided_decoder is None and supports_generation and can_run_general_warmup:
             # The cute_dsl_mla FMHA lib now only support the generation-only batch, we need to
             # warmup the TRTLLM-Gen FMHA lib for the mixed context+generation batch.
             # One MIXED context+generation batch (1 ctx token + 1 gen request).
@@ -1111,7 +1126,11 @@ def _run_autotuner_warmup(
             warmup_configs = [(curr_max_num_tokens, num_gen_requests)]
         else:
             warmup_configs = [(curr_max_num_tokens, 0)]
-            if self.guided_decoder is None and not self.mapping.has_pp():
+            if (
+                self.guided_decoder is None
+                and not self.mapping.has_pp()
+                and not is_context_only_kv_cache(kv_cache_manager)
+            ):
                 # Add generation request to warmup the autotuner cache.
                 warmup_configs.append((1 + self._config.max_total_draft_tokens, 1))
 

tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py

@@ -191,6 +191,28 @@ def get_num_available_tokens(self, **_kwargs: object) -> int:
         return self._available_tokens
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("cache_flag", ["kda_context_only", "fp4_mla_context_only_cache"])
+@pytest.mark.parametrize("context_only", [False, True])
+def test_general_and_memory_pool_warmups_leave_out_generation_shapes_on_context_only_cache(
+    cache_flag: str, context_only: bool
+) -> None:
+    """A context-only cache that is not a Mamba hybrid one runs the general and memory-pool phases.
+
+    Both take their shapes from these lists, so the lists must hold context shapes only.
+    """
+    warmup = decoder_warmup(decoder_config(), agreement=_agreement(world_size=4, dwdp_size=4))
+    resources = _context_only_cache_resources(cache_flag, context_only)
+
+    max_shapes = warmup._get_max_shape_warmup_requests(resources)
+    general_shapes = warmup._get_full_general_warmup_requests(resources)
+
+    for shapes in (max_shapes, general_shapes):
+        has_generation = any(num_gen > 0 for _, num_gen in shapes)
+        assert has_generation is not context_only, shapes
+    assert ((1, 1) in general_shapes) is not context_only
+
+
 def _engine(
     *, world_size: int = 1, dwdp_size: int = 0, config: DecoderRunnerConfig | None = None
 ) -> DecoderWarmup:
@@ -556,6 +578,110 @@ def test_autotuner_warmup_skips_only_completed_token_shape(
     assert engine._completed_autotuner_warmup_num_tokens == 16
 
 
+def _context_only_warmup_requests(context_only: bool) -> Mock:
+    """Dummy requests that, like the real factory, build no generation rows on a context-only cache."""
+    requests = Mock()
+
+    def create_warmup_request(_resources, num_tokens, num_gen_requests):
+        return None if context_only and num_gen_requests else ScheduledRequests()
+
+    requests.create_warmup_request.side_effect = create_warmup_request
+    requests.release_batch_context.side_effect = lambda request, _resources: (
+        contextlib.nullcontext(request)
+    )
+    return requests
+
+
+def _context_only_cache_resources(cache_flag: str, context_only: bool) -> SimpleNamespace:
+    cache = SimpleNamespace(**{cache_flag: context_only}, get_num_available_tokens=lambda **_: 16)
+    return SimpleNamespace(
+        get_resource_manager=lambda key: (
+            cache if key == ResourceManagerType.KV_CACHE_MANAGER else None
+        )
+    )
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("cache_flag", ["kda_context_only", "fp4_mla_context_only_cache"])
+@pytest.mark.parametrize("context_only", [False, True])
+def test_attention_warmup_leaves_out_generation_shapes_on_context_only_cache(
+    cache_flag: str, context_only: bool
+) -> None:
+    """Under DWDP a missing warmup batch is fatal, so a context-only cache must not propose one."""
+    agreement = _agreement(world_size=4, dwdp_size=4)
+    requests = _context_only_warmup_requests(context_only)
+    forward_warmup = Mock()
+    warmup = decoder_warmup(
+        decoder_config(),
+        agreement=agreement,
+        dummy_requests=requests,
+        forward_warmup=forward_warmup,
+    )
+    resources = _context_only_cache_resources(cache_flag, context_only)
+
+    with patch("torch.cuda.synchronize"):
+        warmup._run_attention_warmup(resources, can_run_general_warmup=True)
+
+    num_gen_requests = [
+        kwargs["num_gen_requests"] for _, kwargs in requests.create_warmup_request.call_args_list
+    ]
+    assert num_gen_requests == ([0] if context_only else [1, 0, 1])
+    assert forward_warmup.call_count == len(num_gen_requests)
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("cache_flag", ["kda_context_only", "fp4_mla_context_only_cache"])
+@pytest.mark.parametrize("context_only", [False, True])
+def test_autotuner_warmup_leaves_out_generation_shape_on_context_only_cache(
+    monkeypatch: pytest.MonkeyPatch, cache_flag: str, context_only: bool
+) -> None:
+    """The autotuner tunes only the context batch on a context-only cache, also under DWDP."""
+    requests = _context_only_warmup_requests(context_only)
+    engine = SimpleNamespace(
+        _config=SimpleNamespace(
+            enable_autotuner=True,
+            torch_compile_enabled=False,
+            torch_compile_prefill_only=False,
+            max_num_tokens=16,
+            max_batch_size=16,
+            max_seq_len=2,
+            original_max_draft_len=0,
+            max_total_draft_tokens=0,
+            is_spec_decode=False,
+            max_draft_len=0,
+            spec_config=None,
+        ),
+        _completed_autotuner_warmup_num_tokens=None,
+        _warmup_timer=_WarmupTimer(rank=0),
+        model=SimpleNamespace(modules=lambda: []),
+        mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False),
+        dist=object(),
+        guided_decoder=None,
+        _graphs=SimpleNamespace(no_cuda_graph=lambda: contextlib.nullcontext(), enabled=False),
+        _dummy_requests=requests,
+        _agreement=_agreement(world_size=4, dwdp_size=4),
+        _release_megamoe_profiling_scratch=Mock(),
+        _forward_warmup=Mock(),
+    )
+    resources = _context_only_cache_resources(cache_flag, context_only)
+    monkeypatch.delenv("TLLM_AUTOTUNER_CACHE_PATH", raising=False)
+    monkeypatch.setattr(
+        decoder_warmup_module.AutoTuner, "get", Mock(return_value=Mock(profiling_cache={}))
+    )
+    monkeypatch.setattr(
+        decoder_warmup_module, "autotune", lambda **kwargs: contextlib.nullcontext()
+    )
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda: None)
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
+    monkeypatch.setattr(decoder_warmup_module, "clear_memory_buffers", Mock())
+
+    DecoderWarmup._run_autotuner_warmup(engine, resources)
+
+    shapes = [args[1:] for args, _ in requests.create_warmup_request.call_args_list]
+    assert shapes == ([(16, 0)] if context_only else [(16, 0), (1, 1)])
+    assert engine._forward_warmup.call_count == len(shapes)
+
+
 @pytest.mark.parametrize("phase", ["general", "attention", "mamba", "prefill"])
 def test_warmup_phases_pass_configured_speculation_state(phase: str) -> None:
     spec = DraftTargetDecodingConfig(max_draft_len=3, speculative_model="dummy")