[TRTLLM-14881][feat] qualify Mistral dense for MX

docs/source/features/model-express.md

@@ -28,6 +28,7 @@ profiles:
 | `llama-for-causal-lm-target-v1` | `LlamaForCausalLM` | `LlamaForCausalLM` / `llama` | Target model | 1 | `trtllm-llama-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 | `qwen2-for-causal-lm-bf16-target-v1` | `Qwen2ForCausalLM` | `Qwen2ForCausalLM` / `qwen2` | Target model | 1 | `trtllm-qwen2-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 | `qwen3-for-causal-lm-bf16-target-v1` | `Qwen3ForCausalLM` | `Qwen3ForCausalLM` / `qwen3` | Target model | 1 | `trtllm-qwen3-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused QK-norm/RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
+| `mistral-for-causal-lm-bf16-target-v1` | `MistralForCausalLM` | `MistralForCausalLM` / `mistral` | Target model | 1 | `trtllm-mistral-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, no sliding window, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 
 The registry matches the exact root class, the architecture/model type captured
 from the resolved config before model construction, and any runtime constraints
@@ -40,6 +41,18 @@ the exact `Qwen2ForCausalLM` / `qwen2` pair and satisfy the constraints above.
 Other Qwen roots and variants, including checkpoints with tied embeddings, do
 not match this profile.
 
+The Mistral identity covers dense Hugging Face-format checkpoints that resolve
+to the exact `MistralForCausalLM` / `mistral` pair and whose realized attention
+runs without a sliding window on every layer; Mistral-7B-Instruct-v0.3 is the
+qualified canary. Checkpoints that enable `sliding_window`, including
+Ministral-style `layer_types` mixes, do not match this profile. YaRN scaling
+and tied embeddings are separate constraints that also exclude a checkpoint.
+The native `mistral` checkpoint format (`params.json` with
+`consolidated.safetensors`) is a separate `checkpoint_format` that rewrites the
+model type to `mistral_common`; it cannot be combined with the MX loading path.
+The unregistered Llama-based `MistralForCausalLM` class in `modeling_llama.py`
+shares components with the qualified root but is not qualified.
+
 TensorRT LLM applies two independent compatibility gates:
 
 - The qualification profile records that a model/config/lifecycle combination
@@ -58,18 +71,21 @@ standard checkpoint path. Target-plus-draft post-transform transfer remains
 disabled until layout state is tracked and qualified independently for each
 submodel.
 
-The Llama, Qwen2, and Qwen3 profiles are text-only and do not enable
+The Llama, Qwen2, Qwen3, and Mistral profiles are text-only and do not enable
 reward-model, embedding, MoE, or vision-language roots. FP16, quantized weights
 or KV cache, alternate attention backends, YaRN, tied embeddings, TP greater
 than 2, PP greater than 1, CP greater than 1, LoRA, sparse attention, attention
 DP, multi-node transfer, and speculative decoding require separate
 qualification rows. Each profile also pins its qualified RoPE realization:
-Llama and Qwen2 require the default fused RoPE path, so unfused RoPE requires
-separate qualification for them, while Qwen3 fuses RoPE into the QK-norm
-kernel and therefore requires `rope_fusion=False` in the realized
-configuration. The profiles do not constrain MoE-only backend and mapping
-settings because these dense roots do not consume them. `SourceIdentity`
-still requires donor and receiver configurations to match.
+Llama, Qwen2, and Mistral require the default fused RoPE path, so unfused RoPE
+requires separate qualification for them, while Qwen3 fuses RoPE into the
+QK-norm kernel and therefore requires `rope_fusion=False` in the realized
+configuration. The Mistral profile additionally pins the realized attention
+window: every attention layer must run without a sliding window, so
+sliding-window checkpoints such as Ministral fall back to the Hugging Face path
+until they are qualified separately. The profiles do not constrain MoE-only
+backend and mapping settings because these dense roots do not consume them.
+`SourceIdentity` still requires donor and receiver configurations to match.
 
 ### Adding a Model Family
 
