[None][feat] Afmoe trinity support - #13148
Conversation
📝 WalkthroughWalkthroughIntroduces a new AFMoE (Arcee Foundation MoE) model implementation for TensorRT-LLM, including a model architecture with specialized attention and MoE layers, a Hugging Face weight mapper for model weight conversion, public API exports, and comprehensive unit tests. Configuration updates exclude the test file from ruff formatting. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant AfmoeModel
participant AfmoeDecoderLayer
participant AfmoeAttention
participant AfmoeMoE
participant AfmoeGate
participant Experts
participant SharedMLP
Client->>AfmoeModel: forward(input_ids, position_ids, attn_metadata)
AfmoeModel->>AfmoeModel: Embed tokens + optional muP scaling
loop Each DecoderLayer
AfmoeModel->>AfmoeDecoderLayer: forward(hidden_states, position_ids, attn_metadata)
AfmoeDecoderLayer->>AfmoeAttention: forward(hidden_states, position_ids, attn_metadata)
AfmoeAttention->>AfmoeAttention: Compute Q/K with RMSNorm
AfmoeAttention->>AfmoeAttention: Attention (sliding-window or full per layer_type)
AfmoeAttention->>AfmoeAttention: Gate projection + sigmoid gate
AfmoeAttention->>AfmoeAttention: Output * gate
AfmoeAttention-->>AfmoeDecoderLayer: Gated attention output
alt Dense Layer (idx < num_dense_layers)
AfmoeDecoderLayer->>AfmoeDecoderLayer: GatedMLP forward
else MoE Layer
AfmoeDecoderLayer->>AfmoeMoE: forward(hidden_states, attn_metadata)
AfmoeMoE->>AfmoeGate: Gate forward (TensorRT dsv3_router_gemm_op)
AfmoeGate-->>AfmoeMoE: Router logits + top-k routing
AfmoeMoE->>Experts: Route & load balance
Experts-->>AfmoeMoE: Expert outputs
AfmoeMoE->>SharedMLP: Shared gated MLP forward (optional)
SharedMLP-->>AfmoeMoE: Shared output
AfmoeMoE->>AfmoeMoE: Combine routed + shared outputs
AfmoeMoE->>AfmoeMoE: Optional AllReduce (TP sync)
AfmoeMoE-->>AfmoeDecoderLayer: MoE output
end
AfmoeDecoderLayer->>AfmoeDecoderLayer: Residual connections + RMSNorms
AfmoeDecoderLayer-->>AfmoeModel: Layer output
end
AfmoeModel->>AfmoeModel: Final RMSNorm
AfmoeModel-->>Client: Logits
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_afmoe.py (1)
16-16: Prefer built-in types over typing module equivalents.Per coding guidelines for Python 3.10+, use built-in
dictandlistinstead oftyping.Dictandtyping.List.♻️ Proposed fix
-from typing import Dict, List, Optional +from typing import OptionalThen update the usages:
- Line 115:
List[Dict]→list[dict]- Line 506:
Dict→dictBased on learnings: "In TensorRT-LLM (Python requires >=3.10), you can use Python 3.10+ features (e.g., PEP 585 generics like dict[str, int], list[str], etc.) throughout the codebase."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_afmoe.py` at line 16, Replace typing module generics with built-in generics: in the import line that currently reads "from typing import Dict, List, Optional" remove Dict and List and keep only what's needed (e.g., Optional) and update all usages — change the annotation "List[Dict]" (used around the symbol at line ~115) to "list[dict]" and change standalone "Dict" annotations (around line ~506) to "dict"; ensure imports reflect the removed typing names and run type checks.
🤖 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/models/checkpoints/hf/afmoe_weight_mapper.py`:
- Around line 1-6: This file is missing the required NVIDIA copyright header;
add the standard NVIDIA Python source file header with the current year at the
top of the file before any imports. Ensure the header appears above the existing
imports (before "from torch import nn") and is included in every new Python
module such as this one that defines/uses HfWeightMapper, register_mapper, and
MoE so it complies with the project's licensing/coding guidelines.
In `@tests/unittest/_torch/modeling/test_modeling_afmoe.py`:
- Around line 1-7: Add the required NVIDIA copyright header at the very top of
the test file (above the existing imports such as "import unittest", "import
torch", and "import tensorrt_llm") to comply with repository coding guidelines;
ensure the header text matches the project's standard copyright block and is
placed before any module-level code or docstrings so it applies to the whole
file.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_afmoe.py`:
- Line 16: Replace typing module generics with built-in generics: in the import
line that currently reads "from typing import Dict, List, Optional" remove Dict
and List and keep only what's needed (e.g., Optional) and update all usages —
change the annotation "List[Dict]" (used around the symbol at line ~115) to
"list[dict]" and change standalone "Dict" annotations (around line ~506) to
"dict"; ensure imports reflect the removed typing names and run type checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 5e3ed3d6-51a9-439c-ad84-b7eba71580c1
📒 Files selected for processing (6)
pyproject.tomltensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/checkpoints/__init__.pytensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.pytensorrt_llm/_torch/models/modeling_afmoe.pytests/unittest/_torch/modeling/test_modeling_afmoe.py
svc-trtllm-gh-bot
added
the
Community want to contribute
723da98 to
7ae17a7
Compare
7ae17a7 to
bd342b6
Compare
|
This PR is really huge. Is it possible to break into smaller pieces? |
|
@alyosha-swamy If I understand correctly, Add AFMoE support is the only functional changes if the PR. If you agree, please clean up the PR commit history. |
bd342b6 to
812614a
Compare
This is cleaned now, is there anything else we can do to get this merged? We have a time-sensitive task blocked on it |
|
@
Hi, pinging again for a status update here pleease :) |
|
SimengLiu-nv
left a comment
SimengLiu-nv
left a comment
There was a problem hiding this comment.
Additional Concerns:
- No coverage of end-2end model tests.
- AFMoE attention-DP path shape-mismatches because gate/MLP layers stay TP-sharded:
tensorrt_llm/_torch/models/modeling_afmoe.py:283
AfmoeAttention.gate_proj is constructed with the original model_config.mapping, while the base Attention constructor internally switches the attention projections to effective tp_size=1 when mapping.enable_attention_dp=True. That means in attention-DP mode the attention output is full-width, but gate_proj(hidden_states) is still column-sharded. The multiply at line 343 can then fail with a shape mismatch.
The same attention-DP issue applies to dense/shared GatedMLP construction in this file, which also uses the original TP mapping instead of the overridden_tp_size=1 pattern used by nearby models. Please either explicitly reject attention-DP for AFMoE or wire the AFMoE gate/MLP modules with the ADP-safe mapping behavior used by existing model implementations.
|
/bot run --disable-fail-fast |
|
PR_Github #50827 [ run ] triggered by Bot. Commit: |
1b42f5e to
fba10c3
Compare
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
fba10c3 to
c425f61
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #51263 [ run ] triggered by Bot. Commit: |
|
PR_Github #51263 [ run ] completed with state
|
|
Multiple tests are failing and I can reproduce them locally. Kindly fix them, confirm tests pass locally and post to retrigger CI again. Branch I used to test: https://github.com/brb-nv/TensorRT-LLM/tree/user/brb/shepherd-afmoe-branch |
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
|
/bot run --disable-fail-fast |
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
f1cd731 to
72fe838
Compare
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
|
/bot run --disable-fail-fast |
|
PR_Github #52132 [ run ] triggered by Bot. Commit: |
|
PR_Github #52132 [ run ] completed with state
|
Can't seem to identify this problem, all tests pass for me locally |
|
/bot run --disable-fail-fast |
|
PR_Github #52203 [ run ] triggered by Bot. Commit: |
|
PR_Github #52203 [ run ] completed with state
|
|
Hi, I see the following issue with Afmoe test in CI: |
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
e309efd to
cb2645a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52416 [ run ] triggered by Bot. Commit: |
|
PR_Github #52416 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52457 [ run ] triggered by Bot. Commit: |
|
PR_Github #52457 [ run ] completed with state |
Signed-off-by: Alyosha-Swamy <raghav@arcee.ai>
7 participants
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.