Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .claude/agents/perf-profiling-specialist.md
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,7 @@ FROM CUPTI_ACTIVITY_KIND_KERNEL;
| Attention implementations | `tensorrt_llm/_torch/attention/` |
| Custom ops | `tensorrt_llm/_torch/custom_ops/` |
| Triton kernels | `tensorrt_llm/_torch/triton_kernels/` |
| MoE routing | `tensorrt_llm/_torch/modules/fused_moe/` |
| MoE routing | `tensorrt_llm/_torch/moe/fused_moe/` |

## Structured Output

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,4 +34,4 @@ measured: []
- **Accuracy risk:** lossless (bf16 token exchange; path asserts `dtype==bfloat16` — changes transport, not math).
- **Verify:** end-to-end throughput/latency vs allgather/reducescatter baseline; correctness with `test_fused_moe.py`; sanity accuracy on a small eval.
- **Rollback:** unset `TRTLLM_CAN_USE_DEEP_EP` (falls back to allgather/reducescatter). Trigger: DeepEP build/IBGDA unavailable, regression vs baseline, or CUDA-Graph capture failure on the non-LL method.
- **Prior art:** PR #4792. Files: `_torch/modules/fused_moe/deep_ep_utils.py`, `fused_moe_cutlass.py` (`AlltoallMethodType`), `modeling_deepseekv3.py`, `docker/common/install_deep_ep.sh`. Owning skill: **trtllm-moe-develop**; sweep with **perf-sweep-workflow**.
- **Prior art:** PR #4792. Files: `_torch/moe/fused_moe/deep_ep_utils.py`, `fused_moe_cutlass.py` (`AlltoallMethodType`), `modeling_deepseekv3.py`, `docker/common/install_deep_ep.sh`. Owning skill: **trtllm-moe-develop**; sweep with **perf-sweep-workflow**.
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,4 @@ measured: []
- **Accuracy risk:** lossy — tokens cross the wire in FP4/FP8 (and combine may be low-precision); partial sums/activations lose precision vs bf16 transport. Needs accuracy record + rollback criterion; FP4 combine is the highest-risk leg.
- **Verify:** task-accuracy/perplexity parity vs the bf16-dispatch DeepEP baseline (and vs low-precision-combine off); confirm throughput gain; check FP4 pack/unpack round-trip (SF swizzle) correctness.
- **Rollback:** `TRTLLM_MOE_POST_QUANT_ALLTOALLV=0` and `TRTLLM_MOE_USE_LOW_PRECISION_COMBINE=0` (revert to bf16 dispatch/combine). Trigger: accuracy drop beyond recorded threshold.
- **Prior art:** PRs #5881, #6296, #7927. Files: `_torch/modules/fused_moe/fused_moe_wide_ep.py` (`alltoall_postquant_dispatch`, `low_latency_dispatch_fp4`, `low_latency_combine_low_precision`), `deep_ep_utils.py`, `thop/moeOp.cpp`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
- **Prior art:** PRs #5881, #6296, #7927. Files: `_torch/moe/fused_moe/fused_moe_wide_ep.py` (`alltoall_postquant_dispatch`, `low_latency_dispatch_fp4`, `low_latency_combine_low_precision`), `deep_ep_utils.py`, `thop/moeOp.cpp`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,4 @@ measured: []
- **Accuracy risk:** lossless when the flag is correct (kernel reads SF in its actual layout). Risk = layout/flag mismatch silently corrupting FP4 GEMM output — parity-check on first enable.
- **Verify:** profile shows the swizzle op gone from MoE path; MoE-layer latency/launch count down; NVFP4 MoE output parity vs pre-swizzle path.
- **Rollback:** set `swizzled_input_sf=True` and restore the `swizzle_sf(...)` calls. Trigger: FP4 MoE accuracy mismatch (layout/flag inconsistency).
- **Prior art:** PR #6231. Files: `_torch/custom_ops/torch_custom_ops.py`, `_torch/modules/fused_moe/fused_moe_wide_ep.py`, `cpp/.../thop/moeOp.cpp`, `cutlass_kernels/moe_gemm/moe_kernels.cu`. Owning specialist: **kernel-cuda-specialist**.
- **Prior art:** PR #6231. Files: `_torch/custom_ops/torch_custom_ops.py`, `_torch/moe/fused_moe/fused_moe_wide_ep.py`, `cpp/.../thop/moeOp.cpp`, `cutlass_kernels/moe_gemm/moe_kernels.cu`. Owning specialist: **kernel-cuda-specialist**.
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ log_markers:
- "use the original pytorch implementation" # is_fused=False fallback warning
eligibility:
- "model uses the noaux_tc router (sigmoid gating + per-expert correction bias, grouped top-k)"
- "fused-path bounds live in tensorrt_llm/_torch/modules/fused_moe/routing.py::Deepseekv3RoutingImpl.noaux_tc (the is_fused guard) — read them from YOUR checkout before judging eligibility"
- "fused-path bounds live in tensorrt_llm/_torch/moe/fused_moe/routing.py::Deepseekv3RoutingImpl.noaux_tc (the is_fused guard) — read them from YOUR checkout before judging eligibility"
- "as of 84926bcb6f14, n_group > 1: requires top_k <= 8 and num_experts <= 256 and experts_per_group <= 32 and experts_per_group * topk_group <= 256"
- "as of 84926bcb6f14, n_group == 1: requires num_experts <= 1024 and top_k <= 32"
interactions:
Expand Down Expand Up @@ -105,7 +105,7 @@ measured: []
must stay in sync with the kernel's supported bounds — a mismatch errors
instead of falling back.
- **Prior art:** PRs #7761 (`[TRTLLM-8637]`), #13186 (widens fused-path
eligibility; `tensorrt_llm/_torch/modules/fused_moe/routing.py`
eligibility; `tensorrt_llm/_torch/moe/fused_moe/routing.py`
`Deepseekv3RoutingImpl`);
`cpp/tensorrt_llm/kernels/noAuxTcKernels.cu`, `moeTopKFuncs.cuh`,
`cpp/tensorrt_llm/thop/noAuxTcOp.cpp`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,4 +33,4 @@ measured: []
- **Accuracy risk:** lossy (FP8 e4m3 block-scale; W4A8 = INT4 weights + FP8 activations). Each needs an on-disk accuracy record (GSM8K/MMLU) + rollback criterion before promotion. W4A8 depends on a calibrated act-scale file; bad calibration degrades accuracy.
- **Verify:** MoE GEMM throughput + weight footprint; accuracy parity vs a higher-precision backend (CUTLASS BF16/FP8); Blackwell confirm UE8M0 layout (sm_100 gate), Hopper confirm W4A8 group scaling (sm_90).
- **Rollback:** switch `moe_backend` to `CUTLASS` (or BF16/FP8-per-tensor) / load higher-precision checkpoint. Trigger: accuracy regression beyond threshold, or backend unsupported on the running SM.
- **Prior art:** PRs #6486, #4123. Files (Blackwell): `_torch/modules/fused_moe/fused_moe_deepgemm.py`, `create_moe.py`, `quantization/utils/fp8_utils.py`, `modeling_deepseekv3.py`. (Hopper): `_torch/modules/fused_moe.py` (`has_w4afp8`, `FusedMoEQuantScalesW4A8`), `thop/moeOp.cpp`, `examples/quantization/quantize_mixed_precision_moe.py`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
- **Prior art:** PRs #6486, #4123. Files (Blackwell): `_torch/moe/fused_moe/fused_moe_deepgemm.py`, `create_moe.py`, `quantization/utils/fp8_utils.py`, `modeling_deepseekv3.py`. (Hopper): `_torch/modules/fused_moe.py` (`has_w4afp8`, `FusedMoEQuantScalesW4A8`), `thop/moeOp.cpp`, `examples/quantization/quantize_mixed_precision_moe.py`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,4 @@ measured: []
- **Accuracy risk:** lossy (W4A8_MXFP4_MXFP8 MoE GEMM) but output-equivalent to the existing `TRTLLMGenFusedMoE` W4A8 backend by design (same math, different kernel) — parity is against that backend, not bf16.
- **Verify:** numeric parity vs `TRTLLMGenFusedMoE` same model/quant; task accuracy vs that backend; MoE collapses to a single kernel in nsys.
- **Rollback:** `moe_backend` → `TRTLLM` or `CUTLASS`. Trigger: `can_implement` rejects env (non-SM100, wrong quant, tp>1), missing DeepGEMM symbols, or parity failure.
- **Prior art:** PR #13384. Files: `_torch/modules/fused_moe/mega_moe/backend.py` (`can_implement`), `create_moe.py`, DeepGEMM `fp8_fp4_mega_moe`/`transform_weights_for_mega_moe`. Owning specialist: **kernel-cute-specialist**.
- **Prior art:** PR #13384. Files: `_torch/moe/fused_moe/mega_moe/backend.py` (`can_implement`), `create_moe.py`, DeepGEMM `fp8_fp4_mega_moe`/`transform_weights_for_mega_moe`. Owning specialist: **kernel-cute-specialist**.
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,4 @@ measured: []
- **Accuracy risk:** lossless — removes computation on padding/dummy tokens that never contributed to outputs.
- **Verify:** throughput/MFU up with imbalanced ISL; profile shows reduced dummy/padding work + lower host prep; output parity vs padded path.
- **Rollback:** revert to padded collective / full per-step padding. Trigger: correctness mismatch on ragged batches or collective hang.
- **Prior art:** PRs #6064, #3439, #4664. Files: `cpp/.../thop/{allgatherOp,reducescatterOp}.cpp`, `_torch/distributed/ops.py`, `_torch/pyexecutor/py_executor.py`, `_torch/modules/fused_moe/{fused_moe_cutlass,fused_moe_wide_ep}.py`. Detection: **perf-host-analysis**.
- **Prior art:** PRs #6064, #3439, #4664. Files: `cpp/.../thop/{allgatherOp,reducescatterOp}.cpp`, `_torch/distributed/ops.py`, `_torch/pyexecutor/py_executor.py`, `_torch/moe/fused_moe/{fused_moe_cutlass,fused_moe_wide_ep}.py`. Detection: **perf-host-analysis**.
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,4 @@ measured: []
- **Accuracy risk:** lossless (EPLB relocates expert replicas/slots across ranks; routing math and outputs unchanged — slots map back to the same experts).
- **Verify:** decode throughput + per-rank MoE-time balance (straggler spread) with online EPLB on vs off; confirm MoeBalancer stream concurrent with MoE GEMMs in nsys. Sanity-check accuracy unchanged.
- **Rollback:** `layer_updates_per_iter=0` (static routing) or omit the `moe_load_balancer` config. Trigger: no measured imbalance to recover, overlap not materializing, or instability during weight swap.
- **Prior art:** PRs #6624, #8886. Files: `_torch/modules/fused_moe/moe_load_balancer.py` (`start_wait_gpu_stage`, `update_local_statistic`, `is_static_routing`), `fused_moe/interface.py` (`_init_load_balancer`), `_torch/utils.py` (`AuxStreamType.MoeBalancer`), `llmapi/llm_args.py` (`MoeLoadBalancerConfig`).
- **Prior art:** PRs #6624, #8886. Files: `_torch/moe/fused_moe/moe_load_balancer.py` (`start_wait_gpu_stage`, `update_local_statistic`, `is_static_routing`), `fused_moe/interface.py` (`_init_load_balancer`), `_torch/utils.py` (`AuxStreamType.MoeBalancer`), `llmapi/llm_args.py` (`MoeLoadBalancerConfig`).
30 changes: 15 additions & 15 deletions .claude/skills/trtllm-moe-develop/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ description: >-
Review, design, and refactor TensorRT-LLM PyTorch MoE code for architecture fit,
clean code, maintainability, and testability. Always use for any modification,
review, refactor, or design planning that touches MoE modules, including
tensorrt_llm/_torch/modules/fused_moe, ConfigurableMoE, MoE backends,
tensorrt_llm/_torch/moe/fused_moe, ConfigurableMoE, MoE backends,
MoEScheduler/moe_scheduler.py, forward execution/chunking, communication
strategies, EPLB, quantization/weight
handling, routing, factories, MoE docs, or MoE tests. Also use when the user
Expand All @@ -26,23 +26,23 @@ style cleanup.
Before proposing or editing MoE code, read:

1. `CODING_GUIDELINES.md`
2. `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md`
2. `tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md`
3. The target files being changed
4. The relevant tests under `tests/unittest/_torch/modules/moe/`
4. The relevant tests under `tests/unittest/_torch/moe/`

Also inspect these files when the area is relevant:

- Forward execution/chunking: inspect `moe_scheduler.py`, `configurable_moe.py`,
`interface.py`, backend `run_moe`/`quantize_input` paths, and communication code.
- MegaMoE/fused communication: inspect `moe_scheduler.py`, `mega_moe/`,
`configurable_moe.py`, `quantization.py`, and communication code.
- Communication: `tensorrt_llm/_torch/modules/fused_moe/communication/base.py`
- Communication: `tensorrt_llm/_torch/moe/fused_moe/communication/base.py`
and `communication_factory.py`.
- Quantization and weights: `tensorrt_llm/_torch/modules/fused_moe/quantization.py`.
- Quantization and weights: `tensorrt_llm/_torch/moe/fused_moe/quantization.py`.
- EPLB/load balancing: `interface.py`, `moe_load_balancer.py`, `quantization.py`,
`moe_scheduler.py`, current forward-execution/chunking code, and
`test_moe_module.py`.
- Test matrix/helpers: `tests/unittest/_torch/modules/moe/moe_test_utils.py` and
- Test matrix/helpers: `tests/unittest/_torch/moe/moe_test_utils.py` and
`quantize_utils.py` when adding backend, quantization, skip, or parameter
coverage.

