[TRTLLM-16104][test] Add ModelExpress post-merge accuracy canaries

docs/source/features/model-express.md

@@ -93,7 +93,8 @@ Support for another model family requires a focused qualification change:
 5. Run a real ModelExpress donor/receiver test with the model configurations
    being claimed, including the supported quantization and TP/PP/EP layouts.
    Compare deterministic output token IDs with the standard Hugging Face load
-   path before documenting the family as supported.
+   path, and require the per-rank weight manifests to match (see the
+   qualification test below), before documenting the family as supported.
 
 ### Qualification Test
 
@@ -104,6 +105,35 @@ uses a metadata-only view of the donor's canonical snapshot and contains no
 weight shards. A positive result therefore requires direct transfer; disk
 fallback cannot accidentally satisfy the test.
 
+Each role generates eight fixed prompts for exactly 32 greedy tokens
+(`end_id=-1`, so no sequence stops early), and the test requires exact
+token-ID equality of the donor and the receiver against the HF baseline.
+
+**Weight manifests.** Every rank also writes a SHA-256 manifest of all
+registered parameters and buffers
+(`tensorrt_llm/_torch/weight_sharing/weight_manifest.py`) when
+`MX_WEIGHT_MANIFEST_DIR` is set; the harness sets it for all three roles and
+production loads never write one. Two manifest families are compared byte for
+byte:
+
+- `manifest.final.<role>.rank<N>.json` is written at the end of
+  `ModelLoader.load`, after every post-load hook and MoE load-balancer
+  finalization. Baseline, donor, and receiver must be pairwise identical:
+  same tensor names, dtypes, shapes, strides, storage offsets, digests,
+  skipped-tensor sets, and storage-alias partitions.
+- `manifest.transfer.<role>.rank<N>.json` is written inside the MX checkpoint
+  loader at the donor's publish point and at the receiver's P2P success point
+  (MX roles only). Donor and receiver parameters must be identical at this
+  boundary; derived buffers are enforced at the final tier because the
+  receiver's `cache_derived_state()` runs after the transfer.
+
+A row may list `final_manifest_exempt_patterns` on its `MxE2ECase` to exempt
+named tensors from the final-tier digest comparison. That is never a numeric
+tolerance and never applies to the transfer tier, and every pattern needs a
+code comment explaining why; the current BF16 dense rows exempt nothing.
+Manifests carry a `manifest_format_version`, and manifests of different
+versions never compare.
+
 Run the TP=1 smoke test against an isolated ModelExpress 0.4.1 service with
 NIXL enabled:
 
@@ -112,9 +142,16 @@ TRTLLM_MX_E2E_REQUIRED=1 \
 MODEL_EXPRESS_URL=http://127.0.0.1:8001 \
 LLM_MODELS_ROOT=/path/to/llm-models \
 pytest -v tests/integration/defs/model_express/test_model_express.py \
