diff --git a/cpp/tensorrt_llm/thop/moeUtilOp.cpp b/cpp/tensorrt_llm/thop/moeUtilOp.cpp index e5496a89cdb0..885eef9a24a9 100644 --- a/cpp/tensorrt_llm/thop/moeUtilOp.cpp +++ b/cpp/tensorrt_llm/thop/moeUtilOp.cpp @@ -138,9 +138,13 @@ std::tuple(experts_per_token * num_rows)}, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 57ef2259ebef..a3db4bd145a6 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -234,6 +234,11 @@ class ModelConfig(Generic[TConfig]): max_seq_len: Optional[int] = None moe_max_num_tokens: Optional[int] = None + # Set in __post_init__; a normal field, not init=False, so that + # dataclasses.replace() and copy.copy() carry it. + _moe_max_num_tokens_is_default: Optional[bool] = field(default=None, + repr=False, + compare=False) moe_load_balancer: Optional[MoeLoadBalancerConfig] = None attn_backend: str = 'TRTLLM' @@ -336,9 +341,24 @@ def get_all_reduce_strategy(strategy: str = "AUTO"): # Set default moe_max_num_tokens if not specified # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled + # Record the provenance first: once filled in, a derived size is + # indistinguishable from one a deployment configured to the same number. + if self._moe_max_num_tokens_is_default is None: + self._moe_max_num_tokens_is_default = self.moe_max_num_tokens is None if self.moe_max_num_tokens is None: self.moe_max_num_tokens = self.max_num_tokens * self.mapping.dp_size + def is_moe_max_num_tokens_default(self) -> bool: + """Whether ``moe_max_num_tokens`` was derived rather than configured. + + A MoE backend with a conservative workspace cap uses this to clamp only + the derived size. A config rebuilt from another one's + ``moe_max_num_tokens`` -- the draft configs in + ``modeling_speculative.py`` -- reads as configured, which is safe: the + target's first MoE layer has already capped the forwarded size. + """ + return bool(self._moe_max_num_tokens_is_default) + @property def torch_dtype(self) -> torch.dtype: """Get the torch dtype of the model.""" diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py index 9ee2a8a55e17..4895152fa345 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py @@ -3,6 +3,7 @@ import threading from abc import ABC, abstractmethod +from bisect import bisect_left from typing import Any, Dict, Iterator, Tuple, Union from tensorrt_llm.mapping import Mapping @@ -25,12 +26,15 @@ class ConsumableWeightsDict: def __init__(self, weights: Dict[str, Any]): self._weights = weights self._lock = threading.Lock() + self._key_index: list[str] | None = None def __getitem__(self, key: str) -> Any: return self._weights[key] def __setitem__(self, key: str, value: Any) -> None: with self._lock: + if key not in self._weights: + self._key_index = None self._weights[key] = value def __delitem__(self, key: str) -> None: @@ -68,6 +72,8 @@ def get(self, key: str, default: Any = None) -> Any: def update(self, other: Dict[str, Any]) -> None: with self._lock: + if any(key not in self._weights for key in other): + self._key_index = None self._weights.update(other) def clear(self) -> None: @@ -79,6 +85,60 @@ def clear(self) -> None: """ with self._lock: self._weights.clear() + self._key_index = [] + + @classmethod + def take_ownership(cls, source: Union[Dict[str, Any], + "ConsumableWeightsDict"], + derived: Dict[str, Any]) -> Dict[str, Any]: + """Hand ``derived`` the tensors ``source`` was holding. + + A renamed or filtered mapping aliases the tensors it was built from, so + while the source is alive it holds a second reference to each one and + consuming the alias frees nothing. Emptying the source makes the alias + the last reference, which is what lets the loader release weights + module by module instead of pinning the whole checkpoint. + + A plain dict source is returned unchanged -- it was never doing + incremental release. **The caller must not use ``source`` afterwards.** + """ + if not isinstance(source, cls): + return derived + source.clear() + return cls(derived) + + def filter_prefix(self, prefix: str) -> Dict[str, Any]: + """Same result as a ``startswith(prefix)`` scan, without the scan. + + ``prefix`` must be non-empty. Callers that may pass an empty prefix + keep their own scan; only the loading loop, which always names a + module, comes through here. + """ + with self._lock: + start = len(prefix) + 1 + return { + key[start:]: self._weights[key] + for key in self._keys_with_prefix_locked(prefix) + } + + def _keys_with_prefix_locked(self, prefix: str) -> list[str]: + """Return the live keys starting with ``prefix``, in sorted order. + + Deletions deliberately do not invalidate the index: a stale index is + always a superset of the live keys, so filtering on membership keeps + every reader correct while each lookup stays proportional to the keys + it matched rather than to the size of the checkpoint. + """ + if self._key_index is None: + self._key_index = sorted(self._weights) + begin = bisect_left(self._key_index, prefix) + # Exclusive upper bound: every key starting with the prefix sorts + # before the prefix with its last character incremented. + upper_bound = prefix[:-1] + chr(ord(prefix[-1]) + 1) + end = bisect_left(self._key_index, upper_bound, begin) + return [ + key for key in self._key_index[begin:end] if key in self._weights + ] def mark_consumed_keys(self, keys) -> int: """Delete an exact set of keys to free memory. @@ -107,9 +167,7 @@ def mark_consumed(self, prefix: str) -> int: Thread-safe: uses a lock to prevent concurrent modification issues. """ with self._lock: - keys_to_delete = [ - k for k in self._weights.keys() if k.startswith(prefix + ".") - ] + keys_to_delete = self._keys_with_prefix_locked(prefix + ".") for key in keys_to_delete: del self._weights[key] return len(keys_to_delete) diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index f604923c8799..f1109ec36e4d 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -148,16 +148,17 @@ def rename_by_params_map( Returns - renamed_weights: Mapping[str, torch.Tensor], weight dict with renamed keys and unchanged tensor values. If the input `weights` is a - `ConsumableWeightsDict`, the returned object preserved that type. + `ConsumableWeightsDict`, the returned object preserves that type and + takes the tensors over from it -- the input is emptied, so **the + caller must not use `weights` afterwards**. See + `ConsumableWeightsDict.take_ownership` for why the transfer is what + lets the loader release weights module by module. """ import re from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ ConsumableWeightsDict - # Check if input is a ConsumableWeightsDict to preserve the type - is_consumable = isinstance(weights, ConsumableWeightsDict) - # Create a new dictionary to store the renamed weights renamed_weights = {} @@ -180,10 +181,7 @@ def rename_by_params_map( if key not in matched_keys: renamed_weights[key] = weights[key] - # Preserve ConsumableWeightsDict type if that's what was passed in - if is_consumable: - return ConsumableWeightsDict(renamed_weights) - return renamed_weights + return ConsumableWeightsDict.take_ownership(weights, renamed_weights) def preprocess_weights( self, weights: Mapping[str, @@ -281,6 +279,15 @@ def filter_weights( """ Return only weights that start with the prefix (and with the prefix removed) """ + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + ConsumableWeightsDict + + # The loading loop calls this once per module, so on a large + # checkpoint the scan below is quadratic; a ConsumableWeightsDict can + # answer the same query from its key index instead. + if prefix and isinstance(weights, ConsumableWeightsDict): + return weights.filter_prefix(prefix) + result = {} for k, v in weights.items(): if k.startswith(prefix): diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index b71fe8495928..10499969fac4 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -36,6 +36,7 @@ ) from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types from ..utils import is_nvfp4_marlin_supported_sm +from .checkpoints.base_weight_loader import ConsumableWeightsDict from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from .modeling_qwen3_next import Qwen3NextForCausalLM @@ -80,6 +81,18 @@ def _get_qwen35_moe_model_defaults(llm_args: "TorchLlmArgs") -> dict: return defaults +def _filter_language_model_weights(weights: Dict[str, torch.Tensor]): + """Drop vision weights without disabling incremental weight consumption. + + Ownership: a ConsumableWeightsDict input is emptied, since the returned + mapping aliases its tensors. The caller must use only the return value. + """ + filtered_weights = { + key: value for key, value in weights.items() if not key.startswith("model.visual.") + } + return ConsumableWeightsDict.take_ownership(weights, filtered_weights) + + def _translate_mtp_pattern(name, n_hidden_layers): """Translate an HF ``mtp.*`` exclude pattern to a TRT-LLM module path. @@ -758,7 +771,7 @@ def load_weights( ) if weight_mapper.model is not self.llm: weight_mapper.init_model_and_config(self.llm, self.llm.model_config) - filtered_weights = {k: v for k, v in weights.items() if not k.startswith("model.visual.")} + filtered_weights = _filter_language_model_weights(weights) params_map = { r"^model\.language_model\.(.*)$": r"model.\1", } diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 87808df7ac48..1a140ea15e03 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -1206,6 +1206,48 @@ def filter_weights(prefix, weights: Dict): return result +def _get_load_weights_num_workers() -> Optional[int]: + """Return the per-rank module-loading worker limit, or None for the default. + + Weight loading runs one ThreadPoolExecutor per rank, which without an + explicit limit defaults to as many as 32 workers (CPython's + ``min(32, cpu_count + 4)``; the count is the machine's, not the rank's + share of it). The limit is per rank, so four ranks on a node can have four + times that many module loads in flight. Each one holds its own host-side + working set while it stages and transforms a module's weights, and every + rank's is charged to the same host-memory cgroup -- which is how a large + checkpoint exhausts host memory while the GPUs are nowhere near full. + + ``TLLM_LOAD_WEIGHTS_NUM_WORKERS`` bounds that overlap. Unset or blank keeps + the executor default; a positive integer trades loading parallelism for + host-memory headroom; anything else raises, so a typo cannot look like it + took effect. ``TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL`` takes precedence + and skips the pool entirely -- this variable is the range in between. + + Set it when ranks share a constrained cgroup. Tune it against node + ``memory.peak`` and the slowest rank's init time, not per-process RSS, + which does not see shared page cache. ``4`` measured well on a four-rank + node but is a starting point, not a default; retune per checkpoint and + topology. + """ + env_name = "TLLM_LOAD_WEIGHTS_NUM_WORKERS" + value = os.environ.get(env_name) + if value is None or not value.strip(): + return None + + try: + num_workers = int(value) + except ValueError as error: + raise ValueError( + f"{env_name} must be a positive integer, got {value!r}") from error + if num_workers <= 0: + raise ValueError( + f"{env_name} must be a positive integer, got {value!r}") + logger.info( + f"Limiting concurrent module weight loading to {num_workers} workers") + return num_workers + + def run_concurrently(func, args_list, reduce_func=None, @@ -1404,7 +1446,10 @@ def load_single_module(name, module): for name, module in model.named_modules(remove_duplicate=False) if name not in serial_load_modules ] - run_concurrently(load_single_module, args_list, pbar=pbar) + run_concurrently(load_single_module, + args_list, + pbar=pbar, + num_workers=_get_load_weights_num_workers()) def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM], @@ -1537,4 +1582,7 @@ def load_single_module(name, module): for name, module in model.named_modules(remove_duplicate=False) if name not in serial_load_modules ] - run_concurrently(load_single_module, args_list, pbar=pbar) + run_concurrently(load_single_module, + args_list, + pbar=pbar, + num_workers=_get_load_weights_num_workers()) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index d5861f0798e7..39b713613fce 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -22,6 +22,7 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm import deep_gemm from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo from ...memory_buffer_utils import get_memory_buffers @@ -36,6 +37,48 @@ MoEWeightLoadingMode, UnquantizedFusedMoEMethod) from .routing import BaseMoeRoutingMethod +_DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS = 18688 + + +def _configure_deepgemm_moe_max_num_tokens(model_config: ModelConfig) -> None: + """Cap only the workspace size ModelConfig derived for itself. + + A size the deployment configured is honored as-is; the derived + ``max_num_tokens * dp_size`` is capped to keep the 8k/1k case OOM-safe. + Each outcome is logged once, since ConfigurableMoE reads the size only + after the backend is built and an unexpected one is otherwise invisible. + """ + default = _DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS + moe_max_num_tokens = model_config.moe_max_num_tokens + key = "deepgemm_moe_max_num_tokens" + + if not model_config.is_moe_max_num_tokens_default(): + configured = f"Using the configured moe_max_num_tokens {moe_max_num_tokens}" + if moe_max_num_tokens > default: + logger.warning_once( + f"{configured}, above the DeepGEMM default {default}; this may " + "increase GPU memory usage.", + key=key) + else: + logger.info_once(f"{configured}.", key=key) + return + + derived = (f"derived moe_max_num_tokens (max_num_tokens " + f"{model_config.max_num_tokens} * dp_size " + f"{model_config.mapping.dp_size})") + if moe_max_num_tokens <= default: + logger.info_once(f"Using the {derived}.", key=key) + return + + logger.info_once( + f"Clamping the {derived} to the DeepGEMM default {default}; set " + "moe_config.max_num_tokens to size the workspace explicitly.", + key=key) + was_frozen = model_config._frozen + model_config._frozen = False + model_config.moe_max_num_tokens = default + model_config._frozen = was_frozen + @triton.jit def _masked_index_copy_group_quant_fp8( @@ -619,7 +662,7 @@ def preprocess_after_permute(expert_first_token_offset_tensor, Only the number of permuted (expanded) tokens is needed here, not the permuted activations themselves. Callers that run moe_permute_op with - skip_data_expand=True leave permuted_data_tensor uninitialized, so the count + skip_data_expand=True return an empty permuted_data_tensor, so the count must come from a populated tensor (e.g. permuted_row_to_unpermuted_row_tensor.shape[0]). """ total_tokens = num_permuted_tokens @@ -808,11 +851,10 @@ def __init__( # max_num_tokens = ((mtp+1)*max_batch_size+max_isl+128+63)//64*64 = 9344 # moe_max_num_tokens = max_num_tokens * 2 = 18688 # It can avoid OOM for 8k/1k cases. - default_moe_max_num_tokens = 18688 - if model_config.moe_max_num_tokens > default_moe_max_num_tokens: - model_config._frozen = False - model_config.moe_max_num_tokens = default_moe_max_num_tokens - model_config._frozen = True + # Preserve an explicit deployment value. Only clamp the derived + # default, which keeps the existing OOM-safe behavior when the user + # does not size the DeepGEMM workspace deliberately. + _configure_deepgemm_moe_max_num_tokens(model_config) super().__init__( routing_method=routing_method, @@ -962,20 +1004,16 @@ def run_moe( assert token_selected_experts is not None assert token_final_scales is not None - # Permutation. - # skip_data_expand=True computes the permutation maps but skips the - # data-copy step (expandInputRowsKernel), so permuted_data_tensor and - # permuted_token_final_scales_tensor are returned with UNINITIALIZED - # contents (still full-size, just never written). The fused expand+quant - # kernel re-derives the activations from x via - # permuted_row_to_unpermuted_row_tensor instead, so all unused outputs are - # discarded with `_`. + # Permutation. skip_data_expand=True computes the permutation maps but + # skips the data copy, so the expanded activation and scale returns come + # back empty. The fused expand+quant kernel re-derives the activations + # from x via permuted_row_to_unpermuted_row_tensor instead. ( permuted_row_to_unpermuted_row_tensor, _, # permuted_token_selected_experts_tensor (unused) - _, # permuted_data_tensor (uninitialized under skip_data_expand) + _, # permuted_data_tensor (empty under skip_data_expand) expert_first_token_offset_tensor, - _, # permuted_token_final_scales_tensor (uninitialized under skip_data_expand) + _, # permuted_token_final_scales_tensor (empty under skip_data_expand) unpermuted_row_to_permuted_row_tensor, ) = torch.ops.trtllm.moe_permute_op( x, diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 90828ccd6f5d..d03fbc4175c3 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -15,6 +15,7 @@ import inspect import math +import os import threading from abc import ABC, abstractmethod from enum import Enum, auto @@ -231,6 +232,42 @@ class EplbSupportStatus(Enum): NOT_VERIFIED = auto() +def _host_weight_pageout_mode() -> Optional[str]: + """Select how a consumed expert's host weight pages are reclaimed. + + A loaded MoE layer leaves its source weight pages resident on the host, + charged to the cgroup the ranks share, so on a large checkpoint they pile + up long after the device weights are ready. This picks the ``madvise`` mode + applied once an expert's weights have been consumed; it covers only + contiguous CPU expert projections and biases, and ``madvise`` is a hint, + not a free or an unmap. + + Integrated GPUs share memory with the CPU and always get ``"dontneed"``; + ``TLLM_PAGEOUT_HOST_WEIGHTS`` does not affect them. Discrete GPUs default + to ``None``, because ranks loading the same shards reuse those cached + pages; setting the variable (``True``/``true``/``1``/``yes``/``y``) opts + into ``"pageout"``. + + The opt-in path must use ``MADV_PAGEOUT`` rather than ``MADV_DONTNEED``: + a mapper can rebuild a weight into anonymous memory, where ``DONTNEED`` + discards the contents and a later read sees zeros. ``MADV_PAGEOUT`` is + non-destructive for file-backed and anonymous pages alike. + + Enable it when ranks share a constrained cgroup and host memory is what + runs out. It complements the worker limit rather than replacing it, and + does not pay off everywhere: one measured layout saw a large drop in node + peak, another saw none and loaded slightly slower. Hence opt-in. + """ + if is_device_integrated(): + return "dontneed" + if os.environ.get("TLLM_PAGEOUT_HOST_WEIGHTS", + "False") in ["True", "true", "1", "yes", "y"]: + logger.info_once("Releasing host weight pages after each expert load", + key="moe_host_weight_pageout") + return "pageout" + return None + + class FusedMoEMethodBase(ABC): """ Base class for all fused MoE methods. @@ -346,18 +383,19 @@ def load_expert_weights_to_dst( w2_kargs["allow_partial_loading"] = allow_partial_loading pass_expert_idx_w3w1 = "expert_idx" in w3_w1_args + pageout_mode = _host_weight_pageout_mode() + def maybe_pageout_mmapped_cpu_weights( weight_tensors: List[object]) -> None: - # Integrated GPU systems share physical memory with CPU. After we - # finish copying from mmapped CPU weights, proactively advising the - # kernel to drop those pages reduces shared-memory pressure. - if not is_device_integrated(): + # Once an expert's weights have been copied to the device, advise + # the kernel to release the host pages backing them. + if pageout_mode is None: return for weight in weight_tensors: if (isinstance(weight, torch.Tensor) and weight.device.type == "cpu" and weight.is_contiguous()): - advise_tensor_pageout(weight) + advise_tensor_pageout(weight, mode=pageout_mode) # Multithread weight load is superseded by prefetch_files() in model_engine.py # Also, threading adds overhead in order to protect shuffle index cache with critical section. diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 96ef5223b7bb..df4bc47be9f4 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -31,6 +31,112 @@ REPLAY_WORK_PNAT = 2 REPLAY_WORK_CACHE_BUF_IDX = 3 REPLAY_WORK_ITEM_WIDTH = 4 +_FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE = 256 + + +@triton.jit +def _prepare_gdn_replay_work_items_kernel( + state_indices, + prev_num_accepted_tokens, + cache_buf_idx, + work_items, + n_writes_output, + num_decodes, + replay_step_width: tl.constexpr, + replay_history_size: tl.constexpr, + work_item_width: tl.constexpr, + position_field: tl.constexpr, + cache_slot_field: tl.constexpr, + pnat_field: tl.constexpr, + cache_buf_idx_field: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Build the write-first GDN replay partition in one launch.""" + offsets = tl.arange(0, BLOCK_SIZE) + active = offsets < num_decodes + slots = tl.load(state_indices + offsets, mask=active, other=0) + pnat = tl.load(prev_num_accepted_tokens + slots, mask=active, other=0) + active_buffer = tl.load(cache_buf_idx + slots, mask=active, other=0) + writes = active & (pnat + replay_step_width > replay_history_size) + writes_i32 = writes.to(tl.int32) + inclusive_write_offsets = tl.cumsum(writes_i32, axis=0) + write_offsets = inclusive_write_offsets - writes_i32 + n_writes = tl.sum(writes_i32, axis=0) + no_write_offsets = offsets - write_offsets + output_offsets = tl.where(writes, write_offsets, + n_writes + no_write_offsets) + output_base = work_items + output_offsets * work_item_width + tl.store(output_base + position_field, offsets, mask=active) + tl.store(output_base + cache_slot_field, slots, mask=active) + tl.store(output_base + pnat_field, pnat, mask=active) + tl.store(output_base + cache_buf_idx_field, active_buffer, mask=active) + tl.store(n_writes_output, n_writes) + + +def _build_replay_work_items_triton(state_indices, prev_num_accepted_tokens, + cache_buf_idx, work_items, n_writes, + replay_step_width, replay_history_size): + """Single-launch build of the write-first replay partition. + + Interchangeable with :func:`_build_replay_work_items_torch`; the caller + picks between them. Kept to one CTA because the write-first offsets come + from an in-block ``tl.cumsum``. + """ + num_decodes = state_indices.shape[0] + _prepare_gdn_replay_work_items_kernel[(1, )]( + state_indices, + prev_num_accepted_tokens, + cache_buf_idx, + work_items, + n_writes, + num_decodes, + replay_step_width=replay_step_width, + replay_history_size=replay_history_size, + work_item_width=REPLAY_WORK_ITEM_WIDTH, + position_field=REPLAY_WORK_POSITION_IN_DECODE_BATCH, + cache_slot_field=REPLAY_WORK_CACHE_SLOT, + pnat_field=REPLAY_WORK_PNAT, + cache_buf_idx_field=REPLAY_WORK_CACHE_BUF_IDX, + BLOCK_SIZE=triton.next_power_of_2(num_decodes), + num_warps=4, + ) + + +def _build_replay_work_items_torch(state_indices, prev_num_accepted_tokens, + cache_buf_idx, work_items, n_writes, + replay_step_width, replay_history_size): + """Same partition as :func:`_build_replay_work_items_triton`, in ATen ops. + + Keep field order and write-first partitioning in sync with the AutoDeploy + replay metadata path in shim/interface.py. + """ + num_decodes = state_indices.shape[0] + position_in_decode_batch = torch.arange(num_decodes, + dtype=torch.int32, + device=state_indices.device) + cache_slot_idx = state_indices.to(torch.long) + pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + + writes = (pnat + replay_step_width > replay_history_size) + writes_i32 = writes.to(torch.int32) + write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 + batch_n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) + no_write_offsets = position_in_decode_batch - write_offsets + output_offsets = torch.where(writes, write_offsets, + batch_n_writes + no_write_offsets) + output_offsets = output_offsets.to(torch.long) + + decode_work_items = work_items[:num_decodes] + decode_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH].scatter_( + 0, output_offsets, position_in_decode_batch) + decode_work_items[:, + REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, + state_indices) + decode_work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) + decode_work_items[:, REPLAY_WORK_CACHE_BUF_IDX].scatter_( + 0, output_offsets, active_cache_buf_idx) + n_writes.copy_(batch_n_writes) @triton.jit @@ -299,8 +405,9 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, self.replay_num_decodes = num_decodes if num_decodes == 0: return - if getattr(kv_cache_manager, 'use_gdn_cached_replay_all_layer_commit', - False): + use_gdn_all_layer_commit = getattr( + kv_cache_manager, "use_gdn_cached_replay_all_layer_commit", False) + if use_gdn_all_layer_commit: from tensorrt_llm._torch.modules.fla.cached_replay import \ CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE @@ -310,7 +417,6 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: return - self.replay_n_writes.zero_() if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): raise RuntimeError( "Replay state update is enabled, but the KV cache manager " @@ -327,34 +433,20 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, replay_step_width = replay_metadata.replay_step_width replay_history_size = replay_metadata.replay_history_size - position_in_decode_batch = torch.arange( - num_decodes, dtype=torch.int32, device=self.state_indices.device) - cache_slot = self.state_indices[num_contexts:batch_size] - cache_slot_idx = cache_slot.to(torch.long) - pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) - active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) - - # Keep field order and write-first partitioning in sync with the - # AutoDeploy replay metadata path in shim/interface.py. - writes = (pnat + replay_step_width > replay_history_size) - writes_i32 = writes.to(torch.int32) - write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 - n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) - no_write_offsets = position_in_decode_batch - write_offsets - output_offsets = torch.where(writes, write_offsets, - n_writes + no_write_offsets) - output_offsets = output_offsets.to(torch.long) - - work_items = self.replay_work_items[:num_decodes] - work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH].scatter_( - 0, output_offsets, position_in_decode_batch) - work_items[:, REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, - cache_slot) - work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) - work_items[:, - REPLAY_WORK_CACHE_BUF_IDX].scatter_(0, output_offsets, - active_cache_buf_idx) - self.replay_n_writes.copy_(n_writes) + if (use_gdn_all_layer_commit + and num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE): + build_work_items = _build_replay_work_items_triton + else: + build_work_items = _build_replay_work_items_torch + build_work_items( + self.state_indices[num_contexts:batch_size], + prev_num_accepted_tokens, + cache_buf_idx, + self.replay_work_items, + self.replay_n_writes, + replay_step_width, + replay_history_size, + ) def prepare(self, attn_metadata: AttentionMetadata): batch_size = attn_metadata.seq_lens.shape[0] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 73184a0d046d..7838491f1e8a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -664,6 +664,37 @@ def __init__( if self.attention_dp_enable_balance: self.attention_dp_time_out_iters = self.llm_args.attention_dp_config.timeout_iters self.attention_dp_batching_wait_iters = self.llm_args.attention_dp_config.batching_wait_iters + # ADP balancing withholds context for up to attention_dp_config.timeout_iters + # so all ranks start prefill together, which keeps a few ranks' large prefill + # chunks from setting the DP-padded MoE workload for every rank. When output + # lengths vary, ranks drain asynchronously and that saturated-load timeout + # leaves ranks holding decode work waiting on an alignment that may not come. + # + # TLLM_ADP_BALANCE_MIN_GENERATION_REQUESTS: the per-rank count of real + # decode work -- attention-DP pad dummies excluded -- below which the + # batch counts as under-occupied and the shorter timeout applies. + # This is the only switch; 0 (default) disables the path. Keep it in + # [0, max_batch_size]. + # TLLM_ADP_BALANCE_LOW_OCCUPANCY_TIMEOUT_ITERS: that shorter timeout, in + # scheduler iterations. 0 (default) releases context on the first + # under-occupied check. Keep it in [0, timeout_iters]. + self.attention_dp_min_generation_requests = int( + os.environ.get("TLLM_ADP_BALANCE_MIN_GENERATION_REQUESTS", 0)) + if not 0 <= self.attention_dp_min_generation_requests <= max_batch_size: + raise ValueError( + "TLLM_ADP_BALANCE_MIN_GENERATION_REQUESTS must be between " + f"0 and max_batch_size ({max_batch_size})") + self.attention_dp_low_occupancy_timeout_iters = int( + os.environ.get("TLLM_ADP_BALANCE_LOW_OCCUPANCY_TIMEOUT_ITERS", + 0)) + if not (0 <= self.attention_dp_low_occupancy_timeout_iters <= + self.attention_dp_time_out_iters): + # Above the configured timeout it would wait longer precisely + # when the system is less occupied, which inverts the feature. + raise ValueError( + "TLLM_ADP_BALANCE_LOW_OCCUPANCY_TIMEOUT_ITERS must be " + "between 0 and attention_dp_config.timeout_iters " + f"({self.attention_dp_time_out_iters})") self.batch_wait_timeout_ms = self.llm_args.batch_wait_timeout_ms self.batch_wait_timeout_iters = self.llm_args.batch_wait_timeout_iters self.batch_wait_max_tokens_ratio = self.llm_args.batch_wait_max_tokens_ratio @@ -5995,6 +6026,13 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest], balanced_context_requests = context_requests num_scheduled_context_requests = len(context_requests) num_scheduled_generation_requests = len(generation_requests) + # The low-occupancy test below asks how much real decode work the least + # busy rank has, and an attention-DP pad dummy is not decode work: a + # rank that is padded holds one generation request and no useful work. + # The count above keeps counting the scheduled batch, because the + # predicates built from it are about batch occupancy, not usefulness. + num_real_generation_requests = sum(1 for req in generation_requests + if not req.is_attention_dp_dummy) num_scheduled_tokens = sum( [len(req.get_tokens(0)) for req in context_requests]) + num_scheduled_generation_requests @@ -6002,8 +6040,11 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest], # balance the requests across DP ranks; not CP ranks within those DP ranks. responses_list = self.dist.tp_allgather([ num_scheduled_context_requests, num_scheduled_generation_requests, - num_scheduled_tokens + num_scheduled_tokens, num_real_generation_requests ]) + all_ranks_num_real_generation_requests = [ + response[3] for response in responses_list + ] all_ranks_num_scheduled_context_requests = [ response[0] for response in responses_list ] @@ -6036,7 +6077,14 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest], else: self.adp_ctx_waiting_iters_count += 1 balanced_context_requests = [] - timeout_reached = self.adp_ctx_waiting_iters_count >= self.attention_dp_time_out_iters + # Shorten the wait while the least busy rank is under-occupied; + # see the constructor for what the two knobs mean and how to + # tune them. + timeout_iters = self.attention_dp_time_out_iters + if (min(all_ranks_num_real_generation_requests) + < self.attention_dp_min_generation_requests): + timeout_iters = self.attention_dp_low_occupancy_timeout_iters + timeout_reached = self.adp_ctx_waiting_iters_count >= timeout_iters if timeout_reached or not all_ranks_have_gen_requests: self.adp_ctx_waiting_iters_count = 0 balanced_context_requests = context_requests diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 186af2cb6243..70a04185aa3c 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -3027,3 +3027,92 @@ def test_one_model_mtp_preserves_sampler_draft_tokens(self) -> None: assert gen.py_draft_tokens == sampler_drafts assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS + + +class TestAdpBalanceExcludesPadDummies: + """The low-occupancy test must look at real decode work, not batch size. + + `_pad_attention_dp_dummy_request` gives an otherwise idle rank exactly one + generation request so attention DP can make progress. Counting that dummy + as decode work makes the rank look busy, which is precisely wrong for a + check whose job is to notice idle ranks. + """ + + @staticmethod + def _make_executor(per_rank, threshold): + """per_rank: list of (num_ctx, num_gen, num_real) as allgathered.""" + executor = object.__new__(PyExecutor) + executor.dist = Mock() + executor.dist.tp_allgather = Mock( + return_value=[[ctx, gen, ctx + gen, real] for ctx, gen, real in per_rank] + ) + executor.max_batch_size = 64 + executor.attention_dp_enable_balance = True + executor.attention_dp_time_out_iters = 60 + executor.attention_dp_batching_wait_iters = 10 + executor.attention_dp_min_generation_requests = threshold + executor.attention_dp_low_occupancy_timeout_iters = 0 + executor.adp_ctx_waiting_iters_count = 0 + executor.adp_ctx_batching_wait_iters_count = 0 + return executor + + @staticmethod + def _make_generation_request(is_dummy): + request = Mock() + request.is_attention_dp_dummy = is_dummy + return request + + @staticmethod + def _make_context_request(num_tokens=4): + request = Mock() + request.get_tokens.return_value = [0] * num_tokens + return request + + def test_padded_rank_counts_as_under_occupied(self): + """A rank holding only a pad dummy has zero real decode work. + + Peer rank is padded: it reports one scheduled generation request and + zero real ones. With a threshold of 1 the low-occupancy timeout of 0 + must be selected so the withheld context is released on this very + iteration. Counting the dummy would leave the cross-rank minimum at 1, + keep the 60-iteration timeout, and strand the context batch. + """ + context_requests = [self._make_context_request()] + # This rank has context and decode work; the peer is padded and has + # no context request, so the batch is not aligned and the balancer + # takes the timeout branch. + executor = self._make_executor([(1, 4, 4), (0, 1, 0)], threshold=1) + + balanced = executor._balance_adp_requests( + context_requests, [self._make_generation_request(False)] + ) + + assert balanced == context_requests + assert executor.adp_ctx_waiting_iters_count == 0 + + def test_busy_ranks_still_wait(self): + """No rank is under-occupied, so the configured timeout still applies.""" + executor = self._make_executor([(1, 4, 4), (0, 4, 4)], threshold=1) + + balanced = executor._balance_adp_requests( + [self._make_context_request()], [self._make_generation_request(False)] + ) + + assert balanced == [] + assert executor.adp_ctx_waiting_iters_count == 1 + + def test_real_count_excludes_dummies(self): + """The allgathered real count must not include the pad dummy.""" + executor = self._make_executor([(1, 2, 1)], threshold=0) + + executor._balance_adp_requests( + [self._make_context_request()], + [ + self._make_generation_request(True), + self._make_generation_request(False), + ], + ) + + gathered = executor.dist.tp_allgather.call_args[0][0] + assert gathered[1] == 2, "scheduled count keeps counting the dummy" + assert gathered[3] == 1, "real count must exclude the dummy" diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py index 36c9d8b0c9f7..753bc5b7613a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py @@ -18,9 +18,13 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import Qwen3_5MoeForCausalLM, Qwen3_5MoeVLModel from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM -from tensorrt_llm._torch.models.modeling_qwen3_5 import _normalize_qwen35_moe_vl_config +from tensorrt_llm._torch.models.modeling_qwen3_5 import ( + _filter_language_model_weights, + _normalize_qwen35_moe_vl_config, +) from tensorrt_llm._torch.pyexecutor.config_utils import ( extract_mamba_kv_cache_params, load_pretrained_config, @@ -198,6 +202,23 @@ def test_qwen35_moe_model_defaults( assert llm_args.nvfp4_gemm_config.allowed_backends == expected_gemm_backends +def test_qwen35_vl_filter_preserves_consumable_weights() -> None: + language_weight = torch.tensor([1.0]) + weights = ConsumableWeightsDict( + { + "model.language_model.layers.0.weight": language_weight, + "model.visual.patch_embed.weight": torch.tensor([2.0]), + } + ) + + filtered_weights = _filter_language_model_weights(weights) + + assert isinstance(filtered_weights, ConsumableWeightsDict) + assert len(weights) == 0 + assert len(filtered_weights) == 1 + assert filtered_weights["model.language_model.layers.0.weight"] is language_weight + + def test_qwen35_moe_vl_placeholder_metadata_registered() -> None: metadata = MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata("qwen3_5_moe") diff --git a/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py b/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py new file mode 100644 index 000000000000..0702aed4b326 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper + + +class _Mapper(BaseWeightMapper): + """Concrete stand-in; the methods under test do not use the abstract half.""" + + def map_weights(self) -> None: + pass + + def apply_callbacks(self, module, module_name, module_names_breakdown, weights) -> list[dict]: + raise NotImplementedError + + +def _scan(weights: dict, prefix: str) -> dict: + """The pre-index implementation, kept as the reference for filter_prefix.""" + return {k[len(prefix) + 1 :]: v for k, v in weights.items() if k.startswith(prefix)} + + +def test_filter_prefix_matches_the_scan_it_replaces(): + plain = { + "model.layers.1.weight": 1, + "model.layers.1.self_attn.bias": 2, + "model.layers.10.weight": 3, + "model.norm.weight": 4, + } + weights = ConsumableWeightsDict(dict(plain)) + + for prefix in ("model", "model.layers.1", "model.layers.10", "model.norm", "absent"): + assert weights.filter_prefix(prefix) == _scan(plain, prefix), prefix + + +def test_filter_prefix_spans_keys_that_sort_past_the_separator(): + """The subtree bound must not stop at a key ordered above ``prefix.``.""" + weights = ConsumableWeightsDict({"a.weight": 1, "a/other": 2, "azz": 3, "b.weight": 4}) + + assert weights.filter_prefix("a") == _scan( + {"a.weight": 1, "a/other": 2, "azz": 3, "b.weight": 4}, "a" + ) + + +def test_index_tracks_added_and_consumed_keys(): + weights = ConsumableWeightsDict({"a.weight": 1}) + + assert weights.filter_prefix("a") == {"weight": 1} + weights["a.bias"] = 2 + weights.update({"a.scale": 3, "b.weight": 4}) + assert weights.filter_prefix("a") == {"weight": 1, "bias": 2, "scale": 3} + assert weights.mark_consumed("a") == 3 + assert weights.filter_prefix("a") == {} + assert weights.filter_prefix("b") == {"weight": 4} + + +def test_index_stays_correct_across_deletion_and_clear(): + """Deletions leave the index stale on purpose; it must stay a superset.""" + weights = ConsumableWeightsDict({"a.weight": 1, "a.bias": 2, "b.weight": 3}) + + assert weights.filter_prefix("a") == {"weight": 1, "bias": 2} + del weights["a.weight"] + assert weights.mark_consumed_keys(["a.bias", "missing"]) == 1 + assert weights.filter_prefix("a") == {} + weights.clear() + assert weights.mark_consumed("b") == 0 + + +def test_take_ownership_empties_a_consumable_source(): + """The derived mapping aliases the source, so the source must let go. + + Without this the loader holds a second reference to every tensor and + mark_consumed() on the derived mapping frees nothing -- the whole + checkpoint stays resident for the length of the load. + """ + source = ConsumableWeightsDict({"a.weight": 1, "b.weight": 2}) + + result = ConsumableWeightsDict.take_ownership(source, {"renamed.a.weight": 1}) + + assert isinstance(result, ConsumableWeightsDict) + assert len(source) == 0 + assert result.mark_consumed("renamed.a") == 1 + + +def test_take_ownership_leaves_a_plain_dict_source_alone(): + """A plain dict was never releasing incrementally; do not surprise its owner.""" + source = {"a.weight": 1} + derived = {"renamed.a.weight": 1} + + result = ConsumableWeightsDict.take_ownership(source, derived) + + assert result is derived + assert source == {"a.weight": 1} + + +def test_rename_by_params_map_transfers_ownership(): + tensor = object() + weights = ConsumableWeightsDict( + {"model.language_model.layer.weight": tensor, "lm_head.weight": object()} + ) + + renamed = _Mapper().rename_by_params_map( + {r"^model\.language_model\.(.*)$": r"model.\1"}, weights + ) + + assert isinstance(renamed, ConsumableWeightsDict) + assert len(weights) == 0 + assert renamed["model.layer.weight"] is tensor + assert "lm_head.weight" in renamed + + +def test_rename_by_params_map_leaves_a_plain_dict_alone(): + weights = {"model.language_model.layer.weight": object()} + + renamed = _Mapper().rename_by_params_map( + {r"^model\.language_model\.(.*)$": r"model.\1"}, weights + ) + + assert not isinstance(renamed, ConsumableWeightsDict) + assert len(weights) == 1 diff --git a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py index e1558d0aab4a..42face75bdc8 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py +++ b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py @@ -96,6 +96,45 @@ def _alloc_outputs(num_experts: int, m_max: int, hidden: int, *, device: str): return output_q, output_s +@skip_unsupported +def test_skip_data_expand_omits_unused_outputs() -> None: + num_rows, hidden, num_experts, top_k = 4, 512, 8, 4 + x = torch.randn((num_rows, hidden), device="cuda", dtype=torch.float32) + token_selected_experts = torch.arange(top_k, device="cuda", dtype=torch.int32).repeat( + num_rows, 1 + ) + token_final_scales = torch.full( + (num_rows, top_k), 1.0 / top_k, device="cuda", dtype=torch.float32 + ) + + outputs = torch.ops.trtllm.moe_permute_op( + x, + token_selected_experts, + token_final_scales, + None, + None, + None, + input_sf=None, + num_experts_on_rank=num_experts, + tp_size=1, + tp_rank=0, + ep_size=1, + ep_rank=0, + cluster_size=1, + cluster_rank=0, + min_latency_mode=False, + use_fp8_block_scaling=False, + skip_data_expand=True, + ) + + permuted_row_to_unpermuted_row_tensor = outputs[0] + permuted_data_tensor = outputs[2] + permuted_token_final_scales_tensor = outputs[4] + assert permuted_row_to_unpermuted_row_tensor.numel() == num_rows * top_k + assert permuted_data_tensor.shape == (0, hidden) + assert permuted_token_final_scales_tensor.shape == (0,) + + @skip_unsupported @pytest.mark.parametrize("shape", SHAPES, ids=lambda s: s.name) def test_fused_expand_quant_matches_unfused(shape: ExpandQuantShape) -> None: diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 8aa2a847802b..9c7169e36d73 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -19,12 +19,15 @@ import pytest import torch +from tensorrt_llm._torch.modules.mamba import mamba2_metadata from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( REPLAY_WORK_CACHE_BUF_IDX, REPLAY_WORK_CACHE_SLOT, REPLAY_WORK_PNAT, REPLAY_WORK_POSITION_IN_DECODE_BATCH, Mamba2Metadata, + _build_replay_work_items_torch, + _build_replay_work_items_triton, cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) @@ -39,6 +42,44 @@ ) +class _GdnReplayCacheManager: + use_replay_state_update = True + use_gdn_cached_replay_all_layer_commit = True + + def __init__(self, prev_num_accepted_tokens, cache_buf_idx): + self.prev_num_accepted_tokens = prev_num_accepted_tokens + self.cache_buf_idx = cache_buf_idx + + def get_replay_state_update_metadata(self): + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=6, + replay_history_size=MIN_REPLAY_HISTORY_SIZE, + ) + + +def _torch_reference_work_items(state_indices, prev_num_accepted_tokens, cache_buf_idx): + """Run the production ATen path into fresh buffers. + + The Triton kernel's contract is that it reproduces this path exactly, so + this path -- not a third hand-written copy -- is what it is compared to. + """ + num_decodes = state_indices.numel() + work_items = torch.zeros(num_decodes, 4, dtype=torch.int32, device="cuda") + n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") + _build_replay_work_items_torch( + state_indices, + prev_num_accepted_tokens, + cache_buf_idx, + work_items, + n_writes, + 6, + MIN_REPLAY_HISTORY_SIZE, + ) + return work_items, n_writes + + @skip_no_cuda class TestCuSeqlensToChunkIndicesOffsets: """Tests for cu_seqlens_to_chunk_indices_offsets_triton function.""" @@ -155,6 +196,134 @@ def get_replay_state_update_metadata(self): assert actual[0, REPLAY_WORK_PNAT] == 11 assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 + @pytest.mark.parametrize("num_decodes", [16, 17, 40, 255, 256]) + def test_replay_work_items_triton_matches_torch(self, num_decodes): + """The two builders must agree everywhere the dispatch may pick either. + + num_decodes 17 and 255 are not powers of two, so the kernel runs with + masked-off lanes -- those must contribute nothing to the write-first + prefix sum. + """ + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + state_indices = torch.randperm(num_slots, device="cuda")[:num_decodes].to(torch.int32) + + triton_items = torch.zeros(num_decodes, 4, dtype=torch.int32, device="cuda") + triton_n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") + _build_replay_work_items_triton( + state_indices, + prev_num_accepted_tokens, + cache_buf_idx, + triton_items, + triton_n_writes, + 6, + MIN_REPLAY_HISTORY_SIZE, + ) + torch_items, torch_n_writes = _torch_reference_work_items( + state_indices, prev_num_accepted_tokens, cache_buf_idx + ) + + torch.testing.assert_close(triton_items, torch_items) + torch.testing.assert_close(triton_n_writes, torch_n_writes) + + def test_prepare_replay_work_items_uses_the_selected_builder(self): + """The entry point must feed the builder the decode slice, not the whole batch.""" + num_contexts, num_decodes = 3, 40 + num_slots = num_contexts + num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + state_indices = torch.randperm(num_slots, device="cuda")[: num_contexts + num_decodes].to( + torch.int32 + ) + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + metadata = Mamba2Metadata(max_batch_size=num_contexts + num_decodes, chunk_size=8) + metadata.state_indices[: num_contexts + num_decodes].copy_(state_indices) + + metadata._prepare_replay_work_items(manager, num_contexts + num_decodes, num_contexts) + expected_items, expected_n_writes = _torch_reference_work_items( + state_indices[num_contexts:], prev_num_accepted_tokens, cache_buf_idx + ) + + torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) + torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + + @pytest.mark.parametrize( + ("num_decodes", "expect_fused"), + [ + pytest.param(8, False, id="below-partition-min"), + pytest.param(16, True, id="partition-min"), + pytest.param(256, True, id="fused-max"), + pytest.param(257, False, id="above-fused-max"), + ], + ) + def test_prepare_gdn_replay_work_items_dispatch_boundary( + self, monkeypatch, num_decodes, expect_fused + ): + """Pin which batch sizes reach the single-launch kernel. + + The equivalence tests above pass on either side of this boundary, so + without this the fused launch could silently stop being used. + """ + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) + metadata.state_indices.copy_(torch.arange(num_decodes, dtype=torch.int32, device="cuda")) + + launches = [] + original_kernel = mamba2_metadata._prepare_gdn_replay_work_items_kernel + + class _CountingKernel: + def __getitem__(self, grid): + launches.append(grid) + return original_kernel[grid] + + monkeypatch.setattr( + mamba2_metadata, "_prepare_gdn_replay_work_items_kernel", _CountingKernel() + ) + + metadata._prepare_replay_work_items(manager, num_decodes, 0) + + assert bool(launches) is expect_fused + + def test_prepare_gdn_replay_work_items_cuda_graph_replay(self): + num_decodes = 40 + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) + metadata.state_indices.copy_(torch.arange(num_decodes, dtype=torch.int32, device="cuda")) + + metadata._prepare_replay_work_items(manager, num_decodes, 0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + metadata._prepare_replay_work_items(manager, num_decodes, 0) + + updated_state_indices = torch.arange( + num_slots - 1, + num_slots - num_decodes - 1, + -1, + dtype=torch.int32, + device="cuda", + ) + metadata.state_indices.copy_(updated_state_indices) + prev_num_accepted_tokens.copy_( + (torch.arange(num_slots, dtype=torch.int32, device="cuda") * 7) % 21 + ) + cache_buf_idx.bitwise_xor_(1) + graph.replay() + torch.cuda.synchronize() + + expected_items, expected_n_writes = _torch_reference_work_items( + updated_state_indices, prev_num_accepted_tokens, cache_buf_idx + ) + torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) + torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + def test_single_sequence_unaligned(self): """Test with a single sequence that doesn't align with chunk size.""" cu_seqlens = torch.tensor([0, 10], dtype=torch.int, device="cuda")