From 5e949aa39ce41538f70ca3511fdd3deef4963d89 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:26:30 -0700 Subject: [PATCH 1/2] [None][fix] Warmup: no generation shapes on a context-only KV cache A disaggregated context server with a context-only cache (KDA context-only state, context-only FP4 MLA cache) cannot build generation warmup requests, so create_warmup_request returns None for them. Under DWDP (Kimi K3 CTX) the warmup agreement cannot confirm a unanimous skip and turns that None into a fatal error. The attention-JIT and autotuner warmups now leave out the generation shapes up front when the cache is context-only; that property is the same on every rank of the role. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../engine/runners/decoder/warmup.py | 23 +++- .../engine/runners/decoder/test_warmup.py | 104 ++++++++++++++++++ 2 files changed, 122 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 0fc2fc5f62f2..4b4fe65c3411 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/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, @@ -881,8 +885,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 +907,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 +1120,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)) diff --git a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py index 7f3a33791a88..0f27a867e671 100644 --- a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py +++ b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py @@ -556,6 +556,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") From f2fc07ddf0a75ac3da33e7e1b26d6aa18e5b5caf Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:50:55 -0800 Subject: [PATCH 2/2] [None][fix] Warmup: no generation shapes in the general and memory-pool phases on a context-only KV cache A context-only KV cache that is not a Mamba hybrid manager (e.g. a pure-MLA model on Fp4MlaKVCacheManagerV2 with the context-only FP4 MLA cache) passes can_run_general_warmup, so the general warmup and the memory-pool pre-population still proposed generation shapes: (1, 1) and (max_batch_size, max_batch_size). create_warmup_request returns None for them, and under DWDP should_run_warmup_batch turns that into a fatal error. _get_max_shape_warmup_requests now returns the context shape only on a context-only cache, which covers both phases, and the general plan drops (1, 1). Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../engine/runners/decoder/warmup.py | 18 ++++++++++----- .../engine/runners/decoder/test_warmup.py | 22 +++++++++++++++++++ 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 4b4fe65c3411..c23a29fef747 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -300,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( @@ -320,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 @@ -337,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)] diff --git a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py index 0f27a867e671..ebd68dd12413 100644 --- a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py +++ b/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: