[TRTLLM-12352][feat] complete MX post-transform qualification foundation

docs/source/features/model-express.md

@@ -20,11 +20,30 @@ to the provided Hugging Face checkpoint.
 
 ## Current Support Scope
 
-The post-transform MX receive path currently supports only
-`LlamaForCausalLM` with transform protocol version 1. TensorRT LLM publishes
-post-transform weights together with source-identity and layout metadata. A
-receiver whose model family is not allow-listed does not consume those bytes;
-it falls back to the standard Hugging Face checkpoint path.
+The post-transform MX receive path currently supports one exact qualification
+profile:
+
+| Profile | Root class | Config identity | Scope | Protocol | Transform-layout ABI | Constraints |
+|---------|------------|-----------------|-------|----------|----------------------|-------------|
+| `llama-for-causal-lm-target-v1` | `LlamaForCausalLM` | `LlamaForCausalLM` / `llama` | Target model | 1 | `trtllm-llama-target-layout-v1` | No speculative mode or separately loaded draft model |
+
+The registry matches the exact root class and the architecture/model type
+captured from the resolved config before model construction. An unregistered
+subclass or config alias does not inherit support. It falls back to the
+standard Hugging Face checkpoint path before any P2P transfer starts.
+
+TensorRT LLM applies two independent compatibility gates:
+
+- The qualification profile records that a model/config/lifecycle combination
+  has passed full-load versus staged-load equivalence testing.
+- `SourceIdentity` format version 3 binds two concrete runs to the same
+  checkpoint artifact, runtime layout choices, local shard layout, and
+  transform-layout ABI.
+
+The transfer protocol version identifies the staged receiver protocol. The
+transform-layout ABI identifies the meaning of the transferred tensor names,
+layouts, aliases, and receiver finalization. A pre-version-3 identity, a
+missing ABI, or a different ABI is rejected rather than treated as compatible.
 
 Loads that require a separately loaded draft model also fall back to the
 standard checkpoint path. Target-plus-draft post-transform transfer remains
@@ -42,16 +61,40 @@ Support for another model family requires a focused qualification change:
 2. Verify that every one-time transform is guarded by `_weights_transformed`
    and that the staged receiver can skip `transform_weights()` without
    changing aliases, derived state, tensor layout, or outputs.
-3. Add the model class and transform protocol version to the MX staged-receiver
-   allow-list only after full-load and staged-load equivalence tests pass.
+3. Add an exact qualification profile only after the reusable harness in
+   `tests/unittest/utils/post_transform_qualification.py` proves tensor,
+   alias, transform-guard, derived-state, and deterministic output
+   equivalence. Include an unregistered-root negative control.
 4. Cover compatible transfer, source-identity mismatch, unsupported layout or
-   protocol, and non-allow-listed fallback. Keep target-plus-draft loading
-   disabled unless that combination has its own mixed-layout tests.
+   protocol/ABI, no-disk staged reception, and unqualified-profile fallback.
+   Keep target-plus-draft loading disabled unless that combination has its own
+   mixed-layout tests.
 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.
 
+### Transform-Layout ABI Rules
+
+An existing transform-layout ABI ID is immutable. Introduce a new ID when a
+change affects any transferred tensor name, shape, dtype, packing, sharding,
+alias relationship, one-shot transform result, or receiver-side
+`setup_aliases()`/`cache_derived_state()` interpretation. Keep the existing ID
+for implementation-only changes that preserve all of those observable
+semantics.
+
+When adding an ABI ID:
+
+1. Give the qualified profile the new ID and propagate it through
+   `SourceIdentity` and MX source metadata.
+2. Add matching, missing, and mismatched producer/receiver compatibility
+   tests. ABI mismatches remain incompatible even under the `ENFORCE` identity
+   policy.
+3. Re-run the qualification harness and the real donor/receiver GPU test for
+   every profile that adopts the ID.
+4. Never reinterpret an already published ID. Supporting two ABIs requires an
+   explicit compatibility decision and tests for each producer/receiver pair.
+
 ## Installation
 
 The official TensorRT LLM release container includes the MX Python client. No
@@ -140,7 +183,7 @@ path.
 
 - Post-transform MX reception is currently limited to the Llama model family.
   Other model families safely fall back to Hugging Face loading until they are
-  explicitly qualified and added to the staged-receiver allow-list.
+  explicitly qualified and added as exact capability profiles.
 - The MX server and Redis lifecycle is external to TensorRT LLM. Every
   TensorRT LLM instance must be able to reach the configured MX server URL.
 - The MX server coordinates source discovery but does not store model weights.

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

@@ -66,6 +66,7 @@
 _MX_SOURCE_IDENTITY_METADATA_KEY = "trtllm_source_identity"
 _MX_WEIGHT_LAYOUT_METADATA_KEY = "trtllm_weight_layout"
 _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY = "trtllm_transform_protocol_version"
+_MX_TRANSFORM_ABI_ID_METADATA_KEY = "trtllm_transform_abi_id"
 _MX_WEIGHT_LAYOUT_POST_TRANSFORM = "post_transform"
 _MX_STAGED_TRANSFORM_PROTOCOL_VERSION = 1
 
@@ -425,13 +426,24 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[
                 **kwargs,
             )
 
-        layout_status = _metadata_weight_layout_status(source_metadata)
+        expected_transform_abi_id = (
+            self._local_source_identity.transform_abi_id
+            if self._local_source_identity is not None
+            else None
+        )
+        layout_status = _metadata_weight_layout_status(
+            source_metadata,
+            expected_transform_abi_id=expected_transform_abi_id,
+        )
         if layout_status is _MxWeightLayoutStatus.UNSUPPORTED:
             self._source_identity_compatible_for_last_load = False
             return self._fallback_to_disk(
                 checkpoint_dir,
                 mapping,
-                reason=_metadata_unsupported_layout_reason(source_metadata),
+                reason=_metadata_unsupported_layout_reason(
+                    source_metadata,
+                    expected_transform_abi_id=expected_transform_abi_id,
+                ),
                 **kwargs,
             )
 
@@ -721,6 +733,12 @@ def publish_as_source(
                 "unavailable; receivers cannot safely verify transformed weights."
             )
             return
+        if source_identity.transform_abi_id is None:
+            logger.warning(
+                "Skipping MX post-transform publish because SourceIdentity has "
+                "no qualified transform-layout ABI."
+            )
+            return
 
         try:
             from modelexpress import (
@@ -884,6 +902,8 @@ def _build_mx_source_metadata(source_identity: Optional[SourceIdentity]) -> dict
     }
     if source_identity is not None:
         metadata[_MX_SOURCE_IDENTITY_METADATA_KEY] = _serialize_source_identity(source_identity)
+        if source_identity.transform_abi_id is not None:
+            metadata[_MX_TRANSFORM_ABI_ID_METADATA_KEY] = source_identity.transform_abi_id
     return metadata
 
 
@@ -941,6 +961,7 @@ def _metadata_has_trtllm_key(metadata: dict[str, Any]) -> bool:
             _MX_SOURCE_IDENTITY_METADATA_KEY,
             _MX_WEIGHT_LAYOUT_METADATA_KEY,
             _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY,
+            _MX_TRANSFORM_ABI_ID_METADATA_KEY,
         )
     )
 
@@ -970,13 +991,25 @@ def _source_identity_from_metadata(metadata: Optional[dict[str, Any]]) -> Option
         return None
 
 
-def _metadata_is_post_transform(metadata: Optional[dict[str, Any]]) -> bool:
+def _metadata_is_post_transform(
+    metadata: Optional[dict[str, Any]],
+    *,
+    expected_transform_abi_id: Optional[str],
+) -> bool:
     return (
-        _metadata_weight_layout_status(metadata) is _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED
+        _metadata_weight_layout_status(
+            metadata,
+            expected_transform_abi_id=expected_transform_abi_id,
+        )
+        is _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED
     )
 
 
-def _metadata_weight_layout_status(metadata: Optional[dict[str, Any]]) -> _MxWeightLayoutStatus:
+def _metadata_weight_layout_status(
+    metadata: Optional[dict[str, Any]],
+    *,
+    expected_transform_abi_id: Optional[str],
+) -> _MxWeightLayoutStatus:
     layout = _metadata_get(metadata, _MX_WEIGHT_LAYOUT_METADATA_KEY)
     if layout is None:
         return _MxWeightLayoutStatus.PRE_TRANSFORM
@@ -994,16 +1027,41 @@ def _metadata_weight_layout_status(metadata: Optional[dict[str, Any]]) -> _MxWei
         return _MxWeightLayoutStatus.UNSUPPORTED
     if protocol_version != _MX_STAGED_TRANSFORM_PROTOCOL_VERSION:
         return _MxWeightLayoutStatus.UNSUPPORTED
+
+    source_transform_abi_id = _metadata_get(metadata, _MX_TRANSFORM_ABI_ID_METADATA_KEY)
+    if not isinstance(source_transform_abi_id, str) or not source_transform_abi_id:
+        return _MxWeightLayoutStatus.UNSUPPORTED
+    if expected_transform_abi_id is None or source_transform_abi_id != expected_transform_abi_id:
+        return _MxWeightLayoutStatus.UNSUPPORTED
     return _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED
 
 
-def _metadata_unsupported_layout_reason(metadata: Optional[dict[str, Any]]) -> str:
+def _metadata_unsupported_layout_reason(
+    metadata: Optional[dict[str, Any]],
+    *,
+    expected_transform_abi_id: Optional[str],
+) -> str:
     layout = _metadata_get(metadata, _MX_WEIGHT_LAYOUT_METADATA_KEY)
     if str(layout).lower() == _MX_WEIGHT_LAYOUT_POST_TRANSFORM:
         version = _metadata_get(metadata, _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY)
+        try:
+            protocol_version = int(version)
+        except (TypeError, ValueError):
+            protocol_version = None
+        if protocol_version != _MX_STAGED_TRANSFORM_PROTOCOL_VERSION:
+            return (
+                "source publishes post-transform weights with unsupported "
+                f"transform protocol {version!r}"
+            )
+
+        source_transform_abi_id = _metadata_get(metadata, _MX_TRANSFORM_ABI_ID_METADATA_KEY)
+        if not isinstance(source_transform_abi_id, str) or not source_transform_abi_id:
+            return "source publishes post-transform weights without a transform-layout ABI"
+        if expected_transform_abi_id is None:
+            return "receiver has no qualified transform-layout ABI for post-transform weights"
         return (
-            "source publishes post-transform weights with unsupported "
-            f"transform protocol {version!r}"
+            "source publishes post-transform weights with transform-layout ABI "
+            f"{source_transform_abi_id!r}; receiver requires {expected_transform_abi_id!r}"
         )
     return f"source publishes unsupported MX weight layout {layout!r}"