From 44691420d04b4bfb1a46898bfe713066ffc0d4f7 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:34:19 -0700 Subject: [PATCH 1/4] [TRTLLM-14880][feat] qualify Qwen3 dense for MX Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docs/source/features/model-express.md | 40 +-- .../_torch/pyexecutor/model_loader.py | 23 +- .../_torch/weight_sharing/__init__.py | 2 + .../weight_sharing/post_transform_profiles.py | 3 + .../defs/model_express/test_model_express.py | 40 +++ .../test_lists/test-db/l0_model_express.yml | 4 + .../_torch/executor/test_model_loader_mx.py | 282 ++++++++++++++++-- .../weight_sharing/test_source_identity.py | 2 + 8 files changed, 356 insertions(+), 40 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 21b89c656aa2..aa06b826bced 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -27,6 +27,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 | The registry matches the exact root class, the architecture/model type captured from the resolved config before model construction, and any runtime constraints @@ -57,14 +58,14 @@ standard checkpoint path. Target-plus-draft post-transform transfer remains disabled until layout state is tracked and qualified independently for each submodel. -The Llama and Qwen2 profiles are text-only and do not enable reward-model, MoE, -or vision-language roots. FP16, quantized weights or KV cache, alternate -attention backends, YaRN or unfused RoPE, 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. 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. +The Llama, Qwen2, and Qwen3 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 or unfused RoPE, 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. 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 @@ -112,12 +113,15 @@ 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. `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 row. `TRTLLM_MX_E2E_TIMEOUT_S` controls the 1200-second timeout -used for the baseline worker, receiver worker, and donor-readiness wait; -increase it for slow model storage or startup. +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`. +`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 +row. `TRTLLM_MX_E2E_TIMEOUT_S` controls the 1200-second timeout used for the +baseline worker, receiver worker, and donor-readiness wait; increase it for +slow model storage or startup. The dedicated H100 CI stages own isolated Redis and ModelExpress 0.4.1 sidecars. The two-GPU TP=1 stage is classified as multi-GPU: it runs @@ -251,10 +255,10 @@ path. ## Notes and Limitations -- Post-transform MX reception is currently limited to the exact Llama and - Qwen2/Qwen2.5 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. +- 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. - 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. diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9899ad7426ac..a4ce1ec088d9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -6,6 +6,7 @@ import os import traceback import warnings +from dataclasses import replace from enum import Enum from typing import Callable, Optional, Tuple @@ -16,7 +17,8 @@ from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.weight_sharing import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity, + QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity, IdentityCheckPolicy, PostTransformConfigIdentity, PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, PostTransformQualificationDecision, PostTransformRuntimeConfig, @@ -79,6 +81,10 @@ rope_types=frozenset({"default"}), rope_fusion=frozenset({True}), ) +_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS = replace( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + rope_fusion=frozenset({False}), +) def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, @@ -390,6 +396,7 @@ def _post_transform_profile_registry(cls) -> PostTransformProfileRegistry: if cls._POST_TRANSFORM_PROFILE_REGISTRY is None: from ..models.modeling_llama import LlamaForCausalLM from ..models.modeling_qwen import Qwen2ForCausalLM + from ..models.modeling_qwen3 import Qwen3ForCausalLM cls._POST_TRANSFORM_PROFILE_REGISTRY = PostTransformProfileRegistry( profiles=( PostTransformProfile( @@ -417,6 +424,20 @@ def _post_transform_profile_registry(cls) -> PostTransformProfileRegistry: transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, ), + PostTransformProfile( + profile_id="qwen3-for-causal-lm-bf16-target-v1", + root_model_class=Qwen3ForCausalLM, + architecture="Qwen3ForCausalLM", + model_type="qwen3", + speculative_mode=None, + protocol_version=cls. + _MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id= + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints= + _MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), )) return cls._POST_TRANSFORM_PROFILE_REGISTRY diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index 034f01e17881..1393258294de 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -21,6 +21,7 @@ from tensorrt_llm._torch.weight_sharing.post_transform_profiles import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, PostTransformConfigIdentity, PostTransformFeature, PostTransformProfile, @@ -46,6 +47,7 @@ "ArtifactIdentity", "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1", "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", + "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "SOURCE_IDENTITY_FORMAT_VERSION", "PostTransformConfigIdentity", "PostTransformFeature", diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index e56b6a5b3089..d5542008de0a 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py @@ -42,6 +42,9 @@ # Stable contract for unquantized Qwen2 dense fused-QKV and fused-gate-up # tensors plus the target-only receiver finalization used by its first profile. QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-qwen2-dense-target-layout-v1" +# 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" _MISSING = object() diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index f4edd3a1d4fb..e642a8ae108f 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -116,6 +116,46 @@ class MxE2ECase: id="llama-bf16-tp2", marks=pytest.mark.skip_less_device(4), ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_QWEN2_MODEL", + default_model_subdir="Qwen2-7B-Instruct", + repository_cache_prefix="models--trtllm-mx-e2e--qwen2-tp1", + tp_size=1, + ), + id="qwen2-bf16-tp1", + marks=pytest.mark.skip_less_device(2), + ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_QWEN2_MODEL", + default_model_subdir="Qwen2-7B-Instruct", + repository_cache_prefix="models--trtllm-mx-e2e--qwen2-tp2", + tp_size=2, + ), + id="qwen2-bf16-tp2", + marks=pytest.mark.skip_less_device(4), + ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_QWEN3_MODEL", + default_model_subdir="Qwen3/Qwen3-8B", + repository_cache_prefix="models--trtllm-mx-e2e--qwen3-tp1", + tp_size=1, + ), + id="qwen3-bf16-tp1", + marks=pytest.mark.skip_less_device(2), + ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_QWEN3_MODEL", + default_model_subdir="Qwen3/Qwen3-8B", + repository_cache_prefix="models--trtllm-mx-e2e--qwen3-tp2", + tp_size=2, + ), + id="qwen3-bf16-tp2", + marks=pytest.mark.skip_less_device(4), + ), ) diff --git a/tests/integration/test_lists/test-db/l0_model_express.yml b/tests/integration/test_lists/test-db/l0_model_express.yml index 69adb45bda1b..38eff49c28fe 100644 --- a/tests/integration/test_lists/test-db/l0_model_express.yml +++ b/tests/integration/test_lists/test-db/l0_model_express.yml @@ -19,6 +19,8 @@ l0_model_express: orchestrator: mpi tests: - 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] - condition: ranges: system_gpu_count: @@ -35,3 +37,5 @@ l0_model_express: orchestrator: mpi tests: - 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] diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 2b3a94874967..a6b7839c54a4 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -13,7 +13,7 @@ import pytest import torch from torch import nn -from transformers import LlamaConfig, Qwen2Config +from transformers import LlamaConfig, Qwen2Config, Qwen3Config from utils.post_transform_qualification import ( PostTransformQualificationCase, assert_post_transform_lifecycle_equivalent, @@ -24,6 +24,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_llama as modeling_llama_mod from tensorrt_llm._torch.models import modeling_qwen as modeling_qwen_mod +from tensorrt_llm._torch.models import modeling_qwen3 as modeling_qwen3_mod from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import MXCheckpointLoader from tensorrt_llm._torch.modules import mla as mla_mod from tensorrt_llm._torch.modules.linear import Linear, WeightMode @@ -148,6 +149,10 @@ class _UnqualifiedQwen2ForCausalLM(modeling_qwen_mod.Qwen2ForCausalLM): pass +class _UnqualifiedQwen3ForCausalLM(modeling_qwen3_mod.Qwen3ForCausalLM): + pass + + def _tiny_llama_model( monkeypatch: pytest.MonkeyPatch, *, @@ -267,6 +272,56 @@ def _tiny_qwen2_model( return model +def _tiny_qwen3_model( + monkeypatch: pytest.MonkeyPatch, + *, + model_class: type[nn.Module] = modeling_qwen3_mod.Qwen3ForCausalLM, + tp_size: int = 1, + rank: int = 0, +) -> nn.Module: + monkeypatch.setattr(torch.cuda, "Stream", lambda *_args, **_kwargs: MagicMock()) + monkeypatch.setattr(torch.cuda, "Event", lambda *_args, **_kwargs: MagicMock()) + qwen3_config = Qwen3Config( + architectures=["Qwen3ForCausalLM"], + attention_bias=False, + head_dim=4, + hidden_act="silu", + hidden_size=16, + intermediate_size=32, + max_position_embeddings=16, + mlp_bias=False, + num_attention_heads=4, + num_hidden_layers=2, + num_key_value_heads=2, + rms_norm_eps=1e-6, + rope_scaling=None, + tie_word_embeddings=False, + torch_dtype=torch.bfloat16, + vocab_size=32, + ) + model = model_class( + ModelConfig( + pretrained_config=qwen3_config, + mapping=mapping_mod.Mapping( + world_size=tp_size, + rank=rank, + tp_size=tp_size, + ), + max_num_tokens=16, + max_seq_len=16, + ) + ) + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + values = torch.arange( + parameter.numel(), + dtype=torch.float32, + device=parameter.device, + ).reshape(parameter.shape) + parameter.copy_(((values + index) % 17).to(parameter.dtype) / 17) + return model + + def _bf16_dense_runtime_config(**overrides: object) -> PostTransformRuntimeConfig: values = { "dtype": "bfloat16", @@ -310,7 +365,25 @@ def _qwen2_layout_state(model: nn.Module) -> dict[str, object]: } -def _qwen2_input_embeddings(model: nn.Module) -> torch.Tensor: +def _qwen3_layout_state(model: nn.Module) -> dict[str, object]: + layer = model.model.layers[0] + return { + "attention_type": type(layer.self_attn).__name__, + "qkv_weight_mode": layer.self_attn.qkv_proj.weights_loading_config.weight_mode, + "qkv_weight_shape": tuple(layer.self_attn.qkv_proj.weight.shape), + "gate_up_weight_mode": layer.mlp.gate_up_proj.weights_loading_config.weight_mode, + "gate_up_weight_shape": tuple(layer.mlp.gate_up_proj.weight.shape), + "qkv_bias": layer.self_attn.qkv_proj.bias is not None, + "q_norm_weight_shape": tuple(layer.self_attn.q_norm.weight.shape), + "k_norm_weight_shape": tuple(layer.self_attn.k_norm.weight.shape), + "fuse_qk_norm_rope": layer.self_attn.fuse_qk_norm_rope, + "rope_fusion": layer.self_attn.rope_fusion, + "rotary_embedding_present": layer.self_attn.rotary_emb is not None, + "tied_lm_head": model.lm_head.weight is model.model.embed_tokens.weight, + } + + +def _dense_qwen_input_embeddings(model: nn.Module) -> torch.Tensor: input_ids = torch.tensor( [0, 1, 2], dtype=torch.long, @@ -319,11 +392,11 @@ def _qwen2_input_embeddings(model: nn.Module) -> torch.Tensor: return model.model.embed_tokens(input_ids) -def _qwen2_embedding_logits(model: nn.Module) -> torch.Tensor: - return model.lm_head(_qwen2_input_embeddings(model)) +def _dense_qwen_embedding_logits(model: nn.Module) -> torch.Tensor: + return model.lm_head(_dense_qwen_input_embeddings(model)) -def _qwen2_hidden_states(model: nn.Module) -> torch.Tensor: +def _dense_qwen_hidden_states(model: nn.Module) -> torch.Tensor: qkv_weight = model.model.layers[0].self_attn.qkv_proj.weight values = torch.arange( 3 * model.config.hidden_size, @@ -333,12 +406,12 @@ def _qwen2_hidden_states(model: nn.Module) -> torch.Tensor: return (values % 17).to(qkv_weight.dtype) / 17 -def _qwen2_fused_qkv_output(model: nn.Module) -> torch.Tensor: - return model.model.layers[0].self_attn.qkv_proj(_qwen2_hidden_states(model)) +def _dense_qwen_fused_qkv_output(model: nn.Module) -> torch.Tensor: + return model.model.layers[0].self_attn.qkv_proj(_dense_qwen_hidden_states(model)) -def _qwen2_fused_gate_up_output(model: nn.Module) -> torch.Tensor: - return model.model.layers[0].mlp.gate_up_proj(_qwen2_hidden_states(model)) +def _dense_qwen_fused_gate_up_output(model: nn.Module) -> torch.Tensor: + return model.model.layers[0].mlp.gate_up_proj(_dense_qwen_hidden_states(model)) def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransformProfileRegistry: @@ -465,7 +538,8 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: assert constraints.multi_node == frozenset({False}) assert constraints.tied_word_embeddings == frozenset({False}) assert constraints.rope_types == frozenset({"default"}) - assert constraints.rope_fusion == frozenset({True}) + qwen3_profile = profile.model_type == "qwen3" + assert constraints.rope_fusion == frozenset({not qwen3_profile}) assert constraints.moe_backends is None assert constraints.moe_tp_sizes is None assert constraints.moe_ep_sizes is None @@ -482,9 +556,10 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: attention_backends = _format_documented_values(constraints.attention_backends) tp_sizes = _format_documented_values(constraints.tp_sizes) pp_cp_sizes = _format_documented_values(constraints.pp_sizes) + rope_description = "default fused QK-norm/RoPE" if qwen3_profile else "default fused RoPE" return ( f"Single-node dense {dtypes}, unquantized weights and KV cache, " - f"{attention_backends} attention, default fused RoPE, untied embeddings, " + f"{attention_backends} attention, {rope_description}, untied embeddings, " f"TP={tp_sizes}, PP/CP={pp_cp_sizes}, no LoRA, sparse attention, " "attention DP, speculative mode, or separately loaded draft model" ) @@ -766,9 +841,9 @@ def test_qwen2_dense_profile_qualifies_full_staged_lifecycle() -> None: ), state_probes=(("layout", _qwen2_layout_state),), output_probes=( - ("embedding-logits", _qwen2_embedding_logits), - ("fused-qkv", _qwen2_fused_qkv_output), - ("fused-gate-up", _qwen2_fused_gate_up_output), + ("embedding-logits", _dense_qwen_embedding_logits), + ("fused-qkv", _dense_qwen_fused_qkv_output), + ("fused-gate-up", _dense_qwen_fused_gate_up_output), ), ) @@ -830,8 +905,8 @@ def test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle( ), state_probes=(("layout", _qwen2_layout_state),), output_probes=( - ("fused-qkv", _qwen2_fused_qkv_output), - ("fused-gate-up", _qwen2_fused_gate_up_output), + ("fused-qkv", _dense_qwen_fused_qkv_output), + ("fused-gate-up", _dense_qwen_fused_gate_up_output), ), ) @@ -848,6 +923,96 @@ def test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle( assert _qwen2_layout_state(producer)["gate_up_weight_shape"] == (32, 16) +def test_qwen3_dense_profile_qualifies_full_staged_lifecycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = PostTransformQualificationCase( + profile_id="qwen3-for-causal-lm-bf16-target-v1", + model_factory=lambda: _tiny_qwen3_model(monkeypatch), + unqualified_model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + model_class=_UnqualifiedQwen3ForCausalLM, + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=(("layout", _qwen3_layout_state),), + output_probes=( + ("embedding-logits", _dense_qwen_embedding_logits), + ("fused-qkv", _dense_qwen_fused_qkv_output), + ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ), + ) + + producer, _receiver = assert_post_transform_lifecycle_equivalent(case) + + assert PostTransformRuntimeConfig.from_model_config( + producer.model_config, model=producer + ) == _bf16_dense_runtime_config(rope_fusion=False) + assert _qwen3_layout_state(producer) == { + "attention_type": modeling_qwen3_mod.Qwen3Attention.__name__, + "qkv_weight_mode": WeightMode.FUSED_QKV_LINEAR, + "qkv_weight_shape": (32, 16), + "gate_up_weight_mode": WeightMode.FUSED_GATE_UP_LINEAR, + "gate_up_weight_shape": (64, 16), + "qkv_bias": False, + "q_norm_weight_shape": (4,), + "k_norm_weight_shape": (4,), + "fuse_qk_norm_rope": True, + "rope_fusion": False, + "rotary_embedding_present": True, + "tied_lm_head": False, + } + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( + monkeypatch: pytest.MonkeyPatch, + rank: int, +) -> None: + monkeypatch.setattr(mapping_mod, "mpi_disabled", lambda: False) + monkeypatch.setattr(distributed_mod, "AllReduce", _AllReduceStub) + case = PostTransformQualificationCase( + profile_id="qwen3-for-causal-lm-bf16-target-v1", + model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + tp_size=2, + rank=rank, + ), + unqualified_model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + model_class=_UnqualifiedQwen3ForCausalLM, + tp_size=2, + rank=rank, + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=(("layout", _qwen3_layout_state),), + output_probes=( + ("fused-qkv", _dense_qwen_fused_qkv_output), + ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ), + ) + + producer, _receiver = assert_post_transform_lifecycle_equivalent(case) + + assert PostTransformRuntimeConfig.from_model_config( + producer.model_config, model=producer + ) == _bf16_dense_runtime_config( + tp_size=2, + moe_tp_size=2, + attention_tp_size=2, + rope_fusion=False, + ) + assert _qwen3_layout_state(producer)["qkv_weight_shape"] == (16, 16) + assert _qwen3_layout_state(producer)["gate_up_weight_shape"] == (32, 16) + + @pytest.mark.cpu_only @pytest.mark.parametrize( "overrides, expected_dimensions", @@ -909,24 +1074,32 @@ def test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle( id="tied-embeddings", ), pytest.param({"rope_type": "yarn"}, {"rope_type"}, id="yarn"), - pytest.param({"rope_fusion": False}, {"rope_fusion"}, id="unfused-rope"), ], ) @pytest.mark.parametrize( - "root_model_class, architecture, model_type", + "root_model_class, architecture, model_type, supported_rope_fusion", [ pytest.param( modeling_llama_mod.LlamaForCausalLM, "LlamaForCausalLM", "llama", + True, id="llama", ), pytest.param( modeling_qwen_mod.Qwen2ForCausalLM, "Qwen2ForCausalLM", "qwen2", + True, id="qwen2", ), + pytest.param( + modeling_qwen3_mod.Qwen3ForCausalLM, + "Qwen3ForCausalLM", + "qwen3", + False, + id="qwen3", + ), ], ) def test_bf16_dense_profiles_reject_unqualified_runtime_variants( @@ -935,6 +1108,7 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( root_model_class: type[nn.Module], architecture: str, model_type: str, + supported_rope_fusion: bool, ) -> None: decision = ModelLoader._post_transform_profile_registry().qualify( root_model_class=root_model_class, @@ -943,7 +1117,10 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( speculative_mode=None, protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_config=_bf16_dense_runtime_config(**overrides), + runtime_config=_bf16_dense_runtime_config( + rope_fusion=supported_rope_fusion, + **overrides, + ), ) assert not decision.qualified @@ -951,6 +1128,56 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( assert decision.unsupported_runtime_dimensions == frozenset(expected_dimensions) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "root_model_class, architecture, model_type, unsupported_rope_fusion", + [ + pytest.param( + modeling_llama_mod.LlamaForCausalLM, + "LlamaForCausalLM", + "llama", + False, + id="llama", + ), + pytest.param( + modeling_qwen_mod.Qwen2ForCausalLM, + "Qwen2ForCausalLM", + "qwen2", + False, + id="qwen2", + ), + pytest.param( + modeling_qwen3_mod.Qwen3ForCausalLM, + "Qwen3ForCausalLM", + "qwen3", + True, + id="qwen3", + ), + ], +) +def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( + root_model_class: type[nn.Module], + architecture: str, + model_type: str, + unsupported_rope_fusion: bool, +) -> None: + decision = ModelLoader._post_transform_profile_registry().qualify( + root_model_class=root_model_class, + architecture=architecture, + model_type=model_type, + speculative_mode=None, + protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_config=_bf16_dense_runtime_config( + rope_fusion=unsupported_rope_fusion, + ), + ) + + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert decision.unsupported_runtime_dimensions == frozenset({"rope_fusion"}) + + @pytest.mark.cpu_only @pytest.mark.parametrize( "overrides", @@ -968,20 +1195,29 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( ], ) @pytest.mark.parametrize( - "root_model_class, architecture, model_type", + "root_model_class, architecture, model_type, supported_rope_fusion", [ pytest.param( modeling_llama_mod.LlamaForCausalLM, "LlamaForCausalLM", "llama", + True, id="llama", ), pytest.param( modeling_qwen_mod.Qwen2ForCausalLM, "Qwen2ForCausalLM", "qwen2", + True, id="qwen2", ), + pytest.param( + modeling_qwen3_mod.Qwen3ForCausalLM, + "Qwen3ForCausalLM", + "qwen3", + False, + id="qwen3", + ), ], ) def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( @@ -989,6 +1225,7 @@ def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( root_model_class: type[nn.Module], architecture: str, model_type: str, + supported_rope_fusion: bool, ) -> None: decision = ModelLoader._post_transform_profile_registry().qualify( root_model_class=root_model_class, @@ -997,7 +1234,10 @@ def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( speculative_mode=None, protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_config=_bf16_dense_runtime_config(**overrides), + runtime_config=_bf16_dense_runtime_config( + rope_fusion=supported_rope_fusion, + **overrides, + ), ) assert decision.qualified diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index 2dfb670e92be..a3c6a7cc7e7a 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -36,6 +36,7 @@ from tensorrt_llm._torch.weight_sharing import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, IdentityCheckPolicy, SourceIdentity, SourceIdentityMismatchError, @@ -77,6 +78,7 @@ def test_from_model_config_requires_one_artifact_source() -> None: [ LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ], ) def test_from_model_config_binds_transform_abi(transform_abi_id: str) -> None: From 559a8eb5fa3d368d7035e50f8075c28b1e522c2c Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Fri, 21 Aug 2026 11:50:50 -0700 Subject: [PATCH 2/4] Address review comments Signed-off-by: Michal Guzek --- docs/source/features/model-express.md | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index aa06b826bced..34c02d2ca8f5 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -60,12 +60,16 @@ submodel. The Llama, Qwen2, and Qwen3 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 or unfused RoPE, 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. 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. +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 realized `rope_fusion=False`. 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 From 50ba9f1d03ce22464f89b5b13b6af26fbcbd7bd3 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Fri, 21 Aug 2026 12:01:28 -0700 Subject: [PATCH 3/4] Address review comments Signed-off-by: Michal Guzek --- docs/source/features/model-express.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 34c02d2ca8f5..07d472744866 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -66,10 +66,10 @@ 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 realized `rope_fusion=False`. 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. +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. ### Adding a Model Family From 406790339034f20a464e484ad55835a8ce6b6c70 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 24 Aug 2026 13:15:21 -0700 Subject: [PATCH 4/4] Address review comments Signed-off-by: Michal Guzek --- tensorrt_llm/_torch/weight_sharing/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index 1393258294de..d8d11369a151 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -44,11 +44,14 @@ __all__ = [ "ARTIFACT_IDENTITY_FORMAT_VERSION", - "ArtifactIdentity", "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1", "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "SOURCE_IDENTITY_FORMAT_VERSION", + "ArtifactIdentity", + "IdentityCheckDecision", + "IdentityCheckPolicy", + "IdentityMatchResult", "PostTransformConfigIdentity", "PostTransformFeature", "PostTransformProfile", @@ -59,9 +62,6 @@ "PostTransformRuntimeConstraints", "PostTransformTransferScope", "SourceIdentity", - "IdentityMatchResult", - "IdentityCheckPolicy", - "IdentityCheckDecision", "SourceIdentityMismatchError", "check_weight_sharing_compatibility", ]