Skip to content

[None][perf] Close Mamba hybrid warmup gap in autotuner warmup #16177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenfeiz0326 wants to merge 1 commit into
base: main
Choose a base branch
from
+93 −29
Open

[None][perf] Close Mamba hybrid warmup gap in autotuner warmup #16177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.

import math
import os
from typing import Tuple

import torch
Expand Down Expand Up @@ -414,6 +415,14 @@ def prepare(self, attn_metadata: AttentionMetadata):
device='cpu')

self.has_initial_states_cpu[:num_contexts].copy_(initial_states_cpu)
# Warmup-only override: force HAS_INITSTATES=True path so the
# HAS_INITSTATES=True variants of _state_passing_fwd_kernel,
# _chunk_scan_fwd_kernel, and _chunk_state_varlen_kernel compile
# during warmup instead of the first real-request iter that hits
# chunked prefill with cached tokens.
if os.environ.get(
"TLLM_MAMBA_WARMUP_FORCE_INITSTATES") == "1":
self.has_initial_states_cpu[:num_contexts].fill_(True)
Comment on lines +423 to +425

@coderabbitai coderabbitai Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Keep the init-state override scoped to warmup.

Mamba2Metadata.prepare() now honors TLLM_MAMBA_WARMUP_FORCE_INITSTATES on every call. If that variable is set at process scope, real prefill requests with zero cached tokens are treated as having initial states. Prefer passing an explicit warmup-only flag through metadata instead of reading global env here.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py` around lines 423 - 425,
`Mamba2Metadata.prepare()` is reading `TLLM_MAMBA_WARMUP_FORCE_INITSTATES`
globally on every call, which makes the init-state override affect non-warmup
requests too. Update `prepare()` to take and use an explicit warmup-only flag
from metadata, and apply the `has_initial_states_cpu` override only when that
flag is set; keep the change scoped around the `Mamba2Metadata.prepare` logic
that sets `has_initial_states_cpu[:num_contexts]`.

# Mirror CPU staging flags to the CUDA-side buffer asynchronously.
self.has_initial_states[:num_contexts].copy_(
self.has_initial_states_cpu[:num_contexts], non_blocking=True)
Expand Down
113 changes: 84 additions & 29 deletions tensorrt_llm/_torch/pyexecutor/model_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1238,7 +1238,22 @@ def trtllm_gen_fmha_jit_warmup():
torch.cuda.synchronize()

def _run_autotuner_warmup(self, resource_manager: ResourceManager):
"""Runs a forward pass to populate the autotuner cache."""
"""Runs forward passes over representative shapes to populate the autotuner cache.

Sweeps pure-ctx-max plus a small set of mixed ctx+gen shapes so that a
serve workload's early iterations (context ingestion overlapping with
already-running decode) don't hit an autotune cache miss on their
(M,N,K) tuples and stall the stream doing in-place tuning under load.
Set env TLLM_AUTOTUNE_SWEEP=0 to disable and fall back to the
single-shape (curr_max_num_tokens, 0) sweep.

Mamba hybrid models (e.g. Nemotron 3 Super 120B) also add two extra
passes: a multi-sequence prefill (least_requests=False) and a
multi-sequence prefill with HAS_INITSTATES=True forced via env var,
so the Mamba SSD Triton kernels' multi-seq + chunked-prefill
constexpr variants compile during warmup instead of at iter 3.
Set env TLLM_MAMBA_MULTISEQ_WARMUP=0 to disable.
"""
if not self.llm_args.enable_autotuner:
return
AutoTuner.get().setup_distributed_state(self.mapping, self.dist)
Expand All @@ -1251,39 +1266,79 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
token_num_upper_bound=token_num_upper_bound,
max_num_draft_tokens=self.original_max_draft_len)

# Each shape: (num_tokens, num_gen_requests, least_requests, force_mamba_initstates)
autotune_shapes = [(curr_max_num_tokens, 0, True, False)]
if os.environ.get("TLLM_AUTOTUNE_SWEEP", "1") == "1":
if curr_max_num_tokens >= 4:
autotune_shapes.append((curr_max_num_tokens, 2, True, False))
if curr_max_num_tokens >= 16:
autotune_shapes.append((curr_max_num_tokens, 8, True, False))
Comment on lines +1272 to +1275

@coderabbitai coderabbitai Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Account for speculative tokens before adding mixed warmup shapes.

These guards only check curr_max_num_tokens >= 4/16, but _create_warmup_request() charges each generation request as 1 + self.max_total_draft_tokens. In speculative modes this can make num_gen_tokens > num_tokens_i, so the “mixed ctx+gen” shape has a negative context-token budget and can warm the wrong shape or over-admit generation work.

Suggested fix 
+            gen_tokens_per_request = 1 + self.max_total_draft_tokens
             if curr_max_num_tokens >= 4:
-                autotune_shapes.append((curr_max_num_tokens, 2, True, False))
+                if curr_max_num_tokens > 2 * gen_tokens_per_request:
+                    autotune_shapes.append((curr_max_num_tokens, 2, True, False))
             if curr_max_num_tokens >= 16:
-                autotune_shapes.append((curr_max_num_tokens, 8, True, False))
+                if curr_max_num_tokens > 8 * gen_tokens_per_request:
+                    autotune_shapes.append((curr_max_num_tokens, 8, True, False))
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if curr_max_num_tokens >= 4:
autotune_shapes.append((curr_max_num_tokens, 2, True, False))
if curr_max_num_tokens >= 16:
autotune_shapes.append((curr_max_num_tokens, 8, True, False))
gen_tokens_per_request = 1 + self.max_total_draft_tokens
if curr_max_num_tokens >= 4:
if curr_max_num_tokens > 2 * gen_tokens_per_request:
autotune_shapes.append((curr_max_num_tokens, 2, True, False))
if curr_max_num_tokens >= 16:
if curr_max_num_tokens > 8 * gen_tokens_per_request:
autotune_shapes.append((curr_max_num_tokens, 8, True, False))
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 1272 - 1275, The
mixed warmup shapes added in model_engine.py need to account for speculative
draft tokens before checking the 4/16 thresholds. In the warmup shape logic that
appends the mixed ctx+gen tuples, use the same token accounting as
_create_warmup_request() by subtracting self.max_total_draft_tokens from the
available generation budget or otherwise ensuring the context budget stays
non-negative. Update the guards around the autotune_shapes.append calls so they
only add shapes when the resulting context/gen split is valid in speculative
modes.


is_mamba_hybrid = isinstance(kv_cache_manager, MambaHybridCacheManager)
if is_mamba_hybrid and os.environ.get(
"TLLM_MAMBA_MULTISEQ_WARMUP", "1") == "1":
if curr_max_num_tokens >= 4:
autotune_shapes.append(
(curr_max_num_tokens, 0, False, False))
autotune_shapes.append(
(curr_max_num_tokens, 0, False, True))

logger.info(
f"[Autotuner] Warmup sweep shapes (num_tokens, num_gen_requests, least_requests, force_mamba_initstates): {autotune_shapes}"
)

@contextlib.contextmanager
def _force_mamba_initstates_env():
prev = os.environ.get("TLLM_MAMBA_WARMUP_FORCE_INITSTATES")
os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = "1"
try:
yield
finally:
if prev is None:
os.environ.pop(
"TLLM_MAMBA_WARMUP_FORCE_INITSTATES", None)
else:
os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = prev

cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None)
with self.no_cuda_graph(), autotune(cache_path=cache_path):
warmup_request = self._create_warmup_request(
resource_manager, curr_max_num_tokens, 0)
with self._release_batch_context(warmup_request,
resource_manager) as batch:
if batch is None and self.mapping.tp_size <= 1:
pass # Single rank, safe to skip
else:
self._assert_all_tp_ranks_have_warmup_batch(
batch, curr_max_num_tokens)
if batch is not None:
# Reset the flag is_first_draft for the draft model.
# This is necessary for overlap scheduler.
spec_resource_manager = resource_manager.get_resource_manager(
ResourceManagerType.SPEC_RESOURCE_MANAGER)
if self.is_draft_model and isinstance(
spec_resource_manager, Eagle3ResourceManager):
spec_resource_manager.is_first_draft = True
for (num_tokens_i, num_gen_requests_i, least_req_i,
force_init_i) in autotune_shapes:
init_ctx = (_force_mamba_initstates_env()
if force_init_i else contextlib.nullcontext())
with init_ctx:
warmup_request = self._create_warmup_request(
resource_manager, num_tokens_i, num_gen_requests_i,
least_requests=least_req_i)
with self._release_batch_context(
warmup_request, resource_manager) as batch:
if batch is None and self.mapping.tp_size <= 1:
continue # Single rank, safe to skip this shape
self._assert_all_tp_ranks_have_warmup_batch(
batch, num_tokens_i)
if batch is None:
continue
# Reset the flag is_first_draft for the draft model.
# This is necessary for overlap scheduler.
spec_resource_manager = resource_manager.get_resource_manager(
ResourceManagerType.SPEC_RESOURCE_MANAGER)
if self.is_draft_model and isinstance(
spec_resource_manager, Eagle3ResourceManager):
spec_resource_manager.is_first_draft = True

self.forward(batch,
new_tensors_device=None,
resource_manager=resource_manager)
self.forward(batch,
new_tensors_device=None,
resource_manager=resource_manager)

# pp_recv in AutoTuner choose_one will never be called if there is no tuning op during the forward pass.
# So we need to make an extra call to consume the previous rank's pp_send to guarantee that the previous rank's pp_send is released.
AutoTuner.get().cache_pp_recv()
# Send the cache after the tuning process to the next PP rank
AutoTuner.get().cache_pp_send()
# Clean the pp flag to avoid deadlock with synchronous send/recv
AutoTuner.get().clean_pp_flag()
# pp_recv in AutoTuner choose_one will never be called if there is no tuning op during the forward pass.
# So we need to make an extra call to consume the previous rank's pp_send to guarantee that the previous rank's pp_send is released.
AutoTuner.get().cache_pp_recv()
# Send the cache after the tuning process to the next PP rank
AutoTuner.get().cache_pp_send()
# Clean the pp flag to avoid deadlock with synchronous send/recv
AutoTuner.get().clean_pp_flag()

torch.cuda.synchronize()
torch.cuda.synchronize()

logger.info(
f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}"
Expand Down
Loading