[TRTLLM-15448][perf] Add demand-ordered rank-striped read-ahead

docs/source/features/checkpoint-loading.md

@@ -110,13 +110,27 @@ checkpoint bytes reach host memory:
   is still queued or active can no longer improve first-token latency and is
   stopped. A worker already blocked in a synchronous `pread` must finish its
   current 8 MiB read before shutdown, so degraded storage can delay the join.
+- `demand_ordered_rank_striped_read_ahead` is an explicit experimental variant
+  that uses mapper demand hints to prioritize complete read-ahead chunks across
+  node-local ranks. It retains every chunk, existing full-checkpoint memory
+  admission, and native materialization. It does not add consumption-based
+  pacing, destination-selective reads, or a streaming materializer. `auto`
+  continues to select the original rank-striped policy.
 
 ```yaml
 checkpoint_format: HF
 load_format: auto
 checkpoint_io_policy: rank_striped_read_ahead
 ```
 
+Policy IDs and display metadata are centralized in
+`tensorrt_llm/_checkpoint_io_policy.py`. `auto` is a selector, not a separate
+loading implementation. Display names may evolve without changing the stable
+configuration IDs; future modes become selectable only when implemented.
+The generated LLM-args telemetry manifest records the accepted IDs for tools
+that must not import TRT-LLM. Eligibility, fallback, and CI allocation remain
+separate from the catalog.
+
 The optimized path currently requires identical policy configuration across
 all ranks, the automatically constructed built-in HF loader
 (`checkpoint_loader` must be unset), `load_format: auto`, SafeTensors, and an
@@ -157,6 +171,50 @@ TRT-LLM instances. Policy logs distinguish requested, selected, activated, and
 effective policy, so `auto` selection and native fallback remain observable;
 activation logs also report the local reader assignment.
 
+### Demand-ordering qualification
+
+Select `checkpoint_io_policy: demand_ordered_rank_striped_read_ahead` in the
+normal HF configuration. The same eligibility restrictions apply to both
+rank-striped policies. Demand planning runs after mapper initialization and
+before the new readers are released. Opaque plans use physical order; qualified
+hints prioritize earlier demands while preserving physical locality among
+equal priorities and retaining an unmatched physical-order tail. Advisory
+ordering is not proof that source data can be omitted.
+
+The generic mapper prioritizes explicit native preload groups; concurrently
+loaded modules share a priority. It does not infer an exact tensor-consumption
+timeline or impose a serial layer order. Unsupported ordering stays opaque.
+
+For a planning-cost ablation, set
+`TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER=1` on every rank. This still
+builds and validates the demand plan but applies physical read order. Leave it
+unset for demand ordering. The switch affects only the experimental policy.
+
+Compare four cases using the same TRT-LLM revision, runtime image, checkpoint,
+GPU/node, parallel configuration, and controlled source-cache conditions:
+
+| Case | Policy | Physical-order ablation |
+| --- | --- | --- |
+| Native control | `native` | Unset |
+| Original rank-striped | `rank_striped_read_ahead` | Unset |
+| Planning overhead | `demand_ordered_rank_striped_read_ahead` | `1` |
+| Demand ordering | `demand_ordered_rank_striped_read_ahead` | Unset |
+
+Repeat cold and warm trials with counterbalanced order. Start with Qwen,
+Llama Maverick, and a preprocessing-heavy model; include TP and PP/EP cases.
+Check generation/weight parity, memory peaks, effective policy and applied
+order, and net model-loading latency. Do not infer a speedup from read bandwidth
+or weight population alone. The experimental policy includes its additional
+planning/activation in preparation; existing native/v1 phase boundaries are
+unchanged. Compare total model-loading time as well as the checkpoint phases,
+so unclassified work such as mapper initialization is not hidden. These metrics
+do not measure process launch to first token.
+
+Explicit policy configurations bypass the regular CI `auto`/`native`
+assignment. Keep demand-ordered observations separate from the original
+rank-striped treatment and retain the physical-order ablation setting with
+benchmark results.
+
 ### CI startup experiment
 
 Pre-merge perf-sanity runs that upload telemetry explicitly assign eligible

