[None][chore] Integration tests for MoE lora & bugfixes - #15271
Conversation
brb-nv
changed the title
5889aa7 to
e043499
Compare
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR implements rank-0 handling and CUDA graph cache management for routed-expert MoE LoRA support. CUTLASS kernels conditionally skip output GEMM operations for rank-0 rows; tensor weight pointers are reshaped to correct layout; and CUDA graph caches are refreshed after pointer updates to avoid stale references. Integration tests validate the feature across multiple execution modes. ChangesRouted-Expert MoE LoRA Implementation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 1
🧹 Nitpick comments (1)
tests/unittest/llmapi/test_llm_pytorch.py (1)
1099-1152: ⚡ Quick winAnnotate the new helper and test entrypoint.
Both new functions are introduced untyped, which makes the adapter-shape contract here harder to follow and diverges from the repo’s Python typing rule. Please add parameter and return annotations (
-> Nonefor the test) while this helper is still small.As per coding guidelines, Python functions should always be annotated with explicit return types.
Also applies to: 1162-1163
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_llm_pytorch.py` around lines 1099 - 1152, The helper _write_routed_expert_lora_adapter and the new test entrypoint must have explicit type annotations; update _write_routed_expert_lora_adapter signature to something like (save_dir: str, *, moe_layers: Sequence[int], num_experts: int, hidden_size: int, moe_intermediate_size: int, rank: int, lora_alpha: float, seed: int) -> None and add -> None to the test entrypoint function that calls it (the test function immediately following the helper), and if helpful annotate local helper types (e.g. randn: Callable[[int,int,float], torch.Tensor]) — no behavioral changes, just add parameter and return type annotations to satisfy the repo typing rule.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/llmapi/test_llm_pytorch.py`:
- Around line 1155-1285: The test currently only checks that any adapter changed
output and never mixes a no-LoRA request nor exercises CUDA-graph replay; update
test_qwen_moe_routed_expert_multi_lora_varying_ranks to include at least one
None/no-LoRA request in lora_requests and assert that the corresponding output
token_ids equal base_tokens (verifies rank-0/skip path), change the any_differs
check to require that each non-None LoRA request produces tokens different from
base_tokens (ensuring all adapters applied), and when moe_lora_mode ==
"device_path_cudagraph" invoke llm.generate a second time with a different
adapter ordering/subset and assert the outputs reflect the new slot assignment
(i.e., outputs for the same prompts change according to the new ordering) to
exercise slot-table refresh on CUDA-graph replay.
---
Nitpick comments:
In `@tests/unittest/llmapi/test_llm_pytorch.py`:
- Around line 1099-1152: The helper _write_routed_expert_lora_adapter and the
new test entrypoint must have explicit type annotations; update
_write_routed_expert_lora_adapter signature to something like (save_dir: str, *,
moe_layers: Sequence[int], num_experts: int, hidden_size: int,
moe_intermediate_size: int, rank: int, lora_alpha: float, seed: int) -> None and
add -> None to the test entrypoint function that calls it (the test function
immediately following the helper), and if helpful annotate local helper types
(e.g. randn: Callable[[int,int,float], torch.Tensor]) — no behavioral changes,
just add parameter and return type annotations to satisfy the repo typing rule.
🪄 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: Enterprise
Run ID: 2bd517f4-3bd0-4069-986d-e8bf55e1ff8a
📒 Files selected for processing (4)
cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cutensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.pytensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.pytests/unittest/llmapi/test_llm_pytorch.py
|
PR_Github #53653 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #53675 [ run ] triggered by Bot. Commit: |
|
PR_Github #53653 [ run ] completed with state |
|
PR_Github #53675 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53758 [ run ] triggered by Bot. Commit: |
|
PR_Github #53758 [ run ] completed with state
|
f80ccb9 to
88e868b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53919 [ run ] triggered by Bot. Commit: |
88e868b to
78a80fb
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53939 [ run ] triggered by Bot. Commit: |
78a80fb to
f651594
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53984 [ run ] triggered by Bot. Commit: |
|
PR_Github #53939 [ run ] completed with state
|
|
PR_Github #53984 [ run ] completed with state
|
f651594 to
11c6031
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54154 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #54157 [ run ] triggered by Bot. Commit: |
|
PR_Github #54154 [ run ] completed with state |
11c6031 to
d69f8ad
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54160 [ run ] triggered by Bot. Commit: |
|
PR_Github #54157 [ run ] completed with state |
|
PR_Github #54160 [ run ] completed with state
|
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
d69f8ad to
d398db3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54343 [ run ] triggered by Bot. Commit: |
|
PR_Github #54343 [ run ] completed with state |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
4 participants
Description
This MR adds integration tests and 3 bugfixes for MoE lora. Details about fixes:
weight_pointersfrom flat [num_seqs*3] to [num_seqs, 3] so the MoE op gets the table shape it expects.To be merged after #14881.
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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.Summary by CodeRabbit
Bug Fixes
Tests