Expand Down Expand Up @@ -557,10 +557,10 @@ Role:

Main APIs / references:

- `tests/unittest/_torch/modules/moe/moe_test_utils.py`: `MoeBackendType`,
- `tests/unittest/_torch/moe/moe_test_utils.py`: `MoeBackendType`,
`get_backend_class`, `get_quick_skip_reason`, backend-specific
`should_skip_*`, `iter_base_test_configs`, CI acceleration logic.
- `tests/unittest/_torch/modules/moe/quantize_utils.py`: quantized test weight
- `tests/unittest/_torch/moe/quantize_utils.py`: quantized test weight
generation and quant-parameter setup.
- `test_moe_backend.py`: backend interface tests for `quantize_input` and
`run_moe`.
Expand Down Expand Up @@ -667,13 +667,13 @@ If there are no findings, say so and list remaining test or performance risk.

Prefer the unified MoE tests:

- Shared test matrix/helper changes: inspect `tests/unittest/_torch/modules/moe/moe_test_utils.py` and `quantize_utils.py`, then run the affected backend/module tests below.
- Backend interface changes: `pytest tests/unittest/_torch/modules/moe/test_moe_backend.py -k '<backend or quant>'`.
- Module/create/forward changes: `pytest tests/unittest/_torch/modules/moe/test_moe_module.py -k '<backend or feature>'`.
- Communication changes: `pytest tests/unittest/_torch/modules/moe/test_moe_comm.py -k '<strategy>'`.
- Routing changes: `pytest tests/unittest/_torch/modules/test_moe_routing.py -k '<routing>'`.
- Load balancer changes: `pytest tests/unittest/_torch/modules/test_moe_load_balancer.py -k '<case>'`.
- Multi-GPU EP/all-to-all behavior: `pytest tests/unittest/_torch/multi_gpu/test_moe_a2a.py -k '<case>'`.
- Shared test matrix/helper changes: inspect `tests/unittest/_torch/moe/moe_test_utils.py` and `quantize_utils.py`, then run the affected backend/module tests below.
- Backend interface changes: `pytest tests/unittest/_torch/moe/test_moe_backend.py -k '<backend or quant>'`.
- Module/create/forward changes: `pytest tests/unittest/_torch/moe/test_moe_module.py -k '<backend or feature>'`.
- Communication changes: `pytest tests/unittest/_torch/moe/test_moe_comm.py -k '<strategy>'`.
- Routing changes: `pytest tests/unittest/_torch/moe/test_moe_routing.py -k '<routing>'`.
- Load balancer changes: `pytest tests/unittest/_torch/moe/test_moe_load_balancer.py -k '<case>'`.
- Multi-GPU EP/all-to-all behavior: `pytest tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py -k '<case>'`.

When GPU resources are required, use the TRT-LLM GPU allocation/test-runner
skills first and record skipped tests with reasons.
Loading
Loading