tensorrt_llm/_checkpoint_io_policy.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Checkpoint I/O names and display metadata, independent of loader execution."""
+
+from typing import Final, Literal, NamedTuple, TypeAlias, get_args
+
+CheckpointIoPolicy: TypeAlias = Literal[
+    "auto", "native", "rank_striped_read_ahead", "demand_ordered_rank_striped_read_ahead"
+]
+
+AUTO_IO_POLICY: Final = "auto"
+NATIVE_IO_POLICY: Final = "native"
+RANK_STRIPED_IO_POLICY: Final = "rank_striped_read_ahead"
+DEMAND_ORDERED_IO_POLICY: Final = "demand_ordered_rank_striped_read_ahead"
+
+CHECKPOINT_IO_POLICIES: tuple[CheckpointIoPolicy, ...] = get_args(CheckpointIoPolicy)
+RANK_STRIPED_IO_POLICIES = (RANK_STRIPED_IO_POLICY, DEMAND_ORDERED_IO_POLICY)
+EXECUTABLE_CHECKPOINT_IO_POLICIES = (NATIVE_IO_POLICY,) + RANK_STRIPED_IO_POLICIES
+
+
+class CheckpointIoPolicyInfo(NamedTuple):
+    """Human-readable metadata; eligibility and dispatch remain in the loader."""
+
+    display_name: str
+    description: str
+    kind: Literal["selector", "implementation"]
+
+
+CHECKPOINT_IO_POLICY_INFO: dict[CheckpointIoPolicy, CheckpointIoPolicyInfo] = {
+    AUTO_IO_POLICY: CheckpointIoPolicyInfo(
+        "Automatic selection",
+        "Select basic rank-striped read-ahead where eligible, otherwise native I/O.",
+        "selector",
+    ),
+    NATIVE_IO_POLICY: CheckpointIoPolicyInfo(
+        "Native checkpoint loading",
+        "Preserve the existing checkpoint I/O and materialization path.",
+        "implementation",
+    ),
+    RANK_STRIPED_IO_POLICY: CheckpointIoPolicyInfo(
+        "Rank-cooperative read-ahead",
+        "Read complete checkpoint chunks across node-local ranks while native materialization runs.",
+        "implementation",
+    ),
+    DEMAND_ORDERED_IO_POLICY: CheckpointIoPolicyInfo(
+        "Demand-ordered rank-cooperative read-ahead",
+        "Prioritize complete checkpoint chunks using native materialization demand hints.",
+        "implementation",
+    ),
+}

tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py

@@ -82,6 +82,16 @@ def is_weights_preloaded(self) -> bool:
         """Whether the last load wrote weights directly into the model."""
         return False
 
+    def activate_weight_session(
+            self,
+            readiness_error: BaseException | None = None,
+            *,
+            weight_mapper: BaseWeightMapper | None = None) -> None:
+        """Activate deferred advisory work before model materialization."""
+        del weight_mapper
+        if readiness_error is not None:
+            raise readiness_error
+
     def is_post_transform_weights_preloaded(self) -> bool:
         """Whether the last direct preload delivered post-transform weights.
 

tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py

@@ -1,7 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 from abc import ABC, abstractmethod
+from bisect import bisect_left
 from collections.abc import Callable, Mapping
 
 import torch
@@ -13,7 +15,8 @@
 from tensorrt_llm._torch.models.checkpoints.weight_load_plan import (
     WeightDemand, WeightLoadOrderConfidence, WeightLoadPlan,
     WeightLoadPlanCoverage)
-from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM
+from tensorrt_llm._torch.models.modeling_utils import (DecoderModelForCausalLM,
+                                                       is_moe_weight_owner)
 
 
 class BaseWeightMapper(ABC):
@@ -86,18 +89,202 @@ def cleanup(self) -> None:
 
     def build_weight_load_plan(self,
                                catalog: CheckpointCatalog) -> WeightLoadPlan:
-        """Build a conservative shadow plan for this rank.
-
-        The generic mapper cannot infer model-specific fusion, transformation,
-        or ordering dependencies. Treat the entire catalog as one indivisible
-        demand so consumers may use the plan for all-source read-ahead, but not
-        destination-aware selective I/O. Specialized mappers may override this
-        method to return an exact plan.
+        """Build a conservative, source-neutral load plan for this rank.
+
+        This method mirrors the native module walk without invoking any weight
+        transformation or loading callback. It uses module structure, explicit
+        fusion mappings, and preload hints to identify likely consumer groups.
+        The inference is advisory: every source tensor is retained exactly once,
+        unmatched tensors remain in physical catalog order at the tail, and the
+        plan never permits selective I/O. Specialized mappers may override this
+        method when they can provide stronger guarantees.
         """
         if self._model is None:
             raise ValueError(
                 "weight mapper must be initialized before building a load plan")
         mapping = self._model.model_config.mapping
+
+        # Catalog order is the source adapter's physical order. Preserve it
+        # within each inferred group and in the unmatched tail.
+        physical_names = tuple(tensor.name for tensor in catalog.tensors)
+        if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL",
+                          "False") in ("True", "true", "1", "yes", "y"):
+            # The serial override bypasses preload grouping in the native
+            # materializer; do not apply the concurrent-path ordering hints.
+            return self._build_opaque_weight_load_plan(catalog, physical_names)
+        physical_index = {
+            name: index
+            for index, name in enumerate(physical_names)
+        }
+        sorted_names = tuple(sorted(physical_names))
+
+        named_modules = getattr(self._model, "named_modules", None)
+        if named_modules is None:
+            return self._build_opaque_weight_load_plan(catalog, physical_names)
+
+        modules = tuple(named_modules(remove_duplicate=False))
+        remaining_names = set(physical_names)
+        demands: list[WeightDemand] = []
+
+        def take_prefixes(prefixes: tuple[str, ...]) -> tuple[str, ...]:
+            """Take unassigned tensor names below structural module prefixes."""
+            matched_indexes: set[int] = set()
+            for prefix in prefixes:
+                if not prefix:
+                    continue
+                if prefix in remaining_names:
+                    matched_indexes.add(physical_index[prefix])
+
+                descendant_prefix = f"{prefix}."
+                index = bisect_left(sorted_names, descendant_prefix)
+                while (index < len(sorted_names)
+                       and sorted_names[index].startswith(descendant_prefix)):
+                    name = sorted_names[index]
+                    if name in remaining_names:
+                        matched_indexes.add(physical_index[name])
+                    index += 1
+
+            result = tuple(physical_names[index]
+                           for index in sorted(matched_indexes))
+            remaining_names.difference_update(result)
+            return result
+
+        def take_exact(names: tuple[str, ...]) -> tuple[str, ...]:
+            matched_indexes = {
+                physical_index[name]
+                for name in names
+                if name in remaining_names and name in physical_index
+            }
+            result = tuple(physical_names[index]
+                           for index in sorted(matched_indexes))
+            remaining_names.difference_update(result)
+            return result
+
+        def source_module_path(module_path: str, module: nn.Module) -> str:
+            """Mirror native path normalization without touching payloads."""
+            parts = module_path.split('.') if module_path else []
+            if (parts and parts[-1] == "backend"
+                    and is_moe_weight_owner(module)):
+                return '.'.join(parts[:-1])
+            return module_path
+
+        def infer_source_names(module_path: str,
+                               module: nn.Module) -> tuple[str, ...]:
+            # Match the materializer's early exit: only modules that own at
+            # least one parameter slot can be consumers in the native walk.
+            if not module._parameters or self.should_skip_module(module_path):
+                return ()
+
+            normalized_path = source_module_path(module_path, module)
+            module_parts = normalized_path.split('.') if normalized_path else []
+            module_name = module_parts[-1] if module_parts else ""
+            parent_parts = module_parts[:-1]
+
+            # A fused destination consumes raw tensors from explicitly named
+            # sibling subtrees, for example q_proj/k_proj/v_proj -> qkv_proj.
+            if module_name and self.does_require_special_handling(module_name):
+                prefixes = tuple('.'.join(parent_parts + [source_name])
+                                 for source_name in self.mapping[module_name])
+                return take_prefixes(prefixes)
+
+            # Custom module loaders receive the full structurally matched
+            # subtree. Do not treat the root as a subtree consumer: an empty
+            # prefix contains the whole checkpoint and conveys no useful order.
+            # Other model-specific path rewrites that are not expressed by the
+            # mapper (for example a checkpoint-only namespace prefix)
+            # intentionally miss here and remain in the safe, all-source
+            # unmatched tail.
+            if normalized_path and (self.is_special_instance_module(module)
+                                    or hasattr(module, 'load_weights')):
+                return take_prefixes((normalized_path, ))
+
+            # The native fallback copies only direct parameters. Descendant
+            # parameters belong to their own named-module consumers.
+            parameter_names = tuple(
+                f"{normalized_path}.{name}" if normalized_path else name
+                for name, _ in module.named_parameters(recurse=False))
+            return take_exact(parameter_names)
+
+        preload_suffixes = tuple(
+            getattr(self._model, "preload_weight_modules", None) or ())
+        preload_modules: list[tuple[str, nn.Module]] = []
+        preload_paths: set[str] = set()
+        for suffix in preload_suffixes:
+            for module_path, module in modules:
+                if (module_path.endswith(suffix)
+                        and module_path not in preload_paths):
+                    preload_paths.add(module_path)
+                    preload_modules.append((module_path, module))
+
+        previous_preload_group: str | None = None
+        for module_path, module in preload_modules:
+            source_names = infer_source_names(module_path, module)
+            if not source_names:
+                continue
+            group_id = f"preload:{len(demands):06d}:{module_path}"
+            predecessors = ((previous_preload_group, )
+                            if previous_preload_group is not None else ())
+            demands.append(
+                WeightDemand(
+                    group_id=group_id,
+                    source_names=source_names,
+                    destination_ranks=(mapping.rank, ),
+                    priority=len(demands),
+                    predecessors=predecessors,
+                ))
+            previous_preload_group = group_id
+
+        # Native materialization schedules non-preloaded modules concurrently.
+        # Equal priorities let read-ahead preserve physical locality without
+        # inventing dependencies between these consumers.
+        nonpreload_priority = len(demands)
+        for module_path, module in modules:
+            if module_path in preload_paths:
+                continue
+            source_names = infer_source_names(module_path, module)
+            if not source_names:
+                continue
+            demands.append(
+                WeightDemand(
+                    group_id=
+                    f"module:{len(demands):06d}:{module_path or '<root>'}",
+                    source_names=source_names,
+                    destination_ranks=(mapping.rank, ),
+                    priority=nonpreload_priority,
+                ))
+
+        # If structural inspection found nothing useful, expose no ordering
+        # claim. Otherwise retain every unresolved tensor as a physical tail.
+        if not demands:
+            return self._build_opaque_weight_load_plan(catalog, physical_names)
+
+        unmatched_names = tuple(name for name in physical_names
+                                if name in remaining_names)
+        if unmatched_names:
+            demands.append(
+                WeightDemand(
+                    group_id="unmatched_checkpoint_tensors",
+                    source_names=unmatched_names,
+                    destination_ranks=(mapping.rank, ),
+                    priority=nonpreload_priority + 1,
+                ))
+
+        plan = WeightLoadPlan(
+            catalog_id=catalog.catalog_id,
+            rank=mapping.rank,
+            world_size=mapping.world_size,
+            coverage=WeightLoadPlanCoverage.CONSERVATIVE,
+            ordering=WeightLoadOrderConfidence.ADVISORY,
+            demands=tuple(demands),
+        )
+        plan.validate_against(catalog)
+        return plan
+
+    def _build_opaque_weight_load_plan(
+            self, catalog: CheckpointCatalog,
+            physical_names: tuple[str, ...]) -> WeightLoadPlan:
+        """Build an all-source plan when no useful ordering can be inferred."""
+        mapping = self.model.model_config.mapping
         plan = WeightLoadPlan(
             catalog_id=catalog.catalog_id,
             rank=mapping.rank,
@@ -106,7 +293,7 @@ def build_weight_load_plan(self,
             ordering=WeightLoadOrderConfidence.OPAQUE,
             demands=(WeightDemand(
                 group_id="all_checkpoint_tensors",
-                source_names=tuple(sorted(catalog.tensor_names)),
+                source_names=physical_names,
                 destination_ranks=(mapping.rank, ),
             ), ),
         )