[None][perf] Qwen3.5/3.8 wave-2: MoE, attention-DP, GDN replay, weight loading

cpp/tensorrt_llm/thop/moeUtilOp.cpp

@@ -138,9 +138,13 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Te
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
     auto permuted_token_selected_experts_tensor
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
-    auto permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false));
+    // Skipping the expand leaves these two unwritten, so size them to zero
+    // rather than to the expanded token count: at experts_per_token * num_rows
+    // rows they dominate this operator's allocation.
+    int64_t const num_expanded_rows = skip_data_expand ? 0 : num_moe_inputs;
+    auto permuted_data_tensor = torch::empty({num_expanded_rows, hidden_size}, input.options().requires_grad(false));
     auto permuted_token_final_scales_tensor
-        = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+        = torch::empty({num_expanded_rows}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
     auto expert_first_token_offset_tensor = torch::empty(
         {num_experts_per_node + 1}, torch::dtype(torch::kInt64).device(torch::kCUDA).requires_grad(false));
     auto unpermuted_row_to_permuted_row_tensor = torch::empty({static_cast<int64_t>(experts_per_token * num_rows)},

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."""

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)

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):

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",
         }

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())