@@ -117,9 +133,10 @@ pytest -v tests/integration/defs/model_express/test_model_express.py \
 
 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` /
-`qwen2-bf16-tp2` and `qwen3-bf16-tp1` / `qwen3-bf16-tp2`, with optional model
-path overrides in `TRTLLM_MX_QWEN2_MODEL` and `TRTLLM_MX_QWEN3_MODEL`.
+checkpoint path. The Qwen2, Qwen3, and Mistral profile rows use
+`qwen2-bf16-tp1` / `qwen2-bf16-tp2`, `qwen3-bf16-tp1` / `qwen3-bf16-tp2`, and
+`mistral-bf16-tp1` / `mistral-bf16-tp2`, with optional model path overrides in
+`TRTLLM_MX_QWEN2_MODEL`, `TRTLLM_MX_QWEN3_MODEL`, and `TRTLLM_MX_MISTRAL_MODEL`.
 `TRTLLM_MX_E2E_REQUIRED=1` converts missing service, model, or NIXL
 prerequisites from skips into failures and must be set by a CI qualification
 stage. That stage must also allocate the GPUs declared by the selected test
@@ -260,9 +277,13 @@ path.
 ## Notes and Limitations
 
 - Post-transform MX reception is currently limited to the exact Llama,
-  Qwen2/Qwen2.5 dense, and Qwen3 dense profiles above. Other roots and variants
-  that do not match the documented identity and runtime envelope safely fall
-  back to Hugging Face loading until explicitly qualified.
+  Qwen2/Qwen2.5 dense, Qwen3 dense, and Mistral dense profiles above. Other
+  roots and variants that do not match the documented identity and runtime
+  envelope safely fall back to Hugging Face loading until explicitly qualified.
+- Mistral checkpoints served through the native `mistral` checkpoint format
+  (`mistral_common` model type), sliding-window or Ministral `layer_types`
+  variants, YaRN variants, Mistral3 vision-language roots, and Mistral Large 3
+  are not qualified and use the Hugging Face fallback.
 - 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/pyexecutor/model_loader.py

@@ -17,6 +17,7 @@
 from tensorrt_llm._torch.peft.lora.config import LoraConfig
 from tensorrt_llm._torch.weight_sharing import (
     LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1,
+    MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity,
     IdentityCheckPolicy, LazyRootModelIdentity, PostTransformConfigIdentity,
@@ -85,6 +86,12 @@
     _MX_BF16_DENSE_RUNTIME_CONSTRAINTS,
     rope_fusion=frozenset({False}),
 )
+# Mistral realizes a per-layer attention window from its checkpoint config.
+# The dense profile is qualified only when every layer runs full attention.
+_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS = replace(
+    _MX_BF16_DENSE_RUNTIME_CONSTRAINTS,
+    sliding_windows=frozenset({"none"}),
+)
 
 
 def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig,
@@ -428,6 +435,19 @@ class ModelLoader:
             transfer_scope=PostTransformTransferScope.TARGET_MODEL,
             runtime_constraints=_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS,
         ),
+        PostTransformProfile(
+            profile_id="mistral-for-causal-lm-bf16-target-v1",
+            root_model_class=LazyRootModelIdentity(
+                "tensorrt_llm._torch.models.modeling_mistral",
+                "MistralForCausalLM"),
+            architecture="MistralForCausalLM",
+            model_type="mistral",
+            speculative_mode=None,
+            protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION,
+            transform_abi_id=MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
+            transfer_scope=PostTransformTransferScope.TARGET_MODEL,
+            runtime_constraints=_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS,
+        ),
     ))
 
     @classmethod

tensorrt_llm/_torch/weight_sharing/__init__.py

@@ -20,6 +20,7 @@
 )
 from tensorrt_llm._torch.weight_sharing.post_transform_profiles import (
     LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1,
+    MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     LazyRootModelIdentity,
@@ -46,6 +47,7 @@
 __all__ = [
     "ARTIFACT_IDENTITY_FORMAT_VERSION",
     "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1",
+    "MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "SOURCE_IDENTITY_FORMAT_VERSION",

tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py

@@ -48,6 +48,10 @@
 # Stable contract for unquantized Qwen3 dense fused-QKV and fused-gate-up
 # tensors, Q/K norm state, and target-only receiver finalization.
 QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-qwen3-dense-target-layout-v1"
+# Stable contract for unquantized Mistral dense fused-QKV and fused-gate-up
+# tensors plus target-only receiver finalization. Mistral shares the dense
+# decoder base class with Llama but is a distinct root with its own contract.
+MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-mistral-dense-target-layout-v1"
 _MISSING = object()
 
 
@@ -82,6 +86,34 @@ def _realized_rope_fusion(model: nn.Module) -> bool | None:
     return realized_values.pop() if len(realized_values) == 1 else None
 
 
+def _realized_sliding_window(model: nn.Module) -> str | None:
+    """Summarize the sliding-window attention realized by a constructed model.
+
+    Attention modules that can apply a window expose `attention_window_size`.
+    The result is `none` when every exposing module applies full attention,
+    `uniform` when all exposing modules share one positive window, `mixed` when
+    windowed and full-attention layers coexist or windows differ, and `None`
+    when no module exposes the attribute or a value is not a positive `int`.
+    """
+    realized_values: set[int | None] = set()
+    for module in model.modules():
+        value = getattr(module, "attention_window_size", _MISSING)
+        if value is _MISSING:
+            continue
+        if value is None:
+            realized_values.add(None)
+            continue
+        window_size = _canonical_int(value)
+        if window_size is None or window_size <= 0:
+            return None
+        realized_values.add(window_size)
+    if not realized_values:
+        return None
+    if realized_values == {None}:
+        return "none"
+    return "uniform" if len(realized_values) == 1 else "mixed"
+
+
 def _canonical_optional_string(container: object, attribute: str) -> str | None:
     value = getattr(container, attribute, _MISSING)
     if value is _MISSING:
@@ -139,6 +171,9 @@ class PostTransformRuntimeConfig:
     tied_word_embeddings: bool | None
     rope_type: str | None
     rope_fusion: bool | None
+    # Realized sliding-window attention: `none`, `uniform`, `mixed`, or `None`
+    # when the constructed model does not determine it.
+    sliding_window: str | None
 
     @classmethod
     def from_model_config(
@@ -172,6 +207,7 @@ def from_model_config(
             rope_type = None
 
         rope_fusion = _realized_rope_fusion(model) if model is not None else None
+        sliding_window = _realized_sliding_window(model) if model is not None else None
 
         return cls(
             dtype=_canonical_string(getattr(model_config, "torch_dtype", None)),
@@ -203,6 +239,7 @@ def from_model_config(
             ),
             rope_type=rope_type,
             rope_fusion=rope_fusion,
+            sliding_window=sliding_window,
         )
 
 
@@ -232,6 +269,7 @@ class PostTransformRuntimeConstraints:
         ("tied_word_embeddings", "tied_word_embeddings"),
         ("rope_types", "rope_type"),
         ("rope_fusion", "rope_fusion"),
+        ("sliding_windows", "sliding_window"),
     )
 
     dtypes: frozenset[str | None] | None = None
@@ -255,6 +293,7 @@ class PostTransformRuntimeConstraints:
     tied_word_embeddings: frozenset[bool | None] | None = None
     rope_types: frozenset[str | None] | None = None
     rope_fusion: frozenset[bool | None] | None = None
+    sliding_windows: frozenset[str | None] | None = None
 
     def __post_init__(self) -> None:
         for constraint_name, _runtime_name in self._DIMENSIONS:

tests/integration/defs/model_express/test_model_express.py

@@ -156,6 +156,26 @@ class MxE2ECase:
         id="qwen3-bf16-tp2",
         marks=pytest.mark.skip_less_device(4),
     ),
+    pytest.param(
+        MxE2ECase(
+            model_env="TRTLLM_MX_MISTRAL_MODEL",
+            default_model_subdir="Mistral-7B-Instruct-v0.3",
+            repository_cache_prefix="models--trtllm-mx-e2e--mistral-tp1",
+            tp_size=1,
+        ),
+        id="mistral-bf16-tp1",
+        marks=pytest.mark.skip_less_device(2),
+    ),
+    pytest.param(
+        MxE2ECase(
+            model_env="TRTLLM_MX_MISTRAL_MODEL",
+            default_model_subdir="Mistral-7B-Instruct-v0.3",
+            repository_cache_prefix="models--trtllm-mx-e2e--mistral-tp2",
+            tp_size=2,
+        ),
+        id="mistral-bf16-tp2",
+        marks=pytest.mark.skip_less_device(4),
+    ),
 )
 
 

tests/integration/test_lists/test-db/l0_model_express.yml

@@ -21,6 +21,7 @@ l0_model_express:
   - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1]
   - model_express/test_model_express.py::test_mx_donor_receiver[qwen2-bf16-tp1]
   - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-bf16-tp1]
+  - model_express/test_model_express.py::test_mx_donor_receiver[mistral-bf16-tp1]
 - condition:
     ranges:
       system_gpu_count:
@@ -39,3 +40,4 @@ l0_model_express:
   - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2]
   - model_express/test_model_express.py::test_mx_donor_receiver[qwen2-bf16-tp2]
   - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-bf16-tp2]
+  - model_express/test_model_express.py::test_mx_donor_receiver[mistral-bf16-tp2]