-  -k llama-bf16-tp1
+  -k llama-bf16-tp1 --output-dir /path/to/artifacts
 ```
 
+With `--output-dir` (always set in CI), the worker payloads, worker logs,
+transfer logs, weight manifests, and `timing.json` are copied to
+`model_express/<case-id>/` under that directory, so they are part of the stage
+results archive even when the test fails. The test also prints one
+`MX E2E timing` line per role and rank with the load, generation, and manifest
+durations.
+
 Run the TP=2 rank-mapping qualification on four GPUs by selecting
 `llama-bf16-tp2`. `TRTLLM_MX_LLAMA_MODEL` can override the default TinyLlama
 checkpoint path. The Qwen2 and Qwen3 profile rows use `qwen2-bf16-tp1` /
@@ -150,6 +187,57 @@ or NIXL prerequisites fail instead of skipping. Do not add every model profile
 to recurring coverage: use the harness for representative rows claimed by the
 support table and keep wider matrices in scheduled qualification.
 
+### Accuracy Canaries (Post-Merge)
+
+`tests/integration/defs/model_express/test_model_express_accuracy.py` runs one
+reference-backed accuracy task per qualified family on an MX receiver. The donor
+publishes exactly as in the smoke test and never evaluates; the receiver starts
+from the metadata-only snapshot, self-checks its own transfer logs before
+spending any evaluation time (a fallback exits with status 3), evaluates the
+task with `tensorrt_llm.evaluate` inside its own subprocess, and writes the
+score to JSON. The pytest process never constructs an `LLM`: it loads the
+accuracy reference YAMLs, asserts the same hypothesis-testing threshold as
+`tests/integration/defs/accuracy/`, and additionally requires the transfer
+evidence and the donor/receiver weight manifests to match. There is no paired
+HF baseline evaluation because the reference value is that baseline.
+
+Current rows (all TP=1, BF16, `references/*.yaml` hold the expected values):
+
+| Test ID | Model | Task | Model path override |
+| --- | --- | --- | --- |
+| `llama3-8b-instruct-mmlu-tp1` | `meta-llama/Meta-Llama-3-8B-Instruct` | MMLU | `TRTLLM_MX_LLAMA3_8B_MODEL` |
+| `qwen2.5-7b-instruct-mmlu-tp1` | `Qwen/Qwen2.5-7B-Instruct` | MMLU | `TRTLLM_MX_QWEN25_MODEL` |
+| `qwen3-8b-gsm8k-tp1` | `Qwen3/Qwen3-8B` | GSM8K | `TRTLLM_MX_QWEN3_MODEL` |
+
+The rows are registered as `stage: post_merge` entries in
+`tests/integration/test_lists/test-db/l0_model_express.yml`, so they run in
+`DGX_H100-2_GPUs-PyTorch-ModelExpress-Post-Merge-1` on every main commit and
+never in pre-merge pipelines. Trigger the stage on a pull request with:
+
+```text
+/bot run --stage-list "DGX_H100-2_GPUs-PyTorch-ModelExpress-Post-Merge-1"
+```
+
+GSM8K needs the `lm_eval` package from `requirements-dev.txt`; the test checks
+for it and, under `TRTLLM_MX_E2E_REQUIRED=1`, fails instead of skipping when it
+is absent. Each run records the donor and receiver load times, the evaluation
+time, the score, and the threshold as junit properties and as
+`model_express_accuracy/<test-id>.json` under `--output-dir`; load times are
+observed only, not gated. To run one canary locally:
+
+```bash
+TRTLLM_MX_E2E_REQUIRED=1 \
+MODEL_EXPRESS_URL=http://127.0.0.1:8001 \
+LLM_MODELS_ROOT=/path/to/llm-models \
+pytest -v tests/integration/defs/model_express/test_model_express_accuracy.py \
+  -k llama3-8b-instruct-mmlu-tp1 --output-dir /path/to/artifacts
+```
+
+Adding a canary is one `MxAccuracyCase` row (the model must be inside the
+family's qualified runtime envelope and have a bare reference entry for the
+task), one line in the `post_merge` block of `l0_model_express.yml`, and a row
+in the table above.
+
 ### Transform-Layout ABI Rules
 
 An existing transform-layout ABI ID is immutable. Introduce a new ID when a

jenkins/L0_Test.groovy

@@ -6104,7 +6104,7 @@ def launchTestJobs(pipeline, testFilter, globalVars)
     // IMPORTANT: Stage Configuration Syntax Requirement
     //
     // The test_to_stage_mapping.py script expects stage definitions in the following format:
-    // "Stage-Name": ["platform", "yaml_file", splitId, split_count, gpu_count]
+    // "Stage-Name": ["platform", "yaml_file", splitId, split_count, gpu_count, modelExpress]
     //
     // Where:
     // - Stage-Name: Must be quoted string, used to identify the Jenkins stage
@@ -6113,6 +6113,9 @@ def launchTestJobs(pipeline, testFilter, globalVars)
     // - splitId: Current split number (1-based)
     // - split_count: Total number of splits
     // - gpu_count: Number of GPUs required (optional, defaults to 1)
+    // - modelExpress: Optional boolean; true attaches the Redis + ModelExpress server
+    //   sidecars (and the CI ModelExpress env) to the test pod. The mapping regex in
+    //   scripts/test_to_stage_mapping.py accepts trailing booleans.
     //
     // This format is parsed by scripts/test_to_stage_mapping.py to provide bidirectional
     // mapping between test names and Jenkins stage names. Any changes to this syntax
@@ -6135,6 +6138,8 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         // platform, test DB, split, splits, GPU count, ModelExpress sidecars
         "DGX_H100-2_GPUs-PyTorch-ModelExpress-1": ["dgx-h100-x4", "l0_model_express", 1, 1, 2, true],
         "DGX_H100-4_GPUs-PyTorch-ModelExpress-OnDemand-1": ["dgx-h100-x4", "l0_model_express", 1, 1, 4, true],
+        // Post-merge only: runs the `stage: post_merge` rows of l0_model_express (accuracy canaries).
+        "DGX_H100-2_GPUs-PyTorch-ModelExpress-Post-Merge-1": ["dgx-h100-x4", "l0_model_express", 1, 1, 2, true],
         "RTX5090-PyTorch-1": ["rtx-5090", "l0_gb202", 1, 1],
         "RTX5080-PyTorch-1": ["rtx-5080", "l0_gb203", 1, 2],
         "RTX5080-PyTorch-2": ["rtx-5080", "l0_gb203", 2, 2],

tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py

@@ -36,7 +36,17 @@
 from contextlib import contextmanager
 from enum import Enum
 from pathlib import Path
-from typing import Any, Callable, Iterator, MutableMapping, Optional, Protocol, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterator,
+    MutableMapping,
+    Optional,
+    Protocol,
+    Type,
+    Union,
+)
 
 from tensorrt_llm._torch.models.checkpoints.base_config_loader import BaseConfigLoader
 from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader
@@ -47,10 +57,14 @@
     IdentityCheckPolicy,
     SourceIdentity,
     check_weight_sharing_compatibility,
+    maybe_write_weight_manifest,
 )
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
+if TYPE_CHECKING:
+    from torch import nn
+
 # Defensive default for the upstream `MX_SOURCE_QUERY_TIMEOUT` env var.
 # The upstream `MxLiveWeightLoader` polls the MX server every 5 s for up
 # to `MX_SOURCE_QUERY_TIMEOUT` seconds (default 3600 = 1 hour) waiting
@@ -183,7 +197,7 @@ def _close_mx_client(client: Any) -> None:
 
 
 def _synchronize_cuda_for_mx_publish() -> None:
-    """Finish pending CUDA writes before exposing source buffers through MX."""
+    """Finish pending CUDA writes at an MX transfer boundary (publish or receive)."""
     import torch
 
     if torch.cuda.is_initialized():
@@ -200,6 +214,16 @@ def _enable_mx_transfer_logging() -> None:
         mx_logger.setLevel(logging.INFO)
 
 
+def _maybe_write_mx_transfer_manifest(model: "nn.Module", *, rank: int, boundary: str) -> None:
+    """Fingerprint the bytes at an MX transfer boundary when `MX_WEIGHT_MANIFEST_DIR` is set."""
+    maybe_write_weight_manifest(
+        model,
+        family="transfer",
+        rank=rank,
+        context={"boundary": boundary, "checkpoint_format": "MX"},
+    )
+
+
 @register_checkpoint_loader("MX")
 class MXCheckpointLoader(HfCheckpointLoader):
     """Checkpoint loader for MX (ModelExpress) P2P weight transfer.
@@ -573,6 +597,10 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[
             return fallback_weights
 
         self._p2p_succeeded = True
+        # P2P writes and any upstream dtype casts must be globally visible
+        # before the bytes are fingerprinted or finalized by ModelLoader.
+        _synchronize_cuda_for_mx_publish()
+        _maybe_write_mx_transfer_manifest(model, rank=mapping.rank, boundary="receiver_p2p_success")
         logger.info(
             "MX P2P weight transfer succeeded from %s",
             self._mx_server_url,
@@ -808,13 +836,19 @@ def publish_as_source(
                 "can still observe transient values. Tracked by MX-2.",
                 key="mx_publish_env_threaded_warning",
             )
+        # Post-load transforms may enqueue asynchronous writes. Make the source
+        # buffers globally ready before they are fingerprinted and before MX
+        # publishes their addresses and allows a receiver to issue RDMA reads.
+        # The manifest runs outside the best-effort publish guard below so a
+        # manifest problem is loud; nothing between here and the publish call
+        # touches the weights.
+        _synchronize_cuda_for_mx_publish()
+        _maybe_write_mx_transfer_manifest(
+            model, rank=source_identity.rank, boundary="donor_publish"
+        )
         try:
             with _MX_TRANSFER_STATE_LOCK:
                 resolved_name = self._resolve_publish_name(checkpoint_dir)
-                # Post-load transforms may enqueue asynchronous writes. Make
-                # the source buffers globally ready before MX publishes their
-                # addresses and allows a receiver to issue RDMA reads.
-                _synchronize_cuda_for_mx_publish()
                 with (
                     _temporary_env("MODEL_EXPRESS_URL", self._mx_server_url),
                     _temporary_env("MODEL_NAME", resolved_name),

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -23,7 +23,8 @@
     PostTransformProfile, PostTransformProfileRegistry,
     PostTransformQualificationDecision, PostTransformRuntimeConfig,
     PostTransformRuntimeConstraints, PostTransformTransferScope, SourceIdentity,
-    check_weight_sharing_compatibility)
+    WeightManifestWriteResult, check_weight_sharing_compatibility,
+    maybe_write_weight_manifest)
 from tensorrt_llm._utils import str_dtype_to_torch
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           ExecutorMemoryType,
@@ -377,6 +378,7 @@ class ModelLoaderMetricNames(Enum):
         "draft_checkpoint_preparation_seconds")
     DRAFT_WEIGHT_POPULATION_SECONDS = "draft_weight_population_seconds"
     POST_LOAD_PROCESSING_SECONDS = "post_load_processing_seconds"
+    WEIGHT_MANIFEST_SECONDS = "weight_manifest_seconds"
 
 
 class ModelLoader:
@@ -1146,6 +1148,13 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                 # perturb NVLink barrier synchronization in multi-rank DG init.
                 torch.cuda.empty_cache()
 
+        manifest_result = self._dump_final_weight_manifest(
+            checkpoint_loader, model, weights_preloaded=weights_preloaded)
+        if manifest_result is not None:
+            self._metrics[ModelLoaderMetricNames.WEIGHT_MANIFEST_SECONDS.
+                          value] = (manifest_result.build_seconds +
+                                    manifest_result.write_seconds)
+
         metrics = ", ".join(f"{name}={value:.4f}"
                             for name, value in self._metrics.items())
         logger.info(
@@ -1302,6 +1311,34 @@ def _qualify_post_transform_profile(
                 model.model_config, model=model),
         )
 
+    def _dump_final_weight_manifest(
+            self, checkpoint_loader: BaseCheckpointLoader,
+            model: DecoderModelForCausalLM, *,
+            weights_preloaded: bool) -> Optional[WeightManifestWriteResult]:
+        """Write the env-gated final-state weight manifest for this rank.
+
+        This is a no-op unless `MX_WEIGHT_MANIFEST_DIR` is set. It runs after
+        every post-load hook and after MoE load-balancer finalization, and
+        before engine warmup, so the manifest describes the final post-load
+        state that all roles of the MX qualification harness share.
+        `reload()` is deliberately not covered: incremental weight updates own
+        their own lifecycle.
+        """
+        return maybe_write_weight_manifest(
+            model,
+            family="final",
+            rank=self.mapping.rank,
+            context={
+                "boundary": "model_loader_load_end",
+                "checkpoint_format": checkpoint_loader.checkpoint_format,
+                "weights_preloaded": bool(weights_preloaded),
+                "load_format": str(self.llm_args.load_format),
+                "tp_rank": self.mapping.tp_rank,
+                "pp_rank": self.mapping.pp_rank,
+                "world_size": self.mapping.world_size,
+                "model_class": type(model).__name__,
+            })
+
     def _post_load_publish(
             self, checkpoint_loader: BaseCheckpointLoader,
             model: DecoderModelForCausalLM, *, checkpoint_dir: str,
@@ -1454,7 +1491,8 @@ def reload(self,
         before rebinding fresh weights. Partial reloads keep existing transform
         guards intact because untouched modules may already contain transformed
         live weights. The owner of the update lifecycle is responsible for
-        running post-load processing once all bytes are present.
+        running post-load processing once all bytes are present. Weight
+        manifests are not written here; see `_dump_final_weight_manifest`.
 
         Args:
             model: Model instance receiving the replacement weights.

tensorrt_llm/_torch/weight_sharing/__init__.py

@@ -41,13 +41,40 @@
     SourceIdentityMismatchError,
     check_weight_sharing_compatibility,
 )
+from tensorrt_llm._torch.weight_sharing.weight_manifest import (
+    WEIGHT_MANIFEST_DIR_ENV,
+    WEIGHT_MANIFEST_FAMILIES,
+    WEIGHT_MANIFEST_FILE_PATTERN,
+    WEIGHT_MANIFEST_FORMAT_VERSION,
+    WEIGHT_MANIFEST_KINDS,
+    WEIGHT_MANIFEST_ROLE_ENV,
+    SkippedTensor,
+    WeightManifest,
+    WeightManifestDiff,
+    WeightManifestEntry,
+    WeightManifestWriteResult,
+    build_weight_manifest,
+    canonical_tensor_bytes,
+    compare_weight_manifests,
+    load_weight_manifest,
+    manifest_file_name,
+    maybe_write_weight_manifest,
+    serialize_weight_manifest,
+    write_weight_manifest,
+)
 
 __all__ = [
     "ARTIFACT_IDENTITY_FORMAT_VERSION",
     "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "SOURCE_IDENTITY_FORMAT_VERSION",
+    "WEIGHT_MANIFEST_DIR_ENV",
+    "WEIGHT_MANIFEST_FAMILIES",
+    "WEIGHT_MANIFEST_FILE_PATTERN",
+    "WEIGHT_MANIFEST_FORMAT_VERSION",
+    "WEIGHT_MANIFEST_KINDS",
+    "WEIGHT_MANIFEST_ROLE_ENV",
     "ArtifactIdentity",
     "IdentityCheckDecision",
     "IdentityCheckPolicy",
@@ -61,7 +88,20 @@
     "PostTransformRuntimeConfig",
     "PostTransformRuntimeConstraints",
     "PostTransformTransferScope",
+    "SkippedTensor",
     "SourceIdentity",
     "SourceIdentityMismatchError",
+    "WeightManifest",
+    "WeightManifestDiff",
+    "WeightManifestEntry",
+    "WeightManifestWriteResult",
+    "build_weight_manifest",
+    "canonical_tensor_bytes",
     "check_weight_sharing_compatibility",
+    "compare_weight_manifests",
+    "load_weight_manifest",
+    "manifest_file_name",
+    "maybe_write_weight_manifest",
+    "serialize_weight_manifest",
+    "write_weight_manifest",
 ]