Repository navigation
[TRTLLM-11544][feat] Add Qwen 3.5 supporting(NVFP4). - #12302
Conversation
e7f4d58 to
71b1cab
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
71b1cab to
e5a0605
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #39541 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR adds comprehensive support for Qwen3.5 model architectures (both dense and sparse MoE variants) to TensorRT-LLM's PyTorch backend, including configuration classes, model implementations, weight mappers, and infrastructure updates for execution and caching. Changes
Sequence DiagramsequenceDiagram
participant User
participant Config as Config System
participant ModelBuilder as Model Builder
participant WeightMapper as Weight Mapper
participant Executor as Executor/KV Cache
participant GPU as GPU Execution
User->>Config: Load Qwen3.5 config
Config->>Config: Resolve model_type via _CONFIG_REGISTRY
Config->>Config: Normalize rope_scaling & layer_types
Config-->>ModelBuilder: Return Qwen3_5Config/Qwen3_5MoeConfig
ModelBuilder->>ModelBuilder: Instantiate model from config
ModelBuilder->>ModelBuilder: Build embeddings, decoder layers<br/>(GatedDeltaNet or Attention + MoE)
ModelBuilder-->>WeightMapper: Return model structure
User->>WeightMapper: Load pretrained weights
WeightMapper->>WeightMapper: Strip multimodal wrapper prefix
WeightMapper->>WeightMapper: Filter vision/MTP tensors
WeightMapper->>WeightMapper: Reorder QKV for TP/linear attention
WeightMapper->>WeightMapper: Transpose expert tensors (MoE)
WeightMapper-->>ModelBuilder: Inject processed weights
User->>Executor: Create executor with Qwen3.5 model
Executor->>Executor: Select MambaHybridCacheManager<br/>(is_qwen3_5_hybrid detected)
Executor->>Executor: Configure mRoPE position tracking
Executor-->>GPU: Ready for inference
User->>GPU: Execute requests
GPU->>GPU: Process with layer_types routing<br/>(full_attention or linear_attention)
GPU->>GPU: Apply hybrid attention/MoE blocks
GPU-->>User: Return generated tokens
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment Tip You can enable review details to help with troubleshooting, context usage and more.Enable the |
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (5)
tensorrt_llm/_torch/configs/__init__.py (1)
9-15: Consider sorting__all__alphabetically.Static analysis (RUF022) flags that
__all__is not sorted. Keeping it sorted improves readability and reduces merge conflicts.🔧 Sorted `__all__`
__all__ = [ "DeepseekV3Config", "Qwen3_5Config", - "Qwen3_5TextConfig", "Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig", + "Qwen3_5TextConfig", ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/configs/__init__.py` around lines 9 - 15, The module-level __all__ list is not alphabetically sorted; please reorder the entries in __all__ so they appear in lexicographical order (e.g., DeepseekV3Config, Qwen3_5Config, Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5TextConfig) to satisfy static analysis (RUF022) and reduce merge conflicts—update the __all__ definition accordingly.tensorrt_llm/_torch/model_config.py (1)
784-798: Fix bracket indentation to satisfy linter.The closing bracket on line 790 doesn't align with the opening bracket's line per E123. Consider reformatting for consistency.
🔧 Suggested fix
elif hasattr( self.pretrained_config, "architectures" ) and self.pretrained_config.architectures is not None and self.pretrained_config.architectures[ - 0] in [ - "Qwen3NextForCausalLM", - "Qwen3_5ForCausalLM", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForCausalLM", - "Qwen3_5MoeForConditionalGeneration", - ]: + 0] in ["Qwen3NextForCausalLM", "Qwen3_5ForCausalLM", + "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration"]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/model_config.py` around lines 784 - 798, The list of model names used in the membership check (the literal starting at "in [") has its closing bracket mis-indented; adjust the closing bracket so it aligns with the opening bracket of the list to satisfy E123. Locate the membership check that compares against the list of strings (the code referencing "in [" and the model names "Qwen3NextForCausalLM", "Qwen3_5ForCausalLM", etc.) in the method using self.pretrained_config, and move the closing ] to align with the opening [ (or place it on the same line as the opening token) so the bracket indentation matches the linter expectation.tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2860-2866: Good defensive bounds checking.The bounds check prevents buffer overruns when
mrope_position_idshas fewer tokens than expected. Zero-filling the remainder is a reasonable fallback.Consider adding a debug/warning log when
actual_num_tokens < total_num_tokensto aid debugging, as this condition may indicate unexpected data flow from upstream multimodal processing.Optional: Add debug logging
actual_num_tokens = min(total_num_tokens, mrope_position_ids.shape[-1]) self.mrope_position_ids_cuda[:, :, :actual_num_tokens].copy_( mrope_position_ids[:, :, :actual_num_tokens], non_blocking=True) if actual_num_tokens < total_num_tokens: + logger.debug( + f"mrope_position_ids has {mrope_position_ids.shape[-1]} tokens " + f"but {total_num_tokens} expected; zero-filling remainder") self.mrope_position_ids_cuda[:, :, actual_num_tokens: total_num_tokens].fill_(0)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2860 - 2866, The code correctly bounds-checks and zero-fills when mrope_position_ids has fewer tokens (actual_num_tokens < total_num_tokens); add a debug or warning log inside that conditional to surface unexpected shorter inputs by calling something like self.logger.debug(...) or self.logger.warning(...) (include values: actual_num_tokens, total_num_tokens, and perhaps mrope_position_ids.shape) so it’s easy to trace upstream multimodal data flow; place the log in the branch that fills the remainder right before or after the fill_(0) call referencing mrope_position_ids, mrope_position_ids_cuda, actual_num_tokens, and total_num_tokens.tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)
5826-5829: Avoid mutable class-level dicts for evaluator kwargs.Lines 5826-5829 and 5891-5894 define shared mutable dicts (
RUF012). If any evaluation path mutates them, test state can bleed across cases.♻️ Proposed refactor
class TestQwen3_5_27B(LlmapiAccuracyTestHarness): @@ - GSM8K_EXTRA_EVALUATOR_KWARGS = dict( - apply_chat_template=True, - chat_template_kwargs=dict(enable_thinking=False), - ) + `@staticmethod` + def _gsm8k_extra_evaluator_kwargs(): + return dict( + apply_chat_template=True, + chat_template_kwargs=dict(enable_thinking=False), + ) @@ - extra_evaluator_kwargs=self.GSM8K_EXTRA_EVALUATOR_KWARGS, + extra_evaluator_kwargs=self._gsm8k_extra_evaluator_kwargs(), class TestQwen3_5_122B_A10B(LlmapiAccuracyTestHarness): @@ - GSM8K_EXTRA_EVALUATOR_KWARGS = dict( - apply_chat_template=True, - chat_template_kwargs=dict(enable_thinking=False), - ) + `@staticmethod` + def _gsm8k_extra_evaluator_kwargs(): + return dict( + apply_chat_template=True, + chat_template_kwargs=dict(enable_thinking=False), + ) @@ - extra_evaluator_kwargs=self.GSM8K_EXTRA_EVALUATOR_KWARGS, + extra_evaluator_kwargs=self._gsm8k_extra_evaluator_kwargs(),Also applies to: 5891-5894
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 5826 - 5829, GSM8K_EXTRA_EVALUATOR_KWARGS (and the similar dict at the other location) are class-level mutable dicts that can be mutated across tests; change them to factory functions or always return a fresh dict to avoid shared state (e.g., replace the constant dict with a get_gsm8k_extra_evaluator_kwargs() function that returns a new dict each call, or ensure callers copy() the dict before use). Update references to use the new factory (or copy) so each test gets its own independent kwargs, and keep the keys apply_chat_template and chat_template_kwargs(enable_thinking=False) intact.tensorrt_llm/_torch/configs/qwen3_5.py (1)
168-181: Surfacelayer_typesandfull_attention_intervalon the composite config wrapper.The
_Qwen3_5BaseCompositeConfig.__init__only mirrors scalar fields fromself.text_configbut omitslayer_typesandfull_attention_interval. Callers (e.g.,model_config.pyline 795) access these directly from the top-level config to detect the hybrid schedule; without proxying them, the configuration will not advertise its layer composition.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/configs/qwen3_5.py` around lines 168 - 181, The composite config __init__ currently proxies scalar text config fields but omits layer_types and full_attention_interval, so expose them by copying from self.text_config when present: in _Qwen3_5BaseCompositeConfig.__init__, after setting vocab_size/hidden_size/etc., add logic to set self.layer_types = self.text_config.layer_types (or None if missing) and self.full_attention_interval = self.text_config.full_attention_interval (or None if missing) so callers that inspect the top-level config (e.g., code that checks for a hybrid schedule) can see the layer composition and attention schedule.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/configs/qwen3_5.py`:
- Line 1: Update the SPDX header year in the new file to 2026: edit the
top-of-file copyright line in tensorrt_llm/_torch/configs/qwen3_5.py (the SPDX
header comment) and change the year from 2025 to 2026 so it matches other new
Qwen3.5 sources and complies with the repository copyright header guideline.
- Around line 84-92: Validate the hybrid layer schedule early: if
full_attention_interval is used to auto-generate layer_types, ensure
full_attention_interval is an int > 0 and raise a ValueError with a clear
message if not; if layer_types is provided, validate its length equals
num_hidden_layers and that every entry is a valid member of
_ALL_DECODER_LAYER_TYPES (raise ValueError listing the invalid entries); after
generation also assert len(layer_types) == num_hidden_layers and all entries in
_ALL_DECODER_LAYER_TYPES before assigning self.layer_types so later indexing via
pretrained_config.layer_types[layer_idx] cannot raise
ZeroDivisionError/IndexError/KeyError.
In `@tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_moe_weight_mapper.py`:
- Around line 27-32: The docstring for the Qwen3_5MoeWeightMapper class
incorrectly names a base class; inspect the actual parent classes of
Qwen3_5MoeWeightMapper and update the docstring so the sentence "MoE-specific
weight loading ... is handled by the base X" names the correct base class
(replace the incorrect Qwen2MoeHfWeightMapper with the actual parent that
provides MoE handling, e.g., Qwen3_5HfWeightMapper or the appropriate MoE base),
ensuring the description matches the class inheritance.
In `@tensorrt_llm/_torch/models/modeling_qwen3_5_moe.py`:
- Around line 333-337: The local variable mamba_metadata is left pointing at the
old object after replacing attn_metadata.mamba_metadata, causing subsequent code
(decoder layers) to use stale buffers; after creating and assigning the new
Mamba2Metadata instance (attn_metadata.mamba_metadata = Mamba2Metadata(...))
update the local mamba_metadata variable to reference the new object (e.g.,
mamba_metadata = attn_metadata.mamba_metadata) so the rest of the forward uses
the resized metadata and prepared buffers. Ensure this change is applied in the
same block where Mamba2Metadata is constructed so decoder paths consume the
fresh metadata.
In `@tensorrt_llm/_torch/models/modeling_qwen3_5.py`:
- Around line 722-726: The code creates a new Mamba2Metadata and assigns it to
attn_metadata.mamba_metadata but leaves the local variable mamba_metadata
pointing to the old object and never initializes the new metadata; update the
local mamba_metadata variable to the newly created Mamba2Metadata
(mamba_metadata = attn_metadata.mamba_metadata) and call the metadata's
initialization/preparation routine (the same method used elsewhere to populate
query_start_loc and state_indices) before the decoder layers run so each layer
receives the resized and prepared metadata.
- Around line 419-422: The code unconditionally zeros warm-started SSM states in
the prefill branch; change the logic in the prefill path so
ssm_states[state_indices_p] is only zeroed for requests that do NOT have cached
initial states (use the existing has_initial_states flag per request).
Concretely, compute a boolean mask for the prefill indices (e.g., mask =
~has_initial_states[state_indices_p] or equivalent using cache_indices) and only
assign zeros to ssm_states[state_indices_p][mask], preserving entries where
has_initial_states is true so chunk_gated_delta_rule() can read cached
ssm_states[cache_indices] correctly. Ensure dtype/device match and that the
masked assignment handles empty masks safely.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2610-2614: The loop that builds mrope_position_ids is appending
the same multimodal_params object beam_width times which likely duplicates
entries; in the block that iterates "for beam in range(beam_width)" (where
gen_mrope_position_ids is appended to mrope_position_ids), move the
multimodal_params_list.append(multimodal_params) out of the loop so it is
executed only once per request (matching the non-mrope and context handling
paths), and rename the unused loop variable beam to _ to avoid the
unused-variable warning; if duplication was intentional, add a clarifying
comment instead.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 5822-5823: Update the test harness so TestQwen3_5_27B uses a
distinct MODEL_NAME value ("Qwen/Qwen3.5-27B") and MODEL_PATH accordingly (the
constants in the test file: MODEL_NAME and MODEL_PATH in TestQwen3_5_27B), and
then add the corresponding MMLU baseline entry for "Qwen/Qwen3.5-27B" into
references/mmlu.yaml with the correct accuracy value for the 27B model so MMLU
lookups succeed; ensure GSM8K and MMLU reference keys match the new MODEL_NAME.
---
Nitpick comments:
In `@tensorrt_llm/_torch/configs/__init__.py`:
- Around line 9-15: The module-level __all__ list is not alphabetically sorted;
please reorder the entries in __all__ so they appear in lexicographical order
(e.g., DeepseekV3Config, Qwen3_5Config, Qwen3_5MoeConfig, Qwen3_5MoeTextConfig,
Qwen3_5TextConfig) to satisfy static analysis (RUF022) and reduce merge
conflicts—update the __all__ definition accordingly.
In `@tensorrt_llm/_torch/configs/qwen3_5.py`:
- Around line 168-181: The composite config __init__ currently proxies scalar
text config fields but omits layer_types and full_attention_interval, so expose
them by copying from self.text_config when present: in
_Qwen3_5BaseCompositeConfig.__init__, after setting vocab_size/hidden_size/etc.,
add logic to set self.layer_types = self.text_config.layer_types (or None if
missing) and self.full_attention_interval =
self.text_config.full_attention_interval (or None if missing) so callers that
inspect the top-level config (e.g., code that checks for a hybrid schedule) can
see the layer composition and attention schedule.
In `@tensorrt_llm/_torch/model_config.py`:
- Around line 784-798: The list of model names used in the membership check (the
literal starting at "in [") has its closing bracket mis-indented; adjust the
closing bracket so it aligns with the opening bracket of the list to satisfy
E123. Locate the membership check that compares against the list of strings (the
code referencing "in [" and the model names "Qwen3NextForCausalLM",
"Qwen3_5ForCausalLM", etc.) in the method using self.pretrained_config, and move
the closing ] to align with the opening [ (or place it on the same line as the
opening token) so the bracket indentation matches the linter expectation.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2860-2866: The code correctly bounds-checks and zero-fills when
mrope_position_ids has fewer tokens (actual_num_tokens < total_num_tokens); add
a debug or warning log inside that conditional to surface unexpected shorter
inputs by calling something like self.logger.debug(...) or
self.logger.warning(...) (include values: actual_num_tokens, total_num_tokens,
and perhaps mrope_position_ids.shape) so it’s easy to trace upstream multimodal
data flow; place the log in the branch that fills the remainder right before or
after the fill_(0) call referencing mrope_position_ids, mrope_position_ids_cuda,
actual_num_tokens, and total_num_tokens.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 5826-5829: GSM8K_EXTRA_EVALUATOR_KWARGS (and the similar dict at
the other location) are class-level mutable dicts that can be mutated across
tests; change them to factory functions or always return a fresh dict to avoid
shared state (e.g., replace the constant dict with a
get_gsm8k_extra_evaluator_kwargs() function that returns a new dict each call,
or ensure callers copy() the dict before use). Update references to use the new
factory (or copy) so each test gets its own independent kwargs, and keep the
keys apply_chat_template and chat_template_kwargs(enable_thinking=False) intact.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4a709772-7fd4-41bc-93d4-db5aa48e730c
📒 Files selected for processing (20)
tensorrt_llm/_torch/configs/__init__.pytensorrt_llm/_torch/configs/qwen3_5.pytensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/checkpoints/__init__.pytensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_moe_weight_mapper.pytensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.pytensorrt_llm/_torch/models/modeling_qwen3_5.pytensorrt_llm/_torch/models/modeling_qwen3_5_moe.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/models/__init__.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/references/mmlu.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/test_model_config.py
👮 Files not reviewed due to content moderation or server errors (8)
- tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/integration/defs/accuracy/references/mmlu.yaml
- tensorrt_llm/models/init.py
- tensorrt_llm/_torch/models/init.py
- tests/unittest/_torch/test_model_config.py
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
- tensorrt_llm/_torch/models/checkpoints/init.py
|
PR_Github #39541 [ run ] completed with state
|
e5a0605 to
7a116f8
Compare
b20193c to
1dbf78a
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
1dbf78a to
a2df3f1
Compare
|
PR_Github #39796 [ run ] triggered by Bot. Commit: |
|
PR_Github #39796 [ run ] completed with state |
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #39797 [ run ] triggered by Bot. Commit: |
|
PR_Github #39797 [ run ] completed with state |
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #39825 [ run ] triggered by Bot. Commit: |
|
PR_Github #39825 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #39855 [ run ] triggered by Bot. Commit: |
|
PR_Github #39855 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #39902 [ run ] triggered by Bot. Commit: |
|
PR_Github #39902 [ run ] completed with state |
|
Does this support tensorrt backend? |
Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
|
@nv-guomingz Could you please share your launch commands? Thanks! |
Summary by CodeRabbit
New Features
Tests
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.