diff --git a/.claude/agents/perf-profiling-specialist.md b/.claude/agents/perf-profiling-specialist.md index 03efc802ab32..967766f0f869 100644 --- a/.claude/agents/perf-profiling-specialist.md +++ b/.claude/agents/perf-profiling-specialist.md @@ -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 diff --git a/.claude/skills/perf-optimization-casebook/references/communication/deepep.md b/.claude/skills/perf-optimization-casebook/references/communication/deepep.md index b1407afcdec2..ea1ae54fb032 100644 --- a/.claude/skills/perf-optimization-casebook/references/communication/deepep.md +++ b/.claude/skills/perf-optimization-casebook/references/communication/deepep.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md b/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md index 8ea0ec641034..f80822d41eae 100644 --- a/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md +++ b/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md index 701a892bec02..ae74bc55413a 100644 --- a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md +++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md index 2930d01cb5a2..2ab00eb9567c 100644 --- a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md +++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md @@ -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: @@ -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`, diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md index eb521cdb9f62..d94b7504e9f8 100644 --- a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md +++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md index 0cf219c88911..7a10e4fe6434 100644 --- a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md +++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md index ec4f11af7021..1a6d75492b9e 100644 --- a/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md +++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md @@ -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**. diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md index 95f355c3ddd0..d779fd2185e3 100644 --- a/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md +++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md @@ -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`). diff --git a/.claude/skills/trtllm-moe-develop/SKILL.md b/.claude/skills/trtllm-moe-develop/SKILL.md index ab7b282bc415..0b644259e39d 100644 --- a/.claude/skills/trtllm-moe-develop/SKILL.md +++ b/.claude/skills/trtllm-moe-develop/SKILL.md @@ -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 @@ -26,9 +26,9 @@ 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: @@ -36,13 +36,13 @@ Also inspect these files when the area is relevant: `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. @@ -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`. @@ -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 ''`. -- Module/create/forward changes: `pytest tests/unittest/_torch/modules/moe/test_moe_module.py -k ''`. -- Communication changes: `pytest tests/unittest/_torch/modules/moe/test_moe_comm.py -k ''`. -- Routing changes: `pytest tests/unittest/_torch/modules/test_moe_routing.py -k ''`. -- Load balancer changes: `pytest tests/unittest/_torch/modules/test_moe_load_balancer.py -k ''`. -- Multi-GPU EP/all-to-all behavior: `pytest tests/unittest/_torch/multi_gpu/test_moe_a2a.py -k ''`. +- 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 ''`. +- Module/create/forward changes: `pytest tests/unittest/_torch/moe/test_moe_module.py -k ''`. +- Communication changes: `pytest tests/unittest/_torch/moe/test_moe_comm.py -k ''`. +- Routing changes: `pytest tests/unittest/_torch/moe/test_moe_routing.py -k ''`. +- Load balancer changes: `pytest tests/unittest/_torch/moe/test_moe_load_balancer.py -k ''`. +- Multi-GPU EP/all-to-all behavior: `pytest tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py -k ''`. When GPU resources are required, use the TRT-LLM GPU allocation/test-runner skills first and record skipped tests with reasons. diff --git a/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md b/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md index df4d412c4aca..3751ee8c5d41 100644 --- a/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md +++ b/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md @@ -8,7 +8,7 @@ then inspect the actual repository files. Use these examples to understand wrapper-level lifecycle and composition: -- `tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py` +- `tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py` - `ConfigurableMoE.__init__`: backend construction, communication creation, chunking-related wrapper state, validation, DWDP setup, and weight-removal lifecycle. @@ -38,7 +38,7 @@ Red flags: Use these examples when changing forward policy, chunking, communication order, EPLB hook order, or fused-communication behavior: -- `tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py` +- `tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py` - `MoEScheduler`: abstract forward strategy. - `ExternalCommMoEScheduler`: host-side communication path, chunking, adaptive quantize/dispatch ordering, communication fallback, multi-stream overlap, @@ -48,10 +48,10 @@ EPLB hook order, or fused-communication behavior: no external `Communication`, `ignore_allreduce=False`, and backend `run_moe` invocation. - `create_moe_scheduler`: factory keyed by `backend.scheduler_kind`. -- `tensorrt_llm/_torch/modules/fused_moe/interface.py` +- `tensorrt_llm/_torch/moe/fused_moe/interface.py` - `MoESchedulerKind`: backend-declared selection axis for external vs fused cross-rank exchange. -- `tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py` +- `tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py` - Scheduler is constructed after backend/comm/chunking/DWDP state is ready; wrapper `forward_impl` delegates and then runs shared lifecycle bookkeeping. @@ -80,7 +80,7 @@ Red flags: Use these examples to decide what belongs in a backend: -- `tensorrt_llm/_torch/modules/fused_moe/interface.py` +- `tensorrt_llm/_torch/moe/fused_moe/interface.py` - `MoE`: current transitional base contract shared by legacy standalone MoE modules and newer ConfigurableMoE-compatible backends. A dedicated backend interface is expected in the future; do not treat legacy `forward` methods @@ -101,22 +101,22 @@ Use these examples to decide what belongs in a backend: - `forward` and `forward_impl`: legacy standalone-module entrypoints only; avoid adding them to new ConfigurableMoE-compatible backends unless the user explicitly requests that behavior. -- `tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py` +- `tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py` - Reference backend for normal external communication and broad quant support. - Inspect `_get_quant_method`, `create_weights`, `load_weights`, and `post_load_weights` for the standard backend-to-quantization delegation pattern. -- `tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py` +- `tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py` - Reference backend for TRTLLMGen-specific input/scale contracts. - Inspect `_get_quant_method`, alignment validation, weight lifecycle delegation, and quant scale usage. -- `tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py` +- `tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py` - Reference backend for DeepGEMM workspace and Blackwell FP8 block-scale path. -- `tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py` when present +- `tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py` when present - Reference for a direct `MoE` backend with dense GEMM constraints. - Inspect backend-owned `create_weights`, `load_weights`, and `post_load_weights` delegation. -- `tensorrt_llm/_torch/modules/fused_moe/mega_moe/` when present +- `tensorrt_llm/_torch/moe/fused_moe/mega_moe/` when present - Reference area for a backend whose kernel owns cross-rank exchange. Inspect capability checks, load-balancer support, validation, weight lifecycle delegation, `quantize_input`, and `run_moe`. @@ -162,7 +162,7 @@ the backend section: backend owns the module-level lifecycle entrypoints and method selection; quantization owns the delegated tensor layout, loading details, transforms, and scale registration for a quantization layout: -- `tensorrt_llm/_torch/modules/fused_moe/quantization.py` +- `tensorrt_llm/_torch/moe/fused_moe/quantization.py` - `FusedMoEMethodBase`: common weight create/load/post-load flow, shared weight loading, EPLB hooks, and quant scale setup. - `UnquantizedFusedMoEMethod`: simple supported-EPLB baseline. @@ -197,23 +197,23 @@ Red flags: Use these examples when touching dynamic expert migration or slot routing: -- `tensorrt_llm/_torch/modules/fused_moe/interface.py` +- `tensorrt_llm/_torch/moe/fused_moe/interface.py` - `_using_load_balancer`, `_using_dynamic_load_balancer`. - `_load_balancer_update_statistic` and `_load_balancer_update_statistic_with_gathered_statistic`. - `_load_balancer_route` and slot/expert-ID conversion. - `register_all_parameter_slot_and_to_fix_weight_fns`. -- `tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py` +- `tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py` - Slot registration, expert migration, per-layer update flow, and iteration context. -- `tensorrt_llm/_torch/modules/fused_moe/quantization.py` +- `tensorrt_llm/_torch/moe/fused_moe/quantization.py` - `eplb_support_status`, `supports_online_eplb`, `need_load_shared_weights`, shared-weight registration, and transformed weight fix-up hooks. - Current forward-execution code - EPLB wait/update/route hook order, per-chunk first/last behavior, statistics gathering, and allreduce bypass semantics. -- `tests/unittest/_torch/modules/moe/test_moe_module.py` +- `tests/unittest/_torch/moe/test_moe_module.py` - `_create_moe_load_balancer`, `_run_eplb_test`, `_should_skip_EPLB`, `generate_eplb_test_params`, and backend-specific EPLB param generators when present. @@ -246,7 +246,7 @@ Use these exact locations when adding or reviewing a per-expert Parameter that must survive online EPLB migration. - Bulk FC weights (base-class path, every EPLB-supporting quant method): - - `tensorrt_llm/_torch/modules/fused_moe/quantization.py::FusedMoEMethodBase.load_weights` + - `tensorrt_llm/_torch/moe/fused_moe/quantization.py::FusedMoEMethodBase.load_weights` -- `need_load_shared_weights(module)` branch allocates and fills `module.local_shared_w3_w1_tensors` / `module.local_shared_w2_tensors` (and bias twins when `module.bias`) sized @@ -329,7 +329,7 @@ routed/shared index spaces overlap on the same global expert IDs. Use these examples when changing dispatch/combine behavior: -- `tensorrt_llm/_torch/modules/fused_moe/communication/base.py` +- `tensorrt_llm/_torch/moe/fused_moe/communication/base.py` - `Communication` ABC, `supports_post_quant_dispatch`, `prepare_dispatch`, `dispatch`, and `combine`. - `communication/communication_factory.py` @@ -342,7 +342,7 @@ Use these examples when changing dispatch/combine behavior: - Dispatch-time statistics/workspace pattern. - `communication/deep_ep.py` and `communication/deep_ep_low_latency.py` - DeepEP pre/post quant dispatch constraints. -- `tests/unittest/_torch/modules/moe/test_moe_comm.py` +- `tests/unittest/_torch/moe/test_moe_comm.py` - Focused communication behavior tests. Good uses: @@ -360,18 +360,18 @@ Red flags: Use these examples when wrapper forward policy grows complicated: -- `tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py` +- `tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py` - Scheduler forward entry, external/fused scheduler implementations, routing order, dispatch order, EPLB hook order, zero-token behavior, output truncation, and backend kwargs. -- `tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py` +- `tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py` - Wrapper entry and lifecycle boundary around scheduler execution. - Backend forward paths - Useful only for legacy compatibility context or for identifying policy that should move into the scheduler. -- `tests/unittest/_torch/modules/moe/test_moe_module.py` +- `tests/unittest/_torch/moe/test_moe_module.py` - Module-level multi-GPU, chunking, routing, and EPLB cases. -- `tests/unittest/_torch/multi_gpu/test_moe_a2a.py` +- `tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py` - Multi-GPU all-to-all behavior when relevant. Good uses: @@ -390,12 +390,12 @@ Red flags: Use these examples when changing model-to-backend selection or routing output: -- `tensorrt_llm/_torch/modules/fused_moe/routing.py` +- `tensorrt_llm/_torch/moe/fused_moe/routing.py` - Routing method output shape/dtype and top-k behavior. -- `tensorrt_llm/_torch/modules/fused_moe/create_moe.py` +- `tensorrt_llm/_torch/moe/fused_moe/create_moe.py` - `get_moe_cls`, `create_moe_backend`, `create_moe`, backend fallback, and quantization-specific selection. -- `tests/unittest/_torch/modules/moe/moe_test_utils.py` +- `tests/unittest/_torch/moe/moe_test_utils.py` - Backend enum, class map, quick skip reason, CI/local matrix generation. Good uses: @@ -412,18 +412,18 @@ Red flags: Use these examples before adding backend, quantization, routing, or EPLB tests: -- `tests/unittest/_torch/modules/moe/moe_test_utils.py` +- `tests/unittest/_torch/moe/moe_test_utils.py` - `MoeBackendType`, `get_backend_class`, backend-specific `should_skip_*`, `get_quick_skip_reason`, `supports_autotuner_capture`, `iter_base_test_configs`, and CI acceleration logic. -- `tests/unittest/_torch/modules/moe/quantize_utils.py` +- `tests/unittest/_torch/moe/quantize_utils.py` - Quantized test weight generation, reference module selection, backend-aware weight preparation. -- `tests/unittest/_torch/modules/moe/test_moe_backend.py` +- `tests/unittest/_torch/moe/test_moe_backend.py` - Backend-level `quantize_input` and `run_moe` contracts. -- `tests/unittest/_torch/modules/moe/test_moe_module.py` +- `tests/unittest/_torch/moe/test_moe_module.py` - ConfigurableMoE integration matrix, multi-GPU, and EPLB coverage. -- `tests/unittest/_torch/modules/moe/test_moe_comm.py` +- `tests/unittest/_torch/moe/test_moe_comm.py` - Communication strategy tests. Good uses: diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e27b26e5d88d..1cd30750397f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -365,24 +365,14 @@ /cpp/tensorrt_llm/runtime/moeLoadBalancer @NVIDIA/trt-llm-moe-devs /tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm @NVIDIA/trt-llm-moe-devs /tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4 @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/distributed/moe_alltoall.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/expert_statistic.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/modules/fused_moe @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/modules/fused_shared_expert.py @NVIDIA/trt-llm-moe-devs +/tensorrt_llm/_torch/moe @NVIDIA/trt-llm-moe-devs /tensorrt_llm/deep_ep @NVIDIA/trt-llm-moe-devs /tests/microbenchmarks/bench_moe @NVIDIA/trt-llm-moe-devs /tests/microbenchmarks/bench_moe_comm.py @NVIDIA/trt-llm-moe-devs /tests/microbenchmarks/compare_moe_comm.py @NVIDIA/trt-llm-moe-devs /tests/scripts/cute_dsl_kernels/moe_as_dense_gemm @NVIDIA/trt-llm-moe-devs /tests/scripts/cute_dsl_kernels/moe_workload_generator.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/fused_moe @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/moe @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_fused_moe.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_fused_shared_expert.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_moe_*.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/multi_gpu/test_moe_a2a.py @NVIDIA/trt-llm-moe-devs +/tests/unittest/_torch/moe @NVIDIA/trt-llm-moe-devs /tests/unittest/_torch/thop/parallel/*moe* @NVIDIA/trt-llm-moe-devs /tests/unittest/_torch/thop/serial/*moe* @NVIDIA/trt-llm-moe-devs /tests/unittest/bindings/test_bindings_moe.py @NVIDIA/trt-llm-moe-devs diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2a259562ce95..f8c606585607 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -168,10 +168,8 @@ common-files: &common_files | tensorrt_llm/_torch/device_mesh.py | tensorrt_llm/_torch/distributed/__init__.py | tensorrt_llm/_torch/distributed/communicator.py | - tensorrt_llm/_torch/distributed/moe_alltoall.py | tensorrt_llm/_torch/distributed/ops.py | tensorrt_llm/_torch/distributed/pg_utils.py | - tensorrt_llm/_torch/expert_statistic.py | tensorrt_llm/_torch/flashinfer_utils.py | tensorrt_llm/_torch/hostfunc.py | tensorrt_llm/_torch/llm.py | @@ -253,19 +251,6 @@ common-files: &common_files | tensorrt_llm/_torch/modules/fla/solve_tril.py | tensorrt_llm/_torch/modules/fla/utils.py | tensorrt_llm/_torch/modules/fla/wy_fast.py | - tensorrt_llm/_torch/modules/fused_moe/__init__.py | - tensorrt_llm/_torch/modules/fused_moe/create_moe.py | - tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py | - tensorrt_llm/_torch/modules/fused_moe/interface.py | - tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py | - tensorrt_llm/_torch/modules/fused_moe/quantization.py | - tensorrt_llm/_torch/modules/fused_moe/routing.py | tensorrt_llm/_torch/modules/gated_mlp.py | tensorrt_llm/_torch/modules/layer_norm.py | tensorrt_llm/_torch/modules/linear.py | @@ -289,6 +274,21 @@ common-files: &common_files | tensorrt_llm/_torch/modules/rotary_embedding.py | tensorrt_llm/_torch/modules/swiglu.py | tensorrt_llm/_torch/modules/triton_linear.py | + tensorrt_llm/_torch/moe/expert_statistic.py | + tensorrt_llm/_torch/moe/fused_moe/__init__.py | + tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py | + tensorrt_llm/_torch/moe/fused_moe/create_moe.py | + tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py | + tensorrt_llm/_torch/moe/fused_moe/interface.py | + tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py | + tensorrt_llm/_torch/moe/fused_moe/quantization.py | + tensorrt_llm/_torch/moe/fused_moe/routing.py | tensorrt_llm/_torch/peft/__init__.py | tensorrt_llm/_torch/peft/lora/__init__.py | tensorrt_llm/_torch/peft/lora/config.py | @@ -587,13 +587,14 @@ common-files: &common_files | tests/unittest/_torch/modeling/test_modeling_qwen_moe.py | tests/unittest/_torch/modeling/test_modeling_siglip.py | tests/unittest/_torch/modeling/test_modeling_vila.py | - tests/unittest/_torch/modules/test_fused_moe.py | tests/unittest/_torch/modules/test_group_rmn_norm.py | - tests/unittest/_torch/modules/test_moe_host_sharer.py | - tests/unittest/_torch/modules/test_moe_load_balancer.py | - tests/unittest/_torch/modules/test_moe_routing.py | tests/unittest/_torch/modules/test_rotary_embedding.py | tests/unittest/_torch/modules/test_triton_linear.py | + tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py | + tests/unittest/_torch/moe/test_fused_moe.py | + tests/unittest/_torch/moe/test_moe_host_sharer.py | + tests/unittest/_torch/moe/test_moe_load_balancer.py | + tests/unittest/_torch/moe/test_moe_routing.py | tests/unittest/_torch/multi_gpu/test_allreduce.py | tests/unittest/_torch/multi_gpu/test_alltoall.py | tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py | @@ -602,7 +603,6 @@ common-files: &common_files | tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py | tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py | tests/unittest/_torch/multi_gpu/test_mnnvl_memory.py | - tests/unittest/_torch/multi_gpu/test_moe_a2a.py | tests/unittest/_torch/multi_gpu/test_user_buffers.py | tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py | tests/unittest/_torch/multimodal/test_external_embedding.py | @@ -934,10 +934,8 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/device_mesh.py | tensorrt_llm/_torch/distributed/__init__.py | tensorrt_llm/_torch/distributed/communicator.py | - tensorrt_llm/_torch/distributed/moe_alltoall.py | tensorrt_llm/_torch/distributed/ops.py | tensorrt_llm/_torch/distributed/pg_utils.py | - tensorrt_llm/_torch/expert_statistic.py | tensorrt_llm/_torch/flashinfer_utils.py | tensorrt_llm/_torch/hostfunc.py | tensorrt_llm/_torch/llm.py | @@ -1019,19 +1017,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/modules/fla/solve_tril.py | tensorrt_llm/_torch/modules/fla/utils.py | tensorrt_llm/_torch/modules/fla/wy_fast.py | - tensorrt_llm/_torch/modules/fused_moe/__init__.py | - tensorrt_llm/_torch/modules/fused_moe/create_moe.py | - tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py | - tensorrt_llm/_torch/modules/fused_moe/interface.py | - tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py | - tensorrt_llm/_torch/modules/fused_moe/quantization.py | - tensorrt_llm/_torch/modules/fused_moe/routing.py | tensorrt_llm/_torch/modules/gated_mlp.py | tensorrt_llm/_torch/modules/layer_norm.py | tensorrt_llm/_torch/modules/linear.py | @@ -1055,6 +1040,21 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/modules/rotary_embedding.py | tensorrt_llm/_torch/modules/swiglu.py | tensorrt_llm/_torch/modules/triton_linear.py | + tensorrt_llm/_torch/moe/expert_statistic.py | + tensorrt_llm/_torch/moe/fused_moe/__init__.py | + tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py | + tensorrt_llm/_torch/moe/fused_moe/create_moe.py | + tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py | + tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py | + tensorrt_llm/_torch/moe/fused_moe/interface.py | + tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py | + tensorrt_llm/_torch/moe/fused_moe/quantization.py | + tensorrt_llm/_torch/moe/fused_moe/routing.py | tensorrt_llm/_torch/peft/__init__.py | tensorrt_llm/_torch/peft/lora/__init__.py | tensorrt_llm/_torch/peft/lora/config.py | @@ -1353,13 +1353,14 @@ legacy-files: &legacy_files | tests/unittest/_torch/modeling/test_modeling_qwen_moe.py | tests/unittest/_torch/modeling/test_modeling_siglip.py | tests/unittest/_torch/modeling/test_modeling_vila.py | - tests/unittest/_torch/modules/test_fused_moe.py | tests/unittest/_torch/modules/test_group_rmn_norm.py | - tests/unittest/_torch/modules/test_moe_host_sharer.py | - tests/unittest/_torch/modules/test_moe_load_balancer.py | - tests/unittest/_torch/modules/test_moe_routing.py | tests/unittest/_torch/modules/test_rotary_embedding.py | tests/unittest/_torch/modules/test_triton_linear.py | + tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py | + tests/unittest/_torch/moe/test_fused_moe.py | + tests/unittest/_torch/moe/test_moe_host_sharer.py | + tests/unittest/_torch/moe/test_moe_load_balancer.py | + tests/unittest/_torch/moe/test_moe_routing.py | tests/unittest/_torch/multi_gpu/test_allreduce.py | tests/unittest/_torch/multi_gpu/test_alltoall.py | tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py | @@ -1368,7 +1369,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py | tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py | tests/unittest/_torch/multi_gpu/test_mnnvl_memory.py | - tests/unittest/_torch/multi_gpu/test_moe_a2a.py | tests/unittest/_torch/multi_gpu/test_user_buffers.py | tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py | tests/unittest/_torch/multimodal/test_external_embedding.py | diff --git a/AGENTS.md b/AGENTS.md index 51853d3a963b..519a4bc1affb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -85,7 +85,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy) | `tensorrt_llm/models/automodel.py` | Auto-discovery and model registry | | `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from the top-level `models/` package) | | `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py`, `tensorrt_llm/_torch/modules/mla.py`, or `tensorrt_llm/_torch/attention_backend/`** | -| `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | +| `tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | | `CODING_GUIDELINES.md` | C++ and Python coding standards (referenced throughout, must read before contributing) | ## Design Patterns diff --git a/agent-flow/agent_flow/workflows/modeling_bringup/prompts/_common.py b/agent-flow/agent_flow/workflows/modeling_bringup/prompts/_common.py index f075f7ea8601..fa37b7f59cc4 100644 --- a/agent-flow/agent_flow/workflows/modeling_bringup/prompts/_common.py +++ b/agent-flow/agent_flow/workflows/modeling_bringup/prompts/_common.py @@ -107,7 +107,7 @@ def get_trtllm_test_specialist_invocation() -> str: `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` and reason across module, backend, runtime contract, and KV-cache semantics. - For MoE work, read - `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` and reason + `tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md` and reason across routing, expert parallelism, quantization, and fused-kernel contracts. - Use `KVCacheManagerV2` for new-model bring-up. The TRTLLM and FlashInfer attention backends are both valid targets when the plan tests the selected diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index 1a0020c8fdd7..fce5010dceab 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -487,7 +487,7 @@ void Runner::run(void* hiddenState, void* hiddenStateScale, void* weights, void* // The multiple is no less than 128 as TMA requires it for CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B types // FIXME: enforce valid hidden dim to be multiple of 512 due to unhandled OOB read in routeAct. Please keep this // in sync with - // tensorrt_llm/_torch/modules/fused_moe/quantization.py:MXFP4WeightTRTLLMGenFusedMoEMethod.input_hidden_alignment + // tensorrt_llm/_torch/moe/fused_moe/quantization.py:MXFP4WeightTRTLLMGenFusedMoEMethod.input_hidden_alignment validHiddenSize = tensorrt_llm::common::roundUp(validHiddenSize, 512); } auto maxNumCgasInBatchDim = Routing::getMaxNumCgasInBatchDim(numTokens, topK, numExperts, mTileTokensDim); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index 85333a05945a..9987eece14df 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -65,7 +65,7 @@ namespace Routing { // The type of method in top-K routing, for use in torch custom op -// Please keep this in sync with the counterpart defined in tensorrt_llm/_torch/modules/fused_moe/routing.py +// Please keep this in sync with the counterpart defined in tensorrt_llm/_torch/moe/fused_moe/routing.py enum class RoutingMethodType : int64_t { // Default: Softmax -> TopK diff --git a/docs/source/blogs/tech_blog/blog18_Optimizing_MoE_Communication_with_One_Sided_AlltoAll_Over_NVLink.md b/docs/source/blogs/tech_blog/blog18_Optimizing_MoE_Communication_with_One_Sided_AlltoAll_Over_NVLink.md index 52318a4272dc..5782f7059916 100644 --- a/docs/source/blogs/tech_blog/blog18_Optimizing_MoE_Communication_with_One_Sided_AlltoAll_Over_NVLink.md +++ b/docs/source/blogs/tech_blog/blog18_Optimizing_MoE_Communication_with_One_Sided_AlltoAll_Over_NVLink.md @@ -84,7 +84,7 @@ Rank-major layout has two advantages: - **Smaller buffer.** The pre-allocated recv buffer is `1 / num_experts_per_rank` of an expert-major buffer. - **Save duplicated communication.** When `top_k > ep_size` it is impossible to avoid duplication in an expert-major layout. -The MoE module consumes the rank-major recv buffer directly, and it is up to the MoE module to decide whether explicit expert permutation is needed for GroupGEMM. For example, [trtllm-gen MoE](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py) can efficiently load tokens from the raw rank-major buffer without additional permutation. +The MoE module consumes the rank-major recv buffer directly, and it is up to the MoE module to decide whether explicit expert permutation is needed for GroupGEMM. For example, [trtllm-gen MoE](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py) can efficiently load tokens from the raw rank-major buffer without additional permutation. ### Interface Between Communication and MoE @@ -107,7 +107,7 @@ Here `hidden_states` is the (possibly quantized) activation tensor with `hidden_ The dispatch outputs are **tensor views** into symmetric memory — no allocation, no copy. The recv buffer is pre-allocated with `ep_size * max_tokens_per_rank` slots to accommodate the maximum number of tokens from all ranks, as described in [Rank-Major Buffer Layout](#rank-major-buffer-layout). The MoE module then performs GroupGEMM on the received payloads. Two points are worth noting: - The MoE module only computes the experts on the local rank. For example, if a token's `token_selected_experts` is `[0, 1, 4, 7]` and only experts `[0, 1, 2, 3]` reside locally, the MoE output for that token is the weighted sum of experts `[0, 1]` only. -- Some slots are empty after dispatch. The dispatch kernel sets `token_selected_experts` of empty slots to an invalid expert id (`-1`), so the MoE module knows to skip them. (For instance, [trtllm-gen MoE](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py) consumes the raw-token recv buffer directly without re-permutation.) +- Some slots are empty after dispatch. The dispatch kernel sets `token_selected_experts` of empty slots to an invalid expert id (`-1`), so the MoE module knows to skip them. (For instance, [trtllm-gen MoE](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py) consumes the raw-token recv buffer directly without re-permutation.) The MoE module writes its output for each token to the same slot in the recv buffer. To obtain the final result, each rank combines the partial results from peer ranks. diff --git a/docs/source/blogs/tech_blog/blog24_MoE_as_Dense_GEMM.md b/docs/source/blogs/tech_blog/blog24_MoE_as_Dense_GEMM.md index 6f834e365b8c..fcc7e5b71d8f 100644 --- a/docs/source/blogs/tech_blog/blog24_MoE_as_Dense_GEMM.md +++ b/docs/source/blogs/tech_blog/blog24_MoE_as_Dense_GEMM.md @@ -424,8 +424,8 @@ Future directions include: ## References -- Source code: [`tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py) -- Backend selection: [`tensorrt_llm/_torch/modules/fused_moe/create_moe.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/fused_moe/create_moe.py) +- Source code: [`tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py) +- Backend selection: [`tensorrt_llm/_torch/moe/fused_moe/create_moe.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/moe/fused_moe/create_moe.py) - FC1 kernel: [`tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc1.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc1.py) - FC2 kernel: [`tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py) - Custom op registration: [`tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 5924edc6d4de..66108764bd0c 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1283,7 +1283,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/_torch/distributed/", "tensorrt_llm/_torch/models/modeling_llama.py", "tensorrt_llm/_torch/models/modeling_qwen3_next.py", - "tensorrt_llm/_torch/modules/fused_moe/", + "tensorrt_llm/_torch/moe/", "tensorrt_llm/_torch/pyexecutor/_util.py", "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py", "tensorrt_llm/_torch/pyexecutor/model_engine.py", @@ -1354,6 +1354,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml", "tests/integration/test_lists/test-db/l0_verl.yml", "tests/unittest/_torch/multi_gpu/", + "tests/unittest/_torch/moe/multi_gpu/", "tests/unittest/_torch/multi_gpu_modeling/", "tests/unittest/_torch/visual_gen/multi_gpu/", "tests/unittest/disaggregated/", diff --git a/legacy-files.txt b/legacy-files.txt index 73fe3fee5899..099b17d8e6c3 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -160,10 +160,10 @@ tensorrt_llm/_torch/debug/debug_hook.py tensorrt_llm/_torch/device_mesh.py tensorrt_llm/_torch/distributed/__init__.py tensorrt_llm/_torch/distributed/communicator.py -tensorrt_llm/_torch/distributed/moe_alltoall.py +tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py tensorrt_llm/_torch/distributed/ops.py tensorrt_llm/_torch/distributed/pg_utils.py -tensorrt_llm/_torch/expert_statistic.py +tensorrt_llm/_torch/moe/expert_statistic.py tensorrt_llm/_torch/flashinfer_utils.py tensorrt_llm/_torch/hostfunc.py tensorrt_llm/_torch/llm.py @@ -245,19 +245,19 @@ tensorrt_llm/_torch/modules/fla/op.py tensorrt_llm/_torch/modules/fla/solve_tril.py tensorrt_llm/_torch/modules/fla/utils.py tensorrt_llm/_torch/modules/fla/wy_fast.py -tensorrt_llm/_torch/modules/fused_moe/__init__.py -tensorrt_llm/_torch/modules/fused_moe/create_moe.py -tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py -tensorrt_llm/_torch/modules/fused_moe/interface.py -tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py -tensorrt_llm/_torch/modules/fused_moe/quantization.py -tensorrt_llm/_torch/modules/fused_moe/routing.py +tensorrt_llm/_torch/moe/fused_moe/__init__.py +tensorrt_llm/_torch/moe/fused_moe/create_moe.py +tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py +tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py +tensorrt_llm/_torch/moe/fused_moe/interface.py +tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +tensorrt_llm/_torch/moe/fused_moe/quantization.py +tensorrt_llm/_torch/moe/fused_moe/routing.py tensorrt_llm/_torch/modules/gated_mlp.py tensorrt_llm/_torch/modules/layer_norm.py tensorrt_llm/_torch/modules/linear.py @@ -579,11 +579,11 @@ tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py tests/unittest/_torch/modeling/test_modeling_qwen_moe.py tests/unittest/_torch/modeling/test_modeling_siglip.py tests/unittest/_torch/modeling/test_modeling_vila.py -tests/unittest/_torch/modules/test_fused_moe.py +tests/unittest/_torch/moe/test_fused_moe.py tests/unittest/_torch/modules/test_group_rmn_norm.py -tests/unittest/_torch/modules/test_moe_host_sharer.py -tests/unittest/_torch/modules/test_moe_load_balancer.py -tests/unittest/_torch/modules/test_moe_routing.py +tests/unittest/_torch/moe/test_moe_host_sharer.py +tests/unittest/_torch/moe/test_moe_load_balancer.py +tests/unittest/_torch/moe/test_moe_routing.py tests/unittest/_torch/modules/test_rotary_embedding.py tests/unittest/_torch/modules/test_triton_linear.py tests/unittest/_torch/multi_gpu/test_allreduce.py @@ -594,7 +594,7 @@ tests/unittest/_torch/multi_gpu/test_linear.py tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py tests/unittest/_torch/multi_gpu/test_mnnvl_memory.py -tests/unittest/_torch/multi_gpu/test_moe_a2a.py +tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py tests/unittest/_torch/multi_gpu/test_user_buffers.py tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py tests/unittest/_torch/multimodal/test_external_embedding.py diff --git a/pyproject.toml b/pyproject.toml index f8e1447ca03b..0dbb75731b0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -217,10 +217,8 @@ exclude = [ "tensorrt_llm/_torch/device_mesh.py", "tensorrt_llm/_torch/distributed/__init__.py", "tensorrt_llm/_torch/distributed/communicator.py", - "tensorrt_llm/_torch/distributed/moe_alltoall.py", "tensorrt_llm/_torch/distributed/ops.py", "tensorrt_llm/_torch/distributed/pg_utils.py", - "tensorrt_llm/_torch/expert_statistic.py", "tensorrt_llm/_torch/flashinfer_utils.py", "tensorrt_llm/_torch/hostfunc.py", "tensorrt_llm/_torch/llm.py", @@ -302,19 +300,6 @@ exclude = [ "tensorrt_llm/_torch/modules/fla/solve_tril.py", "tensorrt_llm/_torch/modules/fla/utils.py", "tensorrt_llm/_torch/modules/fla/wy_fast.py", - "tensorrt_llm/_torch/modules/fused_moe/__init__.py", - "tensorrt_llm/_torch/modules/fused_moe/create_moe.py", - "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py", - "tensorrt_llm/_torch/modules/fused_moe/interface.py", - "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py", - "tensorrt_llm/_torch/modules/fused_moe/quantization.py", - "tensorrt_llm/_torch/modules/fused_moe/routing.py", "tensorrt_llm/_torch/modules/gated_mlp.py", "tensorrt_llm/_torch/modules/layer_norm.py", "tensorrt_llm/_torch/modules/linear.py", @@ -338,6 +323,21 @@ exclude = [ "tensorrt_llm/_torch/modules/rotary_embedding.py", "tensorrt_llm/_torch/modules/swiglu.py", "tensorrt_llm/_torch/modules/triton_linear.py", + "tensorrt_llm/_torch/moe/expert_statistic.py", + "tensorrt_llm/_torch/moe/fused_moe/__init__.py", + "tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py", + "tensorrt_llm/_torch/moe/fused_moe/create_moe.py", + "tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py", + "tensorrt_llm/_torch/moe/fused_moe/interface.py", + "tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py", + "tensorrt_llm/_torch/moe/fused_moe/quantization.py", + "tensorrt_llm/_torch/moe/fused_moe/routing.py", "tensorrt_llm/_torch/peft/__init__.py", "tensorrt_llm/_torch/peft/lora/__init__.py", "tensorrt_llm/_torch/peft/lora/config.py", @@ -636,13 +636,14 @@ exclude = [ "tests/unittest/_torch/modeling/test_modeling_qwen_moe.py", "tests/unittest/_torch/modeling/test_modeling_siglip.py", "tests/unittest/_torch/modeling/test_modeling_vila.py", - "tests/unittest/_torch/modules/test_fused_moe.py", "tests/unittest/_torch/modules/test_group_rmn_norm.py", - "tests/unittest/_torch/modules/test_moe_host_sharer.py", - "tests/unittest/_torch/modules/test_moe_load_balancer.py", - "tests/unittest/_torch/modules/test_moe_routing.py", "tests/unittest/_torch/modules/test_rotary_embedding.py", "tests/unittest/_torch/modules/test_triton_linear.py", + "tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py", + "tests/unittest/_torch/moe/test_fused_moe.py", + "tests/unittest/_torch/moe/test_moe_host_sharer.py", + "tests/unittest/_torch/moe/test_moe_load_balancer.py", + "tests/unittest/_torch/moe/test_moe_routing.py", "tests/unittest/_torch/multi_gpu/test_allreduce.py", "tests/unittest/_torch/multi_gpu/test_alltoall.py", "tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py", @@ -651,7 +652,6 @@ exclude = [ "tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py", "tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py", "tests/unittest/_torch/multi_gpu/test_mnnvl_memory.py", - "tests/unittest/_torch/multi_gpu/test_moe_a2a.py", "tests/unittest/_torch/multi_gpu/test_user_buffers.py", "tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py", "tests/unittest/_torch/multimodal/test_external_embedding.py", diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index 0da123f6742f..d24e8a606795 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -214,24 +214,6 @@ "E731": 1, "F821": 5 }, - "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py": { - "F821": 2 - }, - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py": { - "E741": 2 - }, - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py": { - "F821": 1 - }, - "tensorrt_llm/_torch/modules/fused_moe/interface.py": { - "E402": 3 - }, - "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py": { - "E712": 2 - }, - "tensorrt_llm/_torch/modules/fused_moe/quantization.py": { - "E731": 1 - }, "tensorrt_llm/_torch/modules/mamba/selective_state_update.py": { "E731": 1 }, @@ -247,6 +229,24 @@ "tensorrt_llm/_torch/modules/mamba/ssd_state_passing.py": { "E731": 1 }, + "tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py": { + "F821": 2 + }, + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py": { + "E741": 2 + }, + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py": { + "F821": 1 + }, + "tensorrt_llm/_torch/moe/fused_moe/interface.py": { + "E402": 3 + }, + "tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py": { + "E712": 2 + }, + "tensorrt_llm/_torch/moe/fused_moe/quantization.py": { + "E731": 1 + }, "tensorrt_llm/_torch/pyexecutor/model_loader.py": { "F811": 1 }, diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 612b315564ca..ba3ffd5dcc2e 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -177,10 +177,8 @@ include = [ "tensorrt_llm/_torch/device_mesh.py", "tensorrt_llm/_torch/distributed/__init__.py", "tensorrt_llm/_torch/distributed/communicator.py", - "tensorrt_llm/_torch/distributed/moe_alltoall.py", "tensorrt_llm/_torch/distributed/ops.py", "tensorrt_llm/_torch/distributed/pg_utils.py", - "tensorrt_llm/_torch/expert_statistic.py", "tensorrt_llm/_torch/flashinfer_utils.py", "tensorrt_llm/_torch/hostfunc.py", "tensorrt_llm/_torch/llm.py", @@ -262,19 +260,6 @@ include = [ "tensorrt_llm/_torch/modules/fla/solve_tril.py", "tensorrt_llm/_torch/modules/fla/utils.py", "tensorrt_llm/_torch/modules/fla/wy_fast.py", - "tensorrt_llm/_torch/modules/fused_moe/__init__.py", - "tensorrt_llm/_torch/modules/fused_moe/create_moe.py", - "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py", - "tensorrt_llm/_torch/modules/fused_moe/interface.py", - "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py", - "tensorrt_llm/_torch/modules/fused_moe/quantization.py", - "tensorrt_llm/_torch/modules/fused_moe/routing.py", "tensorrt_llm/_torch/modules/gated_mlp.py", "tensorrt_llm/_torch/modules/layer_norm.py", "tensorrt_llm/_torch/modules/linear.py", @@ -298,6 +283,21 @@ include = [ "tensorrt_llm/_torch/modules/rotary_embedding.py", "tensorrt_llm/_torch/modules/swiglu.py", "tensorrt_llm/_torch/modules/triton_linear.py", + "tensorrt_llm/_torch/moe/expert_statistic.py", + "tensorrt_llm/_torch/moe/fused_moe/__init__.py", + "tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py", + "tensorrt_llm/_torch/moe/fused_moe/create_moe.py", + "tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py", + "tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py", + "tensorrt_llm/_torch/moe/fused_moe/interface.py", + "tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py", + "tensorrt_llm/_torch/moe/fused_moe/quantization.py", + "tensorrt_llm/_torch/moe/fused_moe/routing.py", "tensorrt_llm/_torch/peft/__init__.py", "tensorrt_llm/_torch/peft/lora/__init__.py", "tensorrt_llm/_torch/peft/lora/config.py", @@ -596,13 +596,14 @@ include = [ "tests/unittest/_torch/modeling/test_modeling_qwen_moe.py", "tests/unittest/_torch/modeling/test_modeling_siglip.py", "tests/unittest/_torch/modeling/test_modeling_vila.py", - "tests/unittest/_torch/modules/test_fused_moe.py", "tests/unittest/_torch/modules/test_group_rmn_norm.py", - "tests/unittest/_torch/modules/test_moe_host_sharer.py", - "tests/unittest/_torch/modules/test_moe_load_balancer.py", - "tests/unittest/_torch/modules/test_moe_routing.py", "tests/unittest/_torch/modules/test_rotary_embedding.py", "tests/unittest/_torch/modules/test_triton_linear.py", + "tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py", + "tests/unittest/_torch/moe/test_fused_moe.py", + "tests/unittest/_torch/moe/test_moe_host_sharer.py", + "tests/unittest/_torch/moe/test_moe_load_balancer.py", + "tests/unittest/_torch/moe/test_moe_routing.py", "tests/unittest/_torch/multi_gpu/test_allreduce.py", "tests/unittest/_torch/multi_gpu/test_alltoall.py", "tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py", @@ -611,7 +612,6 @@ include = [ "tests/unittest/_torch/multi_gpu/test_lowprecision_allreduce.py", "tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py", "tests/unittest/_torch/multi_gpu/test_mnnvl_memory.py", - "tests/unittest/_torch/multi_gpu/test_moe_a2a.py", "tests/unittest/_torch/multi_gpu/test_user_buffers.py", "tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py", "tests/unittest/_torch/multimodal/test_external_embedding.py", diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py index a148675c63a1..9752bf66a3f7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py @@ -28,7 +28,7 @@ from triton_kernels.tensor import FP4, Tensor, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout -from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_triton import ( TritonEPRouter, TritonRoutingData, combine_expert_outputs, diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py index cc1f75e9d863..6b7cade8f1fc 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py @@ -15,7 +15,7 @@ the TMA layout. Mirrors PT's ``MXFP4WeightTRTLLMGenFusedMoEMethod`` -(``tensorrt_llm/_torch/modules/fused_moe/quantization.py:4135``) — PT helpers +(``tensorrt_llm/_torch/moe/fused_moe/quantization.py:4135``) — PT helpers (``maybe_pad_for_mxfp4``, ``trtllmgen_maybe_get_cached_*``, ``_get_weight_alignment``) are imported directly so the algorithm is byte-identical. """ @@ -25,7 +25,7 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.quantization import ( +from tensorrt_llm._torch.moe.fused_moe.quantization import ( _get_weight_alignment, maybe_pad_for_mxfp4, trtllmgen_maybe_get_cached_w2_permute_indices, diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py index 7987e5dba9ed..f7ad6d0fadc0 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py @@ -17,8 +17,8 @@ import torch -from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll -from tensorrt_llm._torch.modules.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.moe.fused_moe.communication.moe_alltoall import MoeAlltoAll +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType from tensorrt_llm._utils import get_sm_version from tensorrt_llm.mapping import Mapping diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py index ef1b04f0f132..1b1b1ba8d088 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py @@ -313,7 +313,7 @@ def make_mxfp4_sharding_load_hook( if moe_tp_size > 1: # Lazy import: TRT-LLM-only helper. Keeps this module importable in # standalone (no tensorrt_llm) so its transforms still register. - from tensorrt_llm._torch.modules.fused_moe.quantization import _get_weight_alignment + from tensorrt_llm._torch.moe.fused_moe.quantization import _get_weight_alignment alignment_tp = _get_weight_alignment( _WEIGHT_ALIGNMENT, _MXFP4_SCALING_VECTOR_SIZE, moe_tp_size, intermediate_size @@ -801,7 +801,7 @@ def _apply_trtllm( # weight-alignment size that the trtllm-gen runner expects. if moe_tp_size > 1: # Lazy import: TRT-LLM-only helper (see module-level note above). - from tensorrt_llm._torch.modules.fused_moe.quantization import _get_weight_alignment + from tensorrt_llm._torch.moe.fused_moe.quantization import _get_weight_alignment alignment_tp = _get_weight_alignment( _WEIGHT_ALIGNMENT, _MXFP4_SCALING_VECTOR_SIZE, moe_tp_size, i_size diff --git a/tensorrt_llm/_torch/custom_ops/__init__.py b/tensorrt_llm/_torch/custom_ops/__init__.py index de9202100657..513aa921891e 100644 --- a/tensorrt_llm/_torch/custom_ops/__init__.py +++ b/tensorrt_llm/_torch/custom_ops/__init__.py @@ -79,9 +79,11 @@ def inplace_slice_copy(dest: torch.Tensor, src: torch.Tensor, dim1_start: int, # importing the module is safe regardless of the result -- it just # logs and leaves ``IS_MEGAMOE_OP_AVAILABLE = False`` on partial # cutlass-dsl installs so callers can fall back via the factory. - from .cute_dsl_megamoe_custom_op import IS_MEGAMOE_OP_AVAILABLE + from ..moe.custom_ops.cute_dsl_megamoe_custom_op import \ + IS_MEGAMOE_OP_AVAILABLE if IS_MEGAMOE_OP_AVAILABLE: - from .cute_dsl_megamoe_custom_op import cute_dsl_megamoe_nvfp4_blackwell + from ..moe.custom_ops.cute_dsl_megamoe_custom_op import \ + cute_dsl_megamoe_nvfp4_blackwell __all__ += ['cute_dsl_megamoe_nvfp4_blackwell'] if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index d593cc5333ca..0063abfa4f1c 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -238,7 +238,7 @@ def make_selected_dummy_topk( def make_routing_method(): # Lazy import to avoid circular import: fused_moe imports from this module. - from tensorrt_llm._torch.modules.fused_moe.routing import ( + from tensorrt_llm._torch.moe.fused_moe.routing import ( ROUTING_METHOD_TYPE_TO_CLASS, RoutingMethodType) # Get routing method diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py index 30c56a896cf5..81f07e6a48fc 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py @@ -4,7 +4,7 @@ Hosts the MegaMoE fused dispatch + FC1 + activation + FC2 + combine CuteDSL kernel. The package is loaded -lazily by :mod:`tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_cute_dsl` +lazily by :mod:`tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_cute_dsl` through :func:`import_kernel` so environments without a CUDA 13 Cutlass DSL runtime can still import the backend file for capability probing. diff --git a/tensorrt_llm/_torch/distributed/__init__.py b/tensorrt_llm/_torch/distributed/__init__.py index 29564f81ed50..833f78ba1fd1 100644 --- a/tensorrt_llm/_torch/distributed/__init__.py +++ b/tensorrt_llm/_torch/distributed/__init__.py @@ -1,7 +1,6 @@ from tensorrt_llm.functional import AllReduceFusionOp from .communicator import Distributed, MPIDist, TorchDist -from .moe_alltoall import MoeAlltoAll from .ops import (AllReduce, AllReduceParams, AllReduceStrategy, HelixAllToAllNative, MiniMaxAllReduceRMS, MoEAllReduce, MoEAllReduceParams, all_to_all_4d, all_to_all_5d, allgather, @@ -24,7 +23,6 @@ "MoEAllReduce", "MoEAllReduceParams", "MiniMaxAllReduceRMS", - "MoeAlltoAll", "TorchDist", "MPIDist", "Distributed", diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py index a27443ad6274..18fb134fac8c 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py @@ -18,7 +18,7 @@ from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm._torch.modules.fused_moe.weight_owner import is_moe_weight_owner +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner @register_mapper("HF", "AfmoeForCausalLM") diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/exaone_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/exaone_moe_weight_mapper.py index 35149cc809c6..d1419f1f9b42 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/exaone_moe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/exaone_moe_weight_mapper.py @@ -2,7 +2,7 @@ from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm._torch.modules.fused_moe.weight_owner import is_moe_weight_owner +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner @register_mapper("HF", "ExaoneMoEForCausalLM") diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py index 6801377f0f35..31219c332539 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py @@ -4,8 +4,7 @@ from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm._torch.modules.fused_moe.weight_owner import \ - is_moe_weight_owner +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner def _unfuse_moe_expert_weights(weights: dict) -> dict: diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py index e0bda5266b69..830b65e6c3aa 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @@ -13,8 +13,8 @@ Qwen3NextHfWeightMapper, ) from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode -from tensorrt_llm._torch.modules.fused_moe.weight_owner import is_moe_weight_owner +from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner from tensorrt_llm.quantization import QuantAlgo _FP8_2D_BLOCK_SIZE = 128 diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_moe_weight_mapper.py index 98bff8881381..f95dd6ef1d9a 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_moe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_moe_weight_mapper.py @@ -6,7 +6,7 @@ from tensorrt_llm._torch.models.checkpoints.hf.qwen3_moe_weight_mapper import Qwen3MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm._torch.modules.fused_moe.weight_owner import is_moe_weight_owner +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner @register_mapper("HF", "Qwen3VLMoeForConditionalGeneration") diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py index e11ad402409e..0d563f9733d6 100644 --- a/tensorrt_llm/_torch/models/modeling_afmoe.py +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -43,11 +43,11 @@ from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import DeepSeekV3MoeRoutingMethod, create_moe -from ..modules.fused_moe.routing import Deepseekv3RoutingImpl from ..modules.gated_mlp import GatedMLP from ..modules.qk_norm_attention import QKNormRoPEAttention from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import DeepSeekV3MoeRoutingMethod, create_moe +from ..moe.fused_moe.routing import Deepseekv3RoutingImpl from ..utils import AuxStreamType from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index b30c713d69f0..f1cc1af3e3d7 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -60,13 +60,12 @@ maybe_slice_for_helix_cp) from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (DeepSeekV3MoeRoutingMethod, - MoEWeightLoadingMode, create_moe, - is_moe_weight_owner) from ..modules.mla import MLA +from ..moe.fused_moe import (DeepSeekV3MoeRoutingMethod, MoEWeightLoadingMode, + create_moe, is_moe_weight_owner) # isort: off -from ..modules.fused_moe.routing import Deepseekv3RoutingImpl +from ..moe.fused_moe.routing import Deepseekv3RoutingImpl # isort: on from ..modules.gated_mlp import GatedMLP from ..modules.linear import (Linear, TensorParallelMode, WeightsLoadingConfig, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 021823f508bf..54a97a70c3d5 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -67,7 +67,13 @@ from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.engram import Engram, EngramConfig, EngramHashProvider -from ..modules.fused_moe import ( +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig +from ..modules.mhc.hyper_connection import HCHead, HCState, mHC +from ..modules.mla import MLA +from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import ( CutlassFusedMoE, DeepSeekV4MoeRoutingMethod, MoEWeightLoadingMode, @@ -77,13 +83,7 @@ is_moe_weight_owner, resolve_moe_cls, ) -from ..modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE -from ..modules.gated_mlp import GatedMLP -from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig -from ..modules.mhc.hyper_connection import HCHead, HCState, mHC -from ..modules.mla import MLA -from ..modules.multi_stream_utils import maybe_execute_in_parallel -from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from ..peft.lora.layer import LoraLayer from ..speculative import SpecMetadata, get_num_extra_kv_tokens from ..utils import ( diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index c85dd570db93..2e53924e7e00 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -755,7 +755,7 @@ def _active_moe_load_balancer(): which ``MoE._init_load_balancer`` consumes ``model_config.moe_load_balancer``. Gating every DSpark EPLB check on it keeps the non-EPLB path untouched. """ - from ..modules.fused_moe.moe_load_balancer import get_moe_load_balancer + from ..moe.fused_moe.moe_load_balancer import get_moe_load_balancer return get_moe_load_balancer() diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index f535b17e53bf..7e4c11060a59 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -26,10 +26,10 @@ from torch import nn from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe -from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode -from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention +from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe +from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 79461f5cd45c..155618f93262 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -29,12 +29,12 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import MoEWeightLoadingMode, create_moe, is_moe_weight_owner from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.qk_norm_attention import QKNormRoPEAttention from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import MoEWeightLoadingMode, create_moe, is_moe_weight_owner from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, Fp4QuantizedTensor from .modeling_deepseekv3 import DeepseekV3Gate, DeepseekV3MTPHead, moe_reduce_add_shared_output diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 02b03e151c95..5808bcf386ba 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -24,9 +24,8 @@ # isort and yapf will fight against each other here, so we disable isort # isort: off -from ..modules.fused_moe import (MoEWeightLoadingMode, - RenormalizeMoeRoutingMethod, TritonFusedMoE, - create_moe, is_moe_weight_owner) +from ..moe.fused_moe import (MoEWeightLoadingMode, RenormalizeMoeRoutingMethod, + TritonFusedMoE, create_moe, is_moe_weight_owner) # isort: on from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm diff --git a/tensorrt_llm/_torch/models/modeling_hunyuan_moe.py b/tensorrt_llm/_torch/models/modeling_hunyuan_moe.py index 95a697cb15d1..bbf6a665daaf 100644 --- a/tensorrt_llm/_torch/models/modeling_hunyuan_moe.py +++ b/tensorrt_llm/_torch/models/modeling_hunyuan_moe.py @@ -17,12 +17,12 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (RenormalizeMoeRoutingMethod, VanillaMoE, - create_moe, is_moe_weight_owner) from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import (RenormalizeMoeRoutingMethod, VanillaMoE, + create_moe, is_moe_weight_owner) from ..utils import AuxStreamType, Fp4QuantizedTensor from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, duplicate_kv_weight, register_auto_model) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 1fd5e8c47c87..dce32f886909 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -110,9 +110,6 @@ from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceParams from ..model_config import ModelConfig -from ..modules.fused_moe import ConfigurableMoE, create_moe -from ..modules.fused_moe.interface import _compute_ep_partition -from ..modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..modules.gated_mlp import GatedMLP from ..modules.kimi_kda import KimiKDALinearAttention from ..modules.linear import Linear as TrtllmLinear @@ -120,6 +117,9 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul +from ..moe.fused_moe import ConfigurableMoE, create_moe +from ..moe.fused_moe.interface import _compute_ep_partition +from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..utils import ActivationType, ActType_TrtllmGen from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -1196,7 +1196,7 @@ def __init__( "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." ) if routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM": - from ..modules.fused_moe.mega_moe import MegaMoEDeepGemm + from ..moe.fused_moe.mega_moe import MegaMoEDeepGemm if not isinstance(self.routed_experts.backend, MegaMoEDeepGemm): raise RuntimeError( @@ -1204,7 +1204,7 @@ def __init__( f"MoE factory selected {type(self.routed_experts.backend).__name__}." ) if routed_moe_model_config.moe_backend == "MEGAMOE_CUTEDSL": - from ..modules.fused_moe.mega_moe import MegaMoECuteDsl + from ..moe.fused_moe.mega_moe import MegaMoECuteDsl # Same guard as MEGAMOE_DEEPGEMM above, and for the same reason: # create_moe silently falls back when a backend declines the diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 9dacdcc834c9..58ec834ff30a 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -34,7 +34,11 @@ from ..modules.attention import _helix_cp_allgather_input, _helix_cp_output_projection from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import ( +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode +from ..modules.qk_norm_attention import QKNormRoPEAttention +from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import ( MiniMaxM2MoeRoutingMethod, MoEImplClass, RoutingMethodType, @@ -42,12 +46,8 @@ create_moe, resolve_moe_cls, ) -from ..modules.fused_moe.interface import MoEWeightLoadingMode -from ..modules.fused_moe.weight_owner import is_moe_weight_owner -from ..modules.gated_mlp import GatedMLP -from ..modules.linear import Linear, TensorParallelMode -from ..modules.qk_norm_attention import QKNormRoPEAttention -from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe.interface import MoEWeightLoadingMode +from ..moe.fused_moe.weight_owner import is_moe_weight_owner from ..speculative import SpecMetadata from ..utils import AuxStreamType from .checkpoints.hf.weight_mapper import HfWeightMapper diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 9edf8163f0ef..0af7a1e598da 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -36,12 +36,12 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (Llama4RenormalizeMoeRoutingMethod, - MoEWeightLoadingMode, create_moe) from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import (Llama4RenormalizeMoeRoutingMethod, + MoEWeightLoadingMode, create_moe) from ..speculative import SpecMetadata from ..utils import AuxStreamType, Fp4QuantizedTensor from .modeling_multimodal_utils import fuse_input_embeds diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py index 9d48dccb9530..08a84efc0103 100644 --- a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py +++ b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py @@ -29,14 +29,14 @@ from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PredefinedAttentionMask from ..model_config import ModelConfig -from ..modules.fused_moe import (BaseMoeRoutingMethod, ConfigurableMoE, - CutlassFusedMoE, FusedMoEQuantScalesFP8, - Llama4RenormalizeMoeRoutingMethod, - MoEWeightLoadingMode) from ..modules.gated_mlp import GatedMLP, swiglu from ..modules.linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig) from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..moe.fused_moe import (BaseMoeRoutingMethod, ConfigurableMoE, + CutlassFusedMoE, FusedMoEQuantScalesFP8, + Llama4RenormalizeMoeRoutingMethod, + MoEWeightLoadingMode) from ..speculative import SpecMetadata from ..utils import AuxStreamType, Fp4QuantizedTensor from .modeling_llama import Llama4Attention, Llama4DecoderLayer, Llama4MoE diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py index 2ca84ee1f2a0..39d20f45fd3f 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py @@ -30,9 +30,9 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import MiniMaxM2MoeRoutingMethod, create_moe from ..modules.linear import Linear, TensorParallelMode, copy_weight, load_weight_shard from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import MiniMaxM2MoeRoutingMethod, create_moe from ..utils import AuxStreamType from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index ec226cce48a2..4e7162461694 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -51,7 +51,6 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import MiniMaxM3MoeRoutingMethod, create_moe from ..modules.gated_mlp import GatedMLP from ..modules.linear import ( Linear, @@ -63,6 +62,7 @@ ) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, create_moe from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..utils import ( ActivationType, diff --git a/tensorrt_llm/_torch/models/modeling_mistral_large3.py b/tensorrt_llm/_torch/models/modeling_mistral_large3.py index f325dbf97539..29f47b1751a8 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral_large3.py +++ b/tensorrt_llm/_torch/models/modeling_mistral_large3.py @@ -7,7 +7,7 @@ from tensorrt_llm._torch.models.checkpoints.mistral.weight_mapper import MistralLarge3WeightMapper from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM from tensorrt_llm._torch.models.modeling_utils import register_auto_model -from tensorrt_llm._torch.modules.fused_moe import RenormalizeNaiveMoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe import RenormalizeNaiveMoeRoutingMethod from tensorrt_llm.quantization.mode import QuantAlgo diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index 31a4d535a644..ce9afa386f74 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/tensorrt_llm/_torch/models/modeling_mixtral.py @@ -13,9 +13,9 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import RenormalizeMoeRoutingMethod, create_moe from ..modules.linear import Linear from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import RenormalizeMoeRoutingMethod, create_moe from ..utils import AuxStreamType from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 0337ea96d9b8..13cf3a0905c8 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -41,16 +41,16 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import MoEWeightLoadingMode, create_moe -from ..modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE -from ..modules.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod, - W4A16NVFP4CutlassFusedMoEMethod) from ..modules.linear import (Linear, NVFP4LinearMethod, TensorParallelMode, W4A16NVFP4LinearMethod) from ..modules.mamba.mamba2_mixer import Mamba2Mixer from ..modules.mlp import MLP from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import MoEWeightLoadingMode, create_moe +from ..moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from ..moe.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod, + W4A16NVFP4CutlassFusedMoEMethod) from ..peft.lora.layer import LoraLayer, LoraModuleType from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, Fp4QuantizedTensor diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 4c7879858e9e..b406e9802ecc 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -13,14 +13,14 @@ from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (BaseMoeRoutingMethod, CutlassFusedMoE, - MoEImplClass, RenormalizeMoeRoutingMethod, - RenormalizeNaiveMoeRoutingMethod, - RoutingMethodType, TRTLLMGenFusedMoE, - create_moe, resolve_moe_cls) -from ..modules.fused_moe.interface import MoEWeightLoadingMode from ..modules.linear import TensorParallelMode from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import (BaseMoeRoutingMethod, CutlassFusedMoE, + MoEImplClass, RenormalizeMoeRoutingMethod, + RenormalizeNaiveMoeRoutingMethod, + RoutingMethodType, TRTLLMGenFusedMoE, create_moe, + resolve_moe_cls) +from ..moe.fused_moe.interface import MoEWeightLoadingMode from ..speculative import SpecMetadata from ..utils import AuxStreamType from .modeling_qwen3 import Qwen3Attention diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 919e1e6e18ce..c664a000480a 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -35,9 +35,9 @@ from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper -from tensorrt_llm._torch.modules.fused_shared_expert import \ - fused_sigmoid_gate_mul_add from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.moe.fused_shared_expert import \ + fused_sigmoid_gate_mul_add from tensorrt_llm._torch.pyexecutor.config_utils import \ get_qwen3_hybrid_layer_types from tensorrt_llm._utils import get_sm_version @@ -50,15 +50,15 @@ from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (BaseMoeRoutingMethod, MoEWeightLoadingMode, - RenormalizeMoeRoutingMethod, - RenormalizeNaiveMoeRoutingMethod, - RoutingMethodType, create_moe) from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.mamba.gdn_mixer import Qwen3NextGatedDeltaNet from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import (BaseMoeRoutingMethod, MoEWeightLoadingMode, + RenormalizeMoeRoutingMethod, + RenormalizeNaiveMoeRoutingMethod, + RoutingMethodType, create_moe) from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, create_lm_head_tp_mapping from .modeling_qwen3 import Qwen3Attention diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py index 53584ef1a50a..551fccda9af7 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen_moe.py @@ -13,10 +13,10 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import DefaultMoeRoutingMethod, create_moe from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import DefaultMoeRoutingMethod, create_moe from ..utils import AuxStreamType from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index d95cf64349c8..eec349016a63 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -19,12 +19,12 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import moe_load_balancer_set_repeated_for_next_layer from ..modules.gated_mlp import GatedMLP from ..modules.linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig) from ..modules.mla import MLA from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import moe_load_balancer_set_repeated_for_next_layer from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..speculative import (SpecMetadata, get_spec_worker, should_use_separate_draft_kv_cache) diff --git a/tensorrt_llm/_torch/models/modeling_step3p7.py b/tensorrt_llm/_torch/models/modeling_step3p7.py index e9e8378c14fa..ee854f18be6e 100644 --- a/tensorrt_llm/_torch/models/modeling_step3p7.py +++ b/tensorrt_llm/_torch/models/modeling_step3p7.py @@ -55,12 +55,12 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding, LMHead -from ..modules.fused_moe import create_moe -from ..modules.fused_moe.interface import MoEWeightLoadingMode -from ..modules.fused_moe.routing import MiniMaxM2MoeRoutingMethod from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import create_moe +from ..moe.fused_moe.interface import MoEWeightLoadingMode +from ..moe.fused_moe.routing import MiniMaxM2MoeRoutingMethod from ..speculative import SpecMetadata from ..utils import AuxStreamType, create_lm_head_tp_mapping from .modeling_speculative import SpecDecOneEngineForCausalLM, _slice_spec_position_ids diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1a3c48eaa955..169ef12df3d6 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -28,10 +28,10 @@ from ..model_config import ModelConfig, TConfig from ..modules.attention import Attention from ..modules.embedding import Embedding, LMHead -from ..modules.fused_moe import MoE, VanillaMoE, is_moe_weight_owner from ..modules.linear import Linear, TensorParallelMode, WeightMode from ..modules.logits_processor import LogitsProcessor from ..modules.rms_norm import RMSNorm +from ..moe.fused_moe import MoE, VanillaMoE, is_moe_weight_owner from ..speculative import SpecMetadata from ._arch_index import (MODEL_ARCH_TO_MODULE, SPEC_MODE_TO_MODULE, is_builtin_zoo_module) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 8c326ee9db64..13bfd8748d65 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -2118,7 +2118,7 @@ def _restore_output(output: torch.Tensor, original_shape, def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): input, original_shape = self._prepare_input(module, input) - from tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4 import \ + from tensorrt_llm._torch.moe.fused_moe.triton_dequant_nvfp4 import \ dequant_nvfp4_2d_triton weight_deq = dequant_nvfp4_2d_triton( module.weight.view(torch.uint8), diff --git a/tensorrt_llm/_torch/modules/triton_linear.py b/tensorrt_llm/_torch/modules/triton_linear.py index e536e9d66ca4..6ba738c33a5a 100644 --- a/tensorrt_llm/_torch/modules/triton_linear.py +++ b/tensorrt_llm/_torch/modules/triton_linear.py @@ -26,8 +26,8 @@ from tensorrt_llm.mapping import Mapping from ...models.modeling_utils import QuantConfig -from .fused_moe.fused_moe_triton import (swizzle_weight_and_scale, - update_weight_stride) +from ..moe.fused_moe.fused_moe_triton import (swizzle_weight_and_scale, + update_weight_stride) from .linear import (Linear, LinearMethodBase, TensorParallelMode, WeightsLoadingConfig, copy_weight, load_weight_shard, load_weights_fused_gate_up_helper, diff --git a/tensorrt_llm/_torch/moe/__init__.py b/tensorrt_llm/_torch/moe/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/moe/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tensorrt_llm/_torch/moe/custom_ops/__init__.py b/tensorrt_llm/_torch/moe/custom_ops/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/moe/custom_ops/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py similarity index 99% rename from tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py rename to tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index d9a661ed0f61..5b6805d9a25b 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -37,11 +37,7 @@ import torch -from tensorrt_llm.logger import logger - -from ..._utils import get_sm_version -from ...math_utils import ceil_div, pad_up -from ..autotuner import ( +from tensorrt_llm._torch.autotuner import ( AutoTuner, ConstraintSpec, DistributedTuningStrategy, @@ -50,14 +46,20 @@ TunableRunner, TuningConfig, ) -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from ..utils import get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2 +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.utils import ( + get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2, +) +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.math_utils import ceil_div, pad_up def _import_megamoe_kernel(): """Lazy import so non-SM100 / no-cutlass-dsl envs can still import this module.""" - from ..cute_dsl_kernels.mega_moe_nvfp4 import import_kernel - from ..cute_dsl_kernels.mega_moe_nvfp4.token_comm import CombineFormat + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4 import import_kernel + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4.token_comm import CombineFormat return import_kernel(), CombineFormat @@ -315,7 +317,7 @@ def validate_megamoe_tactic(tactic: Tuple) -> None: ``ValueError`` with a clear message on failure; the caller (``get_valid_tactics`` / ``forward``) catches and filters. """ - from ..cute_dsl_kernels.mega_moe_nvfp4 import ( + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4 import ( Nvfp4BlockSize, SupportedMmaTileM, SupportedMmaTileN, @@ -492,7 +494,7 @@ def enumerate_megamoe_candidate_tactics(num_tokens: int) -> List[Tuple]: # ``Nvfp4BlockSize`` is probe-only at module load to fail fast # when the kernel package is partially installed; it is consumed # by the lazy import inside ``validate_megamoe_tactic``. - from ..cute_dsl_kernels.mega_moe_nvfp4 import ( + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4 import ( Nvfp4BlockSize, # noqa: F401 SfPaddingBlock, ) @@ -541,7 +543,7 @@ def _cute_launch_helpers(): import cutlass.utils as cutlass_utils from cutlass.cute.typing import AddressSpace - from ..cute_dsl_kernels.mega_moe_nvfp4.sym_buffer import SymBufferHost + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4.sym_buffer import SymBufferHost def to_cute(t, assumed_align=16): ct = cutlass_torch.from_dlpack(t, assumed_align=assumed_align) diff --git a/tensorrt_llm/_torch/expert_statistic.py b/tensorrt_llm/_torch/moe/expert_statistic.py similarity index 100% rename from tensorrt_llm/_torch/expert_statistic.py rename to tensorrt_llm/_torch/moe/expert_statistic.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md rename to tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 669407ce85bd..f7f19f81e62b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -454,6 +454,6 @@ Five backends declare `MoEImplBase` directly — `CutlassFusedMoE`, `TRTLLMGenFu - **Do NOT pick `scheduler_kind` opportunistically** — Use `EXTERNAL_COMM` (default) unless your backend's fused kernel genuinely owns cross-rank exchange via SymmBuffer / equivalent in-kernel collective; `FUSED_COMM` brings hard invariants (no host comm, lockstep launches, no multi-stream overlap) - **Schedulers MUST NOT write `moe.repeat_idx`** — `repeat_idx` is wrapper state advanced once per `forward_impl` regardless of chunk count - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. -- **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. See `tests/unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_quantize_input_zero_tokens` and `test_megamoe_cutedsl_quantize_input_zero_tokens` for the pattern. -- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py` for the 8-tuple tactic pattern (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. +- **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. See `tests/unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_quantize_input_zero_tokens` and `test_megamoe_cutedsl_quantize_input_zero_tokens` for the pattern. +- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the 8-tuple tactic pattern (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. - **Use `distributed_tuning_strategy=DistributedTuningStrategy.MERGE` on a multi-rank `FUSED_COMM` backend's `TuningConfig`** — Every EP rank must converge on the same compiled tactic for every chunk, otherwise the in-kernel NVLink dispatch barrier deadlocks. `PARALLEL` can profile different tactics on different ranks and is unsafe for fused collectives. Reference: `Sm100MegaMoENvfp4Runner.get_tuning_config`. diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/__init__.py rename to tensorrt_llm/_torch/moe/fused_moe/__init__.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/communication/__init__.py similarity index 94% rename from tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/__init__.py index 9858693c5f37..f9416e998985 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py b/tensorrt_llm/_torch/moe/fused_moe/communication/allgather_reducescatter.py similarity index 96% rename from tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/allgather_reducescatter.py index 706de4247a84..91ca38f2bf84 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/allgather_reducescatter.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py b/tensorrt_llm/_torch/moe/fused_moe/communication/base.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/communication/base.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/base.py index 0cfe8d5e5d7b..03944d1cbf5d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/base.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py b/tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep.py index b7d962806ef4..ae71d3f60653 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep.py @@ -26,7 +26,7 @@ import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed +from tensorrt_llm._torch.moe.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed from tensorrt_llm._utils import local_mpi_size from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py b/tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep_low_latency.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep_low_latency.py index a9923d61ce7d..ff65b1f79d4e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep_low_latency.py @@ -26,7 +26,7 @@ import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed +from tensorrt_llm._torch.moe.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed from tensorrt_llm._utils import get_sm_version from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py similarity index 100% rename from tensorrt_llm/_torch/distributed/moe_alltoall.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nccl_ep.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nccl_ep.py similarity index 97% rename from tensorrt_llm/_torch/modules/fused_moe/communication/nccl_ep.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/nccl_ep.py index 53cf15b2469c..1c54713df7f0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nccl_ep.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nccl_ep.py @@ -60,7 +60,7 @@ def __init__( ): super().__init__(mapping) - from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import is_nccl_ep_installed + from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import is_nccl_ep_installed if not is_nccl_ep_installed(): raise RuntimeError("nccl-ep is not installed.") @@ -118,7 +118,7 @@ def __init__( @staticmethod def is_platform_supported() -> bool: - from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import is_nccl_ep_installed + from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import is_nccl_ep_installed return is_nccl_ep_installed() @@ -138,7 +138,7 @@ def _get_context(self): ) from nccl.ep import Layout - from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import get_nccl_ep_context + from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import get_nccl_ep_context self._ctx = get_nccl_ep_context( self.mapping, @@ -389,7 +389,7 @@ def destroy(self): logger.warning(f"Handle.destroy error during destroy: {e}") self._handle = None - from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import release_nccl_ep_context + from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import release_nccl_ep_context if self._ctx is not None: release_nccl_ep_context(self._ctx) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided.py index 61d03b3a973b..6dbac8b289ce 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided_flashinfer.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py rename to tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided_flashinfer.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py rename to tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 438dfa49d56b..bd44358d74fa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -34,15 +34,15 @@ import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.impl_base import MoEImplBase -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.impl_base import MoEImplBase +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEDeployment, MoEEligibility, MoEProblem, MoERejectReason, ) -from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoESchedulerKind, _reject -from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoESchedulerKind, _reject +from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager from tensorrt_llm._torch.utils import ( ActType_TrtllmGen, @@ -313,7 +313,7 @@ def _create_and_sync_backend( / etc., which only become known after the sync. Defer weight creation to the explicit ``backend.create_weights()`` call below. """ - from tensorrt_llm._torch.modules.fused_moe.create_moe import ( + from tensorrt_llm._torch.moe.fused_moe.create_moe import ( create_moe_backend, infer_swiglu_gptoss_style, resolve_moe_cls, diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/create_moe.py rename to tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 899a6722fa8a..66d9871792d0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -4,10 +4,11 @@ import torch +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.utils import (ActivationType, ActType_TrtllmGen, + AuxStreamType) from tensorrt_llm.models.modeling_utils import QuantConfig -from ...model_config import ModelConfig -from ...utils import ActivationType, ActType_TrtllmGen, AuxStreamType from .configurable_moe import ConfigurableMoE from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE diff --git a/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py b/tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py rename to tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py b/tensorrt_llm/_torch/moe/fused_moe/ep_group_health.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py rename to tensorrt_llm/_torch/moe/fused_moe/ep_group_health.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py index 9ee20b480a91..6c32d00df4aa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py @@ -21,13 +21,14 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm import deep_gemm +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.utils import (ActivationType, AuxStreamType, + Fp4QuantizedTensor) from tensorrt_llm._utils import nvtx_range from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo -from ...memory_buffer_utils import get_memory_buffers -from ...model_config import ModelConfig -from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor from .fused_moe_cutlass import CutlassFusedMoE from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py index 5826dfe27aa8..ff3551875d07 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py @@ -25,10 +25,15 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.utils import Fp4QuantizedTensor, is_nvfp4_marlin_supported_sm +from tensorrt_llm._torch.utils import ( + ActivationType, + Fp4QuantizedTensor, + is_gated_activation, + is_nvfp4_marlin_supported_sm, + relu2, +) from tensorrt_llm.models.modeling_utils import QuantAlgo -from ...utils import ActivationType, is_gated_activation, relu2 from .fused_moe_cutlass import CutlassFusedMoE from .impl_contract import ( MoEDeployment, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py index 122ff85280fc..ac7ff83a586d 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py @@ -35,11 +35,12 @@ from triton_kernels.tensor import convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.linear import (TensorParallelMode, + load_weight_shard) +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo -from ...model_config import ModelConfig -from ...utils import ActivationType -from ..linear import TensorParallelMode, load_weight_shard from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, MoERejectReason) from .interface import MoE, _reject diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py index c043857cda05..7d837e9576e0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py @@ -524,7 +524,7 @@ def run_triton_fp8_block_scale_moe( Returns: (T, H) output tensor. """ - from tensorrt_llm._torch.modules.fused_moe.interface import ActivationType + from tensorrt_llm._torch.moe.fused_moe.interface import ActivationType num_tokens, hidden = x.shape num_experts, gate_up_size, _ = w3_w1.shape diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py index 0978e3b45760..2888daae3916 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py @@ -27,9 +27,9 @@ from ...custom_ops.trtllm_gen_custom_ops import \ fp4_block_scale_fake_output_without_finalize from ...model_config import ModelConfig +from ...modules.gated_mlp import GatedMLP from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType, Fp4QuantizedTensor, MxFp8QuantizedTensor) -from ..gated_mlp import GatedMLP from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py rename to tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py index 0bf5e9b7ff85..5c3d4b5375c9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py @@ -5,14 +5,14 @@ import torch from torch import nn +from tensorrt_llm._torch.distributed import allgather, reducescatter +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import ActivationType, is_gated_activation, relu2 from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.utils import fp4_utils -from ...distributed import allgather, reducescatter -from ...model_config import ModelConfig -from ...utils import ActivationType, is_gated_activation, relu2 -from ..gated_mlp import GatedMLP -from ..mlp import MLP from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, MoERejectReason, MoEStaticCapability) from .interface import MoEWeightLoadingMode, _reject @@ -100,7 +100,7 @@ def __init__( layer_idx: Optional[int] = None, activation_type: ActivationType = ActivationType.Swiglu, ): - from ...distributed import AllReduce + from tensorrt_llm._torch.distributed import AllReduce super().__init__() self.routing_method = routing_method @@ -509,7 +509,7 @@ def pack_params(self, experts, module_name: str, weight_name: str): def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): - from ...models.modeling_utils import filter_weights + from tensorrt_llm._torch.models.modeling_utils import filter_weights assert not allow_partial_loading, "Partial loading is not supported for vanilla MoE now" assert self._weights_created diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_base.py b/tensorrt_llm/_torch/moe/fused_moe/impl_base.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/impl_base.py rename to tensorrt_llm/_torch/moe/fused_moe/impl_base.py index debfd5a0ce80..aa8687486542 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_base.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_base.py @@ -31,7 +31,8 @@ from .routing import BaseMoeRoutingMethod if TYPE_CHECKING: - from ...utils import Fp4QuantizedTensor + from tensorrt_llm._torch.utils import Fp4QuantizedTensor + from .impl_identity import MoEImplDescriptor diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_blocks.py b/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/impl_blocks.py rename to tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/impl_contract.py rename to tensorrt_llm/_torch/moe/fused_moe/impl_contract.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_environment.py b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/impl_environment.py rename to tensorrt_llm/_torch/moe/fused_moe/impl_environment.py index 4dbbe1f19bfe..c9811c266caa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_environment.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py @@ -96,7 +96,7 @@ def _probe_megamoe_cutedsl_runtime() -> Tuple[bool, str]: def _probe_megamoe_cutedsl_op() -> Tuple[bool, str]: # Read the module because registration updates this flag after import. - from ...custom_ops import cute_dsl_megamoe_custom_op as megamoe_op + from ..custom_ops import cute_dsl_megamoe_custom_op as megamoe_op if megamoe_op.IS_MEGAMOE_OP_AVAILABLE: return True, "" diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_identity.py b/tensorrt_llm/_torch/moe/fused_moe/impl_identity.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/impl_identity.py rename to tensorrt_llm/_torch/moe/fused_moe/impl_identity.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/interface.py rename to tensorrt_llm/_torch/moe/fused_moe/interface.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py rename to tensorrt_llm/_torch/moe/fused_moe/mega_moe/__init__.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py rename to tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 43869e7cea39..ec526e78173d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -43,7 +43,7 @@ ``peer_rank_ptr_mapper.map(local_addr, 0, off) == local_addr + off`` to a self-mapped pointer (NVSHMEM degenerate convention). * ``ep_size > 1``: regions carved out of the build-time-rendezvous'd - :class:`~tensorrt_llm._torch.custom_ops.cute_dsl_megamoe_custom_op.MegaMoeSymmMemProvider` + :class:`~tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op.MegaMoeSymmMemProvider` symmetric buffer; ``peer_offsets[r] = peer_base[r] - local_base`` enables in-kernel cross-GPU NVLink load/store via ``peer_rank_ptr_mapper.map``. @@ -77,19 +77,20 @@ import torch import torch.distributed as dist +# ``megamoe_activation_sf_bytes_per_row`` lives at module top of the +# custom-op file (NOT inside its ``IS_MEGAMOE_OP_AVAILABLE`` gate), so +# it is always importable. The provider / shared-workspace helpers used +# in ``_alloc_symm_provider`` and ``_ensure_local_staging`` ARE inside +# that gate and therefore stay lazy at the call site. +from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( + megamoe_activation_sf_bytes_per_row, +) from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantAlgo from ....autotuner import AutoTuner - -# ``megamoe_activation_sf_bytes_per_row`` lives at module top of the -# custom-op file (NOT inside its ``IS_MEGAMOE_OP_AVAILABLE`` gate), so -# it is always importable. The provider / shared-workspace helpers used -# in ``_alloc_symm_provider`` and ``_ensure_local_staging`` ARE inside -# that gate and therefore stay lazy at the call site. -from ....custom_ops.cute_dsl_megamoe_custom_op import megamoe_activation_sf_bytes_per_row from ....cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ....model_config import ModelConfig from ....utils import ActivationType, AuxStreamType, Fp4QuantizedTensor @@ -195,7 +196,7 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: try: # mega_moe_cute_dsl.py lives at - # tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py; + # tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py; # four dots take us back to tensorrt_llm._torch where # cute_dsl_kernels.mega_moe_nvfp4 is registered. from ....cute_dsl_kernels.mega_moe_nvfp4 import ( # noqa: F401 @@ -334,7 +335,7 @@ class MegaMoECuteDsl(MoEImplBase): which collapses the kernel's ``peer_rank_ptr_mapper.map(local, 0, off)`` to a self-mapped pointer (NVSHMEM degenerate). * ``ep_size > 1``: regions carved out of - :class:`~tensorrt_llm._torch.custom_ops.cute_dsl_megamoe_custom_op.MegaMoeSymmMemProvider`'s + :class:`~tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op.MegaMoeSymmMemProvider`'s rendezvous'd symmetric buffer. ``create_weights`` performs the (collective) ``torch_symm_mem.rendezvous`` at build time so forward time stays free of cross-rank IPC. ``run_moe`` raises @@ -980,7 +981,7 @@ def _alloc_symm_provider(self): makes only the first layer pay. Raises MegaMoeCuteDslUnavailable when no ProcessGroup is available (hard error for multi-rank). """ - from ....custom_ops.cute_dsl_megamoe_custom_op import ( + from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( get_megamoe_symm_provider, query_megamoe_shared_workspace_bytes, ) @@ -1294,7 +1295,7 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype ) # Shared-workspace probe lives behind ``IS_MEGAMOE_OP_AVAILABLE`` # in cute_dsl_megamoe_custom_op so the import stays lazy. - from ....custom_ops.cute_dsl_megamoe_custom_op import ( + from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( query_megamoe_shared_workspace_bytes, ) @@ -1461,7 +1462,7 @@ def _launch_megamoe_kernel( # op for tuning-mode profiling launches; cleared in the ``finally`` so # a later same-process call for a different shape never sees a stale # scratch (ownership: see ``_profiling_scratch_provider_ref``). - from ....custom_ops.cute_dsl_megamoe_custom_op import ( + from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( set_active_megamoe_profiling_scratch, set_active_megamoe_profiling_scratch_factory, ) @@ -1484,7 +1485,9 @@ def _launch_megamoe_kernel( # must never allocate the multi-GiB scratch. It runs in the # op's profiling pre-hook (where the collective rendezvous is # lockstep-safe) and refreshes this module's weakref. - from ....custom_ops.cute_dsl_megamoe_custom_op import get_megamoe_profiling_scratch + from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( + get_megamoe_profiling_scratch, + ) def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): provider = _get(**_self._profiling_scratch_kwargs) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py rename to tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py rename to tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index b622effe7c47..a25c624a0b41 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -11,14 +11,13 @@ import tensorrt_llm import tensorrt_llm.bindings.internal.runtime as _tbr +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.mmap_utils import advise_tensor_pageout +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.utils import EventType from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from ...distributed import AllReduce -from ...mmap_utils import advise_tensor_pageout -from ...utils import EventType -from ..multi_stream_utils import do_multi_stream - def _tensor_to_weight(t: torch.Tensor) -> _tbr.MoeWeight: """ diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py rename to tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py index 06e6ef3827b6..50874b453995 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py @@ -568,7 +568,7 @@ def __init__(self): from flashinfer.tllm_enums import ActivationType as _flashinfer_activation_type from flashinfer.tllm_enums import RoutingMethodType as _flashinfer_routing_method_type - from ..fused_moe.routing import RoutingMethodType as _trtllmgen_routing_method_type + from .routing import RoutingMethodType as _trtllmgen_routing_method_type self._trtllmgen_routing_method_type = _trtllmgen_routing_method_type diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py similarity index 98% rename from tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py rename to tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index 6d368dc524e1..e374984cc6ba 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -22,12 +22,12 @@ import torch +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.peft.lora.validation import has_moe_lora_targets +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from ...model_config import ModelConfig -from ...peft.lora.validation import has_moe_lora_targets -from ...utils import ActivationType from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from .fused_moe_cutlass import CutlassFusedMoE diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py rename to tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index b0ee6fadec21..c9b93a64945b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -47,7 +47,7 @@ import torch -from tensorrt_llm._torch.expert_statistic import ExpertStatistic +from tensorrt_llm._torch.moe.expert_statistic import ExpertStatistic from tensorrt_llm._torch.utils import EventType, Fp4QuantizedTensor from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator diff --git a/tensorrt_llm/_torch/modules/fused_moe/nccl_ep_utils.py b/tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/nccl_ep_utils.py rename to tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py similarity index 99% rename from tensorrt_llm/_torch/modules/fused_moe/quantization.py rename to tensorrt_llm/_torch/moe/fused_moe/quantization.py index cdea7e6013e6..d3149adf14ec 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -25,6 +25,13 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._torch.mmap_utils import advise_tensor_pageout +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import (TensorParallelMode, + load_weight_shard) +from tensorrt_llm._torch.utils import (ActivationType, + replace_parameter_and_save_metadata, + swizzle_sf, unswizzle_sf) from tensorrt_llm._utils import get_sm_version, is_device_integrated, is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -37,11 +44,6 @@ from tensorrt_llm.quantization.utils.fp8_utils import ( resmooth_to_fp8_e8m0, transform_sf_into_required_layout) -from ...mmap_utils import advise_tensor_pageout -from ...utils import (ActivationType, replace_parameter_and_save_metadata, - swizzle_sf, unswizzle_sf) -from ..gated_mlp import GatedMLP -from ..linear import TensorParallelMode, load_weight_shard from .interface import MoEWeightLoadingMode # The declarations aligns with moe_kernels.h @@ -4616,7 +4618,7 @@ def _build_mega_sf(raw_sf: torch.Tensor, *, num_slots: int, along the last axis to ``flat_size`` so the registered Parameter shape matches. """ - from ...cute_dsl_kernels.mega_moe_nvfp4 import ( + from tensorrt_llm._torch.cute_dsl_kernels.mega_moe_nvfp4 import ( stack_byte_reinterpretable_tensors, to_blocked) device = raw_sf.device diff --git a/tensorrt_llm/_torch/modules/fused_moe/routing.py b/tensorrt_llm/_torch/moe/fused_moe/routing.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/routing.py rename to tensorrt_llm/_torch/moe/fused_moe/routing.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py b/tensorrt_llm/_torch/moe/fused_moe/triton_dequant_nvfp4.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py rename to tensorrt_llm/_torch/moe/fused_moe/triton_dequant_nvfp4.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/weight_owner.py b/tensorrt_llm/_torch/moe/fused_moe/weight_owner.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/weight_owner.py rename to tensorrt_llm/_torch/moe/fused_moe/weight_owner.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/moe/fused_moe/wide_ep_ft.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py rename to tensorrt_llm/_torch/moe/fused_moe/wide_ep_ft.py diff --git a/tensorrt_llm/_torch/modules/fused_shared_expert.py b/tensorrt_llm/_torch/moe/fused_shared_expert.py similarity index 100% rename from tensorrt_llm/_torch/modules/fused_shared_expert.py rename to tensorrt_llm/_torch/moe/fused_shared_expert.py diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 382119e108e4..a9808bf9c8b1 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -15,9 +15,9 @@ from ..attention_backend.trtllm import TrtllmAttentionMetadata from ..distributed import Distributed -from ..expert_statistic import ExpertStatistic from ..memory_buffer_utils import Buffers, get_memory_buffers from ..modules.multi_stream_utils import with_multi_stream +from ..moe.expert_statistic import ExpertStatistic from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8d2a55365026..e0ab193c50fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -56,7 +56,6 @@ from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import Distributed from ..distributed.communicator import init_pp_comm -from ..expert_statistic import ExpertStatistic from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader @@ -66,9 +65,10 @@ _build_request_multimodal_input) from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids from ..models.modeling_utils import DecoderModelForCausalLM -from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer, - MoeLoadBalancerIterContext) from ..modules.mamba.mamba2_metadata import Mamba2Metadata +from ..moe.expert_statistic import ExpertStatistic +from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, + MoeLoadBalancerIterContext) from ..peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager from ..speculative import (SpecMetadata, get_draft_kv_cache_manager, get_num_extra_kv_tokens, get_spec_metadata, @@ -2011,7 +2011,7 @@ def _release_megamoe_profiling_scratch(): # MegaMoE tuning resources are shared across layers, so only the engine # can release them after its full autotune warmup and before graph # capture. Later eviction could invalidate a captured workspace pointer. - from ..custom_ops import cute_dsl_megamoe_custom_op as _megamoe_op + from ..moe.custom_ops import cute_dsl_megamoe_custom_op as _megamoe_op release_megamoe_scratch = getattr(_megamoe_op, "release_megamoe_profiling_scratch", None) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9899ad7426ac..9b111279934d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -41,9 +41,9 @@ from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, get_registered_model_class, timing_metric) -from ..modules.fused_moe.moe_load_balancer import ( - MoeLoadBalancer, maybe_create_moe_load_balancer) from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm +from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, + maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope from .config_utils import (is_hybrid_linear, resolve_hf_torch_dtype, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4b50c73d7abb..ec2b35a8a105 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -57,11 +57,11 @@ from ..disaggregation.executor.transfer_manager import AsyncTransferManager from ..distributed import Distributed from ..distributed.communicator import ReduceOp -from ..expert_statistic import ExpertStatistic from ..models.modeling_multimodal_mixin import \ maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer +from ..moe.expert_statistic import ExpertStatistic from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.speculation_gate import SpeculationGate diff --git a/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md b/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md index 240e5833ae0c..2dbcea388dad 100644 --- a/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md +++ b/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md @@ -110,7 +110,7 @@ when to reach for which knob — material that survives releases. 5. **Developer guides co-locate with code.** Long-form subsystem guidance lives next to the code it governs (precedent: [`ATTENTION_DEVELOPER_GUIDE.md`](../modules/ATTENTION_DEVELOPER_GUIDE.md), - [`MOE_DEVELOPER_GUIDE.md`](../modules/fused_moe/MOE_DEVELOPER_GUIDE.md)). + [`MOE_DEVELOPER_GUIDE.md`](../moe/fused_moe/MOE_DEVELOPER_GUIDE.md)). ## 5. Tests diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index 57ac5b7731f8..bf876058bf4b 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -8,7 +8,7 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import MoeLoadBalancer from tensorrt_llm._torch.utils import get_device_uuid from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import serialization diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py index 17f91cab2f72..08eafe62b978 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py @@ -16,12 +16,12 @@ Qwen3NextSparseMoeBlock, ) from tensorrt_llm._torch.modules.attention import Attention -from tensorrt_llm._torch.modules.fused_moe.interface import MoE from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention from tensorrt_llm._torch.modules.mamba.mamba2_mixer import Mamba2Mixer from tensorrt_llm._torch.modules.mhc.hyper_connection import mHC from tensorrt_llm._torch.modules.mla import MLA +from tensorrt_llm._torch.moe.fused_moe.interface import MoE def mark_ranges(): diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index e2fa43776a11..c4bad7a0724c 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -18,8 +18,8 @@ from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, remove_weights, skip_forward -from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._torch.pyexecutor._util import _mamba_conv_layout_kwargs, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import ( extract_mamba_kv_cache_params, diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 41a8605c21d4..0692fb820713 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1125,8 +1125,8 @@ "unittest/_torch/models/test_qwen3_next_eager_fusion.py": 23.635, "unittest/_torch/models/test_qwen3_next_moe_quant.py": 23.4044400921659, "unittest/_torch/modules": 25.175281045751635, - "unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py": 21.28216552901024, - "unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py": 21.208106382978723, + "unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py": 21.28216552901024, + "unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py": 21.208106382978723, "unittest/_torch/modules/fused_ops/test_fused_qk_norm_rope_gate.py": 22.95076049382716, "unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py": 20.575498666666668, "unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py": 22.0858407960199, @@ -1137,93 +1137,93 @@ "unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py": 595.4572222222223, "unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py": 1098.244823033708, "unittest/_torch/modules/mamba": 748.922897574124, - "unittest/_torch/modules/moe/test_communication_factory.py": 22.613205673758866, - "unittest/_torch/modules/moe/test_kimi_k3_mlp.py": 19.246746987951806, - "unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py": 18.849098591549296, - "unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py": 19.2915, - "unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py": 74.95904438642297, - "unittest/_torch/modules/moe/test_megamoe_streaming_load.py": 19.82725685785536, - "unittest/_torch/modules/moe/test_moe_backend.py -k \"kimi_k3_shape\"": 780.451, - "unittest/_torch/modules/moe/test_moe_backend.py::test_create_moe_forwards_megamoe_activation_options": 21.215417582417583, - "unittest/_torch/modules/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api": 20.14109677419355, - "unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config": 20.977383720930234, - "unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config": 21.372352941176473, - "unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims": 20.38539743589744, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"CUTEDSL and not MEGAMOE_CUTEDSL\"": 93.08139999999999, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"CUTEDSL\"": 184.9769, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"CUTLASS\"": 1113.6267575757577, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 44.5111, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"DEEPGEMM\"": 59.3355, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"DENSEGEMM\"": 656.5955555555555, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"MEGAMOE_CUTEDSL\"": 79.28139999999999, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"MEGAMOE_DEEPGEMM\"": 27.77722222222222, - "unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k \"TRTLLM\"": 131.89833333333334, - "unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe": 966.4413781965008, - "unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout": 20.104456962025317, - "unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts": 32.36005505761844, - "unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm": 218.63211029411764, - "unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary": 190.55733333333333, - "unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant": 107.34633858267718, - "unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing": 23.028669230769232, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTEDSL\"": 92.07561224489795, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8\"": 72.952, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8 and not FP8_BLOCK_SCALES\"": 66.51533333333333, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8_BLOCK_SCALES\"": 57.0945, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and MXFP8 and not W4A8\"": 65.64705102040817, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and NVFP4\"": 66.17890721649485, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A16_MXFP4\"": 66.76862765957448, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_AWQ\"": 58.06602127659575, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_MXFP4_FP8\"": 76.91766326530612, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_MXFP4_MXFP8\"": 96.356, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W8A16\"": 74.23807329842931, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 73.045, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MARLIN\"": 77.85662365591398, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MEGAMOE_CUTEDSL\"": 62.69460869565218, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MEGAMOE_DEEPGEMM\"": 49.36988659793814, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and FP8_BLOCK_SCALES\"": 55.24917391304348, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and NVFP4 and not W4A8\"": 50.80643877551021, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A16_MXFP4\"": 60.62337113402062, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A8_MXFP4_MXFP8\"": 47.92419387755102, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A8_NVFP4_FP8\"": 147.5529693877551, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize]": 60.821448979591835, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize]": 44.66628571428572, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 52.82661224489796, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize]": 78.25108163265307, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 46.35937755102041, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=DeepSeekV3]": 59.55205154639175, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3]": 50.64664285714286, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize]": 25.353724489795916, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 25.412591836734695, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 25.429469387755102, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 45.83761224489796, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 44.14304081632653, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 46.21804081632653, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb": 45.74682978723404, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"MEGAMOE_CUTEDSL\"": 55.514857142857146, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"MEGAMOE_DEEPGEMM\"": 50.229214285714285, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"not None\"": 24.391927835051543, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 51.30888775510204, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 44.745326530612246, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 45.29019387755102, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTEDSL and not MEGAMOE_CUTEDSL\"": 87.83459452736318, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTLASS and None\"": 25.84276, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTLASS and not None\"": 145.1052183006536, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 47.573130102040814, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"DENSEGEMM\"": 446.285817948718, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"MEGAMOE_CUTEDSL\"": 54.1455625, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"MEGAMOE_DEEPGEMM\"": 38.492025445292626, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize]": 62.71827792207792, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize]": 31.517495912806538, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 30.223971428571428, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 27.863624324324324, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize]": 29.289306493506494, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 34.10716486486486, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 34.477355091383814, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 32.093349862258954, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 31.61841038961039, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize]": 56.81982872928177, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 20.939743999999997, - "unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 20.638391, + "unittest/_torch/moe/test_communication_factory.py": 22.613205673758866, + "unittest/_torch/moe/test_kimi_k3_mlp.py": 19.246746987951806, + "unittest/_torch/moe/test_kimi_k3_moe_gate.py": 18.849098591549296, + "unittest/_torch/moe/test_kimi_k3_situ_and_mul.py": 19.2915, + "unittest/_torch/moe/test_kimi_k3_situ_moe.py": 74.95904438642297, + "unittest/_torch/moe/test_megamoe_streaming_load.py": 19.82725685785536, + "unittest/_torch/moe/test_moe_backend.py -k \"kimi_k3_shape\"": 780.451, + "unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_megamoe_activation_options": 21.215417582417583, + "unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api": 20.14109677419355, + "unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config": 20.977383720930234, + "unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config": 21.372352941176473, + "unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims": 20.38539743589744, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"CUTEDSL and not MEGAMOE_CUTEDSL\"": 93.08139999999999, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"CUTEDSL\"": 184.9769, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"CUTLASS\"": 1113.6267575757577, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 44.5111, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"DEEPGEMM\"": 59.3355, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"DENSEGEMM\"": 656.5955555555555, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"MEGAMOE_CUTEDSL\"": 79.28139999999999, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"MEGAMOE_DEEPGEMM\"": 27.77722222222222, + "unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k \"TRTLLM\"": 131.89833333333334, + "unittest/_torch/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe": 966.4413781965008, + "unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout": 20.104456962025317, + "unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts": 32.36005505761844, + "unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm": 218.63211029411764, + "unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary": 190.55733333333333, + "unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant": 107.34633858267718, + "unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing": 23.028669230769232, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTEDSL\"": 92.07561224489795, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8\"": 72.952, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8 and not FP8_BLOCK_SCALES\"": 66.51533333333333, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and FP8_BLOCK_SCALES\"": 57.0945, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and MXFP8 and not W4A8\"": 65.64705102040817, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and NVFP4\"": 66.17890721649485, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A16_MXFP4\"": 66.76862765957448, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_AWQ\"": 58.06602127659575, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_MXFP4_FP8\"": 76.91766326530612, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W4A8_MXFP4_MXFP8\"": 96.356, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"CUTLASS and W8A16\"": 74.23807329842931, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 73.045, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MARLIN\"": 77.85662365591398, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MEGAMOE_CUTEDSL\"": 62.69460869565218, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"MEGAMOE_DEEPGEMM\"": 49.36988659793814, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and FP8_BLOCK_SCALES\"": 55.24917391304348, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and NVFP4 and not W4A8\"": 50.80643877551021, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A16_MXFP4\"": 60.62337113402062, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A8_MXFP4_MXFP8\"": 47.92419387755102, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k \"TRTLLM and W4A8_NVFP4_FP8\"": 147.5529693877551, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize]": 60.821448979591835, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize]": 44.66628571428572, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 52.82661224489796, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize]": 78.25108163265307, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 46.35937755102041, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=DeepSeekV3]": 59.55205154639175, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3]": 50.64664285714286, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize]": 25.353724489795916, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 25.412591836734695, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 25.429469387755102, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 45.83761224489796, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 44.14304081632653, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 46.21804081632653, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb": 45.74682978723404, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"MEGAMOE_CUTEDSL\"": 55.514857142857146, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"MEGAMOE_DEEPGEMM\"": 50.229214285714285, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k \"not None\"": 24.391927835051543, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 51.30888775510204, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 44.745326530612246, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 45.29019387755102, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTEDSL and not MEGAMOE_CUTEDSL\"": 87.83459452736318, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTLASS and None\"": 25.84276, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"CUTLASS and not None\"": 145.1052183006536, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"DEEPGEMM and not MEGAMOE_DEEPGEMM\"": 47.573130102040814, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"DENSEGEMM\"": 446.285817948718, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"MEGAMOE_CUTEDSL\"": 54.1455625, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k \"MEGAMOE_DEEPGEMM\"": 38.492025445292626, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize]": 62.71827792207792, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize]": 31.517495912806538, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 30.223971428571428, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize]": 27.863624324324324, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize]": 29.289306493506494, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 34.10716486486486, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize]": 34.477355091383814, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize]": 32.093349862258954, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize]": 31.61841038961039, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize]": 56.81982872928177, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 20.939743999999997, + "unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize]": 20.638391, "unittest/_torch/modules/test_awq_quantization.py": 18.658090671316476, "unittest/_torch/modules/test_engram.py": 98.23857815442561, "unittest/_torch/modules/test_fp4_num_tokens_slice.py": 20.46693969849246, @@ -1234,8 +1234,8 @@ "unittest/_torch/modules/test_group_rmn_norm.py": 23.72472727272727, "unittest/_torch/modules/test_mhc.py": 27.251203629032258, "unittest/_torch/modules/test_mla_helix.py": 28.243782258064517, - "unittest/_torch/modules/test_moe_load_balancer.py": 20.485004336513445, - "unittest/_torch/modules/test_moe_routing.py": 37.94783109919571, + "unittest/_torch/moe/test_moe_load_balancer.py": 20.485004336513445, + "unittest/_torch/moe/test_moe_routing.py": 37.94783109919571, "unittest/_torch/modules/test_rotary_embedding.py": 19.833860736747532, "unittest/_torch/modules/test_triton_linear.py": 19.619371231696814, "unittest/_torch/modules/test_w4a16_nvfp4_linear.py": 18.517960339943343, diff --git a/tests/integration/defs/.test_durations_aws_dfw b/tests/integration/defs/.test_durations_aws_dfw index ca3362ebd062..1b4de81b6b81 100644 --- a/tests/integration/defs/.test_durations_aws_dfw +++ b/tests/integration/defs/.test_durations_aws_dfw @@ -47,10 +47,10 @@ "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM]": 432.65749452402815, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler]": 491.4517015229212, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False]": 132.3602316979086, - "test_unittests.py::test_unittests_v2[unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm]": 1953.8634579040809, + "test_unittests.py::test_unittests_v2[unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm]": 1953.8634579040809, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True]": 295.83772118296474, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler]": 550.6863905609935, "accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP_MTP]": 759.9382735569961, - "test_unittests.py::test_unittests_v2[unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary]": 418.16275210899767, - "test_unittests.py::test_unittests_v2[unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant]": 211.54122115299106 + "test_unittests.py::test_unittests_v2[unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary]": 418.16275210899767, + "test_unittests.py::test_unittests_v2[unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant]": 211.54122115299106 } diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 9d4ed6926934..8a8c66d6df7f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -2125,7 +2125,7 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, @pytest.mark.skip( reason="CuteDslFusedMoE declines FP8 block scales: it has no FP8 " "block-scale kernel, only a torch.einsum reference. See the ‡ footnote " - "in tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md. " + "in tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md. " "Re-enable against DEEPGEMM / TRTLLM once this checkpoint has an " "owner backend on SM100.") @skip_pre_blackwell diff --git a/tests/integration/defs/test_unittests.py b/tests/integration/defs/test_unittests.py index 9f2b6314735e..daf643da79c5 100644 --- a/tests/integration/defs/test_unittests.py +++ b/tests/integration/defs/test_unittests.py @@ -245,7 +245,9 @@ def test_unittests_v2(llm_root, llm_venv, case: str, output_dir, request): # the overall TRT-LLM test quality metrics. Skip per-sub-test reporting # for MoE by prefixing with "moe-" (CI only collects files starting # with "results" for JUnit reporting). - if case.startswith("unittest/_torch/modules/moe/"): + case_path = case.rstrip("/") + if case_path == "unittest/_torch/moe" or case_path.startswith( + "unittest/_torch/moe/"): output_xml = os.path.join(output_dir, f'moe-results-sub-unittests-{case_fn}.xml') else: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 0e45cb165fce..2ca62d3f15b1 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -34,7 +34,7 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py - - unittest/_torch/modules/moe/test_communication_factory.py + - unittest/_torch/moe/test_communication_factory.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). - unittest/_torch/models/checkpoints diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8d47a2e993e5..dbe2ed4f3de7 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -109,8 +109,8 @@ l0_b200: # GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the # stage's "not cpu_only" markexpr and run on the CPU-Generic stage instead). - unittest/disaggregated/test_kda_mamba_transfer.py - - unittest/_torch/modules/moe/test_kimi_k3_mlp.py - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py + - unittest/_torch/moe/test_kimi_k3_mlp.py + - unittest/_torch/moe/test_kimi_k3_situ_moe.py - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15) # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_mla_helix.py @@ -131,34 +131,34 @@ l0_b200: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules # ------------- MoE components tests --------------- - - unittest/_torch/modules/test_moe_load_balancer.py - - unittest/_torch/modules/test_moe_routing.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py + - unittest/_torch/moe/test_moe_load_balancer.py + - unittest/_torch/moe/test_moe_routing.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- - - unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py::test_fused_route_quant_matches_unfused_chain - - unittest/_torch/modules/moe/test_moe_backend.py::test_kimi_fused_route_quant_skips_prequantized_input - - unittest/_torch/modules/moe/test_moe_backend.py::test_kimi_mxfp8_quantized_tensor_handoff - - unittest/_torch/modules/moe/test_megamoe_streaming_load.py - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout - - unittest/_torch/modules/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api - - unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims - - unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config - - unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config - - unittest/_torch/modules/moe/test_moe_backend.py::test_create_moe_forwards_megamoe_activation_options - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins + - unittest/_torch/moe/test_kimi_k3_moe_gate.py::test_fused_route_quant_matches_unfused_chain + - unittest/_torch/moe/test_moe_backend.py::test_kimi_fused_route_quant_skips_prequantized_input + - unittest/_torch/moe/test_moe_backend.py::test_kimi_mxfp8_quantized_tensor_handoff + - unittest/_torch/moe/test_megamoe_streaming_load.py + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout + - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config + - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_megamoe_activation_options + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins # ------------- MoE: test_single_gpu (by backend) --------------- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTEDSL and not MEGAMOE_CUTEDSL" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "DENSEGEMM" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_CUTEDSL" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTEDSL and not MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "DENSEGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_DEEPGEMM" # ------------- MoE: FlashInfer & TRTLLM symbol collision tests --------------- - unittest/_torch/flashinfer/test_trtllm_flashinfer_symbol_collision.py # --- MoE end @@ -433,12 +433,12 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTEDSL-mtp_nextn=2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL and not MEGAMOE_CUTEDSL" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DENSEGEMM" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "MEGAMOE_CUTEDSL" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL and not MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DENSEGEMM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "MEGAMOE_DEEPGEMM" # ---- non-quantized (quant=None) moved to post-merge ---- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4e8996d43ba1..79f37a2e49d1 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -36,54 +36,54 @@ l0_b300: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules # ------------- MoE components tests --------------- - - unittest/_torch/modules/test_moe_load_balancer.py - - unittest/_torch/modules/test_moe_routing.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py + - unittest/_torch/moe/test_moe_load_balancer.py + - unittest/_torch/moe/test_moe_routing.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- Kimi K3 NVFP4: expert loading + SiTU --------------- # Listed test-by-test on purpose. The rest of this file has 19 failures that # predate this work (verified by running the base revision's copy against the # same source), so registering the whole module would import them into CI as # if they were ours. Those are reported separately; these are green. - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_kimi_k3_expert_ckpt_spec_selection - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_fp8_block_scaled_dequantization - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_materialize_handles_scalar_lazy_safetensors - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads1] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads4] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_w1_w3_order_is_load_bearing - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_rejects_duplicate_slot - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_drains_staging_per_expert - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_forward_runs - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_cutlass_situ_bf16_matches_reference - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[CUTLASS] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[TRTLLM] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[CUTLASS-static_1.0] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[TRTLLM-static_1.0] - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_megamoe_streamed_coverage_survives_per_expert_drain - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_megamoe_overrides_finalize_streamed_expert - - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_mega_format_transform_is_slot_blockwise + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_kimi_k3_expert_ckpt_spec_selection + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_fp8_block_scaled_dequantization + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_materialize_handles_scalar_lazy_safetensors + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads1] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads4] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_w1_w3_order_is_load_bearing + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_rejects_duplicate_slot + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_drains_staging_per_expert + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_forward_runs + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutlass_situ_bf16_matches_reference + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[CUTLASS] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[TRTLLM] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[CUTLASS-static_1.0] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[TRTLLM-static_1.0] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_megamoe_streamed_coverage_survives_per_expert_drain + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_megamoe_overrides_finalize_streamed_expert + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_mega_format_transform_is_slot_blockwise # ------------- MoE: test_moe_backend (by backend) --------------- - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale - - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale + - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins # ------------- MoE: test_single_gpu (specific quant per backend) --------------- # CUTLASS backend: FP8, NVFP4, W4A8_MXFP4_MXFP8, W8A16 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize] # TRTLLM backend: NVFP4, FP8_BLOCK_SCALES, W4A8_NVFP4_FP8, W4A16_MXFP4 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] # CUTEDSL backend: NVFP4 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize] # DEEPGEMM backend: FP8_BLOCK_SCALES - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize] # ---- end MoE tests ---- - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] @@ -107,7 +107,7 @@ l0_b300: backend: pytorch tests: # ---- moved to post-merge (MoE CI optimization) ---- - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 945515cc2ce6..c0071ffaf48e 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -33,6 +33,7 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/checkpoints - unittest/_torch/modules + - unittest/_torch/moe - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d9e79298e1bd..fb4bf93e7a9b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -81,26 +81,26 @@ l0_dgx_b200: tests: # ------------- MoE: test_multi_gpu (by backend x quant) --------------- # --- CUTLASS --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_MXFP4_FP8" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and NVFP4" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_MXFP4_MXFP8" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and MXFP8 and not W4A8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_MXFP4_FP8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and NVFP4" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_MXFP4_MXFP8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and MXFP8 and not W4A8" # --- TRTLLM --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and NVFP4 and not W4A8" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and FP8_BLOCK_SCALES" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A8_NVFP4_FP8" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A16_MXFP4" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A8_MXFP4_MXFP8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and NVFP4 and not W4A8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and FP8_BLOCK_SCALES" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A8_NVFP4_FP8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A16_MXFP4" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and W4A8_MXFP4_MXFP8" # --- CUTEDSL (NVFP4 only) --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTEDSL" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTEDSL" # --- DEEPGEMM (FP8_BLOCK_SCALES only) --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" # --- MEGAMOE_DEEPGEMM (W4A8_MXFP4_MXFP8 only) --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_DEEPGEMM" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_CUTEDSL" # ------------- MoE: test_multi_gpu_eplb --------------- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "not None" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "not None" - condition: ranges: system_gpu_count: @@ -361,4 +361,4 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_mamba_conc_greater_than_mbs[NVIDIA-Nemotron-3-Super-120B-A12B-FP8] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_attn_multi_gpus TIMEOUT (60) # ---- FP8 per-tensor (QDQ) moved to post-merge; block-scale/W4A8 stay in pre ---- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8" diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml index e8368175313c..2e9baee7978a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml @@ -89,30 +89,30 @@ l0_dgx_b300: # ------------- MoE components tests (multi-GPU) --------------- # ------------- MoE: multi-GPU module tests (DEP parallel, per backend per quant) --------------- # CUTLASS backend: FP8, NVFP4, W4A8_MXFP4_MXFP8, W8A16 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=W8A16-routing=Renormalize] # TRTLLM backend: NVFP4, FP8_BLOCK_SCALES, W4A8_NVFP4_FP8, W4A16_MXFP4, W4A8_MXFP4_MXFP8 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=FP8_BLOCK_SCALES-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_NVFP4_FP8-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e256_k6_h4096_i2048-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A8_MXFP4_MXFP8-routing=Renormalize] # CUTEDSL backend: NVFP4 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize] # DEEPGEMM backend: FP8_BLOCK_SCALES - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize] # MEGAMOE_DEEPGEMM backend: W4A8_MXFP4_MXFP8 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3] # MEGAMOE_CUTEDSL backend: NVFP4 - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=DeepSeekV3] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=DeepSeekV3] # ------------- MoE: EPLB (Expert Load Balancing) tests --------------- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_DEEPGEMM" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_CUTEDSL" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_DEEPGEMM" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_CUTEDSL" - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_nixl[DeepSeek-V3-Lite-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-trtllm-auto] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-trtllm-auto] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 799fbc321cbd..43ab57a68f7a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -18,6 +18,7 @@ l0_dgx_h100: - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" - unittest/llmapi/test_additional_model_outputs.py -m "gpu2" - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) + - unittest/_torch/moe/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2[v1] - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2[v2] @@ -163,14 +164,14 @@ l0_dgx_h100: # ------------- MoE: test_multi_gpu (by backend x quant) --------------- # Only CUTLASS and MARLIN backends run on H100 (SM90). TRTLLM/CUTEDSL/DEEPGEMM require SM100+. # --- CUTLASS --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8_BLOCK_SCALES" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16_MXFP4" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_AWQ" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8_BLOCK_SCALES" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16_MXFP4" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_AWQ" # --- MARLIN (SM90-only; focused DEP + ALLGATHER x NVFP4 matrix) --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MARLIN" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MARLIN" # ------------- MoE: test_multi_gpu_eplb --------------- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb - condition: ranges: system_gpu_count: @@ -337,7 +338,7 @@ l0_dgx_h100: orchestrator: mpi tests: # ---- FP8 per-tensor (QDQ) moved to post-merge; block-scale stays in pre ---- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES" - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-python_mamba_cache] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-cpp_mamba_cache] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache] diff --git a/tests/integration/test_lists/test-db/l0_gb10.yml b/tests/integration/test_lists/test-db/l0_gb10.yml index da77f6e22ace..98a2cb83f014 100644 --- a/tests/integration/test_lists/test-db/l0_gb10.yml +++ b/tests/integration/test_lists/test-db/l0_gb10.yml @@ -40,5 +40,5 @@ l0_gb10: # Below cases which are commented out due to they failed on gb10 # - unittest/_torch/modeling -k "modeling_mllama" - unittest/_torch/modeling -k "modeling_out_of_tree" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 7c00cdd722ca..3a474f887b87 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -52,9 +52,9 @@ l0_gb200_multi_gpus: - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform] - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_b_overlap] - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform_contention_opt] - - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm - - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing - - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant + - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm + - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing + - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_fabric_memory[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_bounce[TinyLlama-1.1B-Chat-v1.0] @@ -120,4 +120,4 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-trtllm-auto] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4a16[dp4-auto] # ---- moved to post-merge (MoE CI optimization) ---- - - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary + - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_boundary diff --git a/tests/integration/test_lists/test-db/l0_gb202.yml b/tests/integration/test_lists/test-db/l0_gb202.yml index acd4ef500947..79bd803a1890 100644 --- a/tests/integration/test_lists/test-db/l0_gb202.yml +++ b/tests/integration/test_lists/test-db/l0_gb202.yml @@ -16,8 +16,8 @@ l0_gb202: tests: # ------------- PyTorch tests --------------- - unittest/_torch/modeling -k "modeling_out_of_tree" - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] # - unittest/_torch/modeling -k "modeling_qwen3" # https://nvbugs/5234573 - unittest/_torch/attention/test_attention_mla.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index c4cc4f4de2f1..ddc4bf45cbe3 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -31,8 +31,8 @@ l0_gb300_multi_gpus: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules # ------------- MoE components tests --------------- - - unittest/_torch/modules/test_moe_load_balancer.py - - unittest/_torch/modules/test_moe_routing.py + - unittest/_torch/moe/test_moe_load_balancer.py + - unittest/_torch/moe/test_moe_routing.py - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_gpt_oss" # Kimi K3 speculative decoding: exact kernel-level verify parity (random @@ -50,7 +50,7 @@ l0_gb300_multi_gpus: - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py # Kimi K3 MoE routing shape on the TRTLLM backend (random weights, no # checkpoint; the case moved out of the skipped TestMoeFp4 class). - - unittest/_torch/modules/moe/test_moe_backend.py -k "kimi_k3_shape" + - unittest/_torch/moe/test_moe_backend.py -k "kimi_k3_shape" - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=True] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 800c465dd6e6..b9de32ffbe84 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -34,21 +34,21 @@ l0_h100: - unittest/_torch/modules/tests_lora_modules - unittest/_torch/modules/test_w4a16_nvfp4_linear.py # ------------- MoE components tests --------------- - - unittest/_torch/modules/test_moe_load_balancer.py - - unittest/_torch/modules/test_moe_routing.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py + - unittest/_torch/moe/test_moe_load_balancer.py + - unittest/_torch/moe/test_moe_routing.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py + - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- # ------------- MoE: test_single_gpu (by backend) --------------- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" # ------------- Kimi K3 MoE unit tests --------------- # Any-CUDA (Triton trtllm::situ_and_mul op + torch reference, requires_cuda # only; not SM100/SM103-gated), so they run on Hopper rather than consuming # scarce Blackwell pre-merge capacity. The Blackwell-gated K3 suites live in # l0_b200.yml / l0_gb300_multi_gpus.yml. (TRTLLM-15177) - - unittest/_torch/modules/moe/test_kimi_k3_mlp.py - - unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py - - unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py + - unittest/_torch/moe/test_kimi_k3_mlp.py + - unittest/_torch/moe/test_kimi_k3_moe_gate.py + - unittest/_torch/moe/test_kimi_k3_situ_and_mul.py - unittest/_torch/multimodal -k "not nemotron_nano_v2_vl_fp8" # test_speculative_d2h_parity_real_predictor (both no_stop_token and # stop_token_15) is a HW-agnostic sampler D2H-parity invariant (plain CUDA @@ -450,9 +450,9 @@ l0_h100: - test_e2e.py::test_draft_token_tree_quickstart_advanced_eagle3[Llama-3.1-8b-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct-EAGLE3-LLaMA3.1-Instruct-8B] - test_e2e.py::test_draft_token_tree_quickstart_advanced_eagle3_depth_1_tree[Llama-3.1-8b-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct-EAGLE3-LLaMA3.1-Instruct-8B] # ---- moved to post-merge (MoE CI optimization) ---- - - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" + - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" # ---- non-quantized (quant=None) moved to post-merge ---- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" # Documentation URL validation (CPU-only, no GPU needed) - test_doc.py::test_url_validity - condition: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c9a3a49b00a9..07bedfcbf24e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -333,9 +333,9 @@ unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwe unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_block_reuse_cache_hit[TorchSampler] SKIP (https://nvbugs/6608387) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py::test_nemotron_nano_v2_vl_video_batch_equivalence SKIP (https://nvbugs/6625695) -unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py::TestQwen3LoRA::test_qwen3_fp8_lora SKIP (https://nvbugs/6668777) +unittest/_torch/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819) unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) diff --git a/tests/microbenchmarks/bench_moe/backend.py b/tests/microbenchmarks/bench_moe/backend.py index cc7ce9f990c1..b50c92a06182 100644 --- a/tests/microbenchmarks/bench_moe/backend.py +++ b/tests/microbenchmarks/bench_moe/backend.py @@ -75,7 +75,7 @@ def ensure_cute_dsl_importable_for_benchmark() -> None: do not package CUTLASS DSL; keep that fallback local to the benchmark rather than weakening the production module. """ - module_name = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl" + module_name = "tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl" if module_name in sys.modules: return try: @@ -87,7 +87,7 @@ def ensure_cute_dsl_importable_for_benchmark() -> None: class CuteDslFusedMoE: @classmethod def can_implement(cls, p, d): - from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEEligibility, MoERejectReason, ) @@ -107,31 +107,31 @@ def __init__(self, *_args, **_kwargs): def get_backend_class(backend_type: MoeBackendType): """Import and return the concrete backend class for ``backend_type`` lazily.""" if backend_type == MoeBackendType.CUTLASS: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE return CutlassFusedMoE if backend_type == MoeBackendType.TRTLLM: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE + from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE return TRTLLMGenFusedMoE if backend_type == MoeBackendType.CUTEDSL: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE return CuteDslFusedMoE if backend_type == MoeBackendType.DEEPGEMM: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE + from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE return DeepGemmFusedMoE if backend_type == MoeBackendType.DENSEGEMM: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE + from tensorrt_llm._torch.moe.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE return DenseGEMMFusedMoE if backend_type == MoeBackendType.MEGAMOE_DEEPGEMM: - from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm + from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoEDeepGemm return MegaMoEDeepGemm if backend_type == MoeBackendType.MEGAMOE_CUTEDSL: - from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl + from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoECuteDsl return MegaMoECuteDsl raise ValueError(f"unknown MoE backend {backend_type!r}") diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index dd1c7e3057be..1021349a25be 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -29,7 +29,7 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.interface import ( +from tensorrt_llm._torch.moe.fused_moe.interface import ( MoESchedulerKind, MoEWeightLoadingMode, _compute_ep_partition, @@ -184,7 +184,7 @@ def _situ_kwargs( def _create_moe_for_benchmark(**kwargs): ensure_cute_dsl_importable_for_benchmark() - from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe + from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe return create_moe(**kwargs) diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index 0cfbcf411133..ea5174dbf5f8 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/tests/microbenchmarks/bench_moe/case_runner.py @@ -27,7 +27,7 @@ import torch from tensorrt_llm._torch.autotuner import AutoTuner -from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._utils import mpi_allgather from .build import ( diff --git a/tests/microbenchmarks/bench_moe/cli.py b/tests/microbenchmarks/bench_moe/cli.py index 29edbed1e72e..9d6b854242b6 100644 --- a/tests/microbenchmarks/bench_moe/cli.py +++ b/tests/microbenchmarks/bench_moe/cli.py @@ -25,7 +25,7 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from tensorrt_llm.models.modeling_utils import QuantAlgo from .backend import MoeBackendType diff --git a/tests/microbenchmarks/bench_moe/mapping.py b/tests/microbenchmarks/bench_moe/mapping.py index 1011fc852cce..7e9c0502fd96 100644 --- a/tests/microbenchmarks/bench_moe/mapping.py +++ b/tests/microbenchmarks/bench_moe/mapping.py @@ -39,7 +39,7 @@ from transformers.configuration_utils import PretrainedConfig from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.routing import ( +from tensorrt_llm._torch.moe.fused_moe.routing import ( DeepSeekV3MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, diff --git a/tests/microbenchmarks/bench_moe/quantize.py b/tests/microbenchmarks/bench_moe/quantize.py index ad6e7cefebe5..94e1791aadb2 100644 --- a/tests/microbenchmarks/bench_moe/quantize.py +++ b/tests/microbenchmarks/bench_moe/quantize.py @@ -32,7 +32,7 @@ def _make_moe_test_utils_stub() -> types.ModuleType: - stub = types.ModuleType("_torch.modules.moe.moe_test_utils") + stub = types.ModuleType("_torch.moe.moe_test_utils") def skip_if_insufficient_gpu_memory(*_args, **_kwargs) -> None: return None @@ -43,13 +43,13 @@ def skip_if_insufficient_gpu_memory(*_args, **_kwargs) -> None: def _load_get_test_quant_params(): ensure_cute_dsl_importable_for_benchmark() - dependency_name = "_torch.modules.moe.moe_test_utils" + dependency_name = "_torch.moe.moe_test_utils" original_dependency = sys.modules.get(dependency_name) inserted_stub = original_dependency is None if inserted_stub: sys.modules[dependency_name] = _make_moe_test_utils_stub() try: - quantize_utils = importlib.import_module("_torch.modules.moe.quantize_utils") + quantize_utils = importlib.import_module("_torch.moe.quantize_utils") finally: if inserted_stub: sys.modules.pop(dependency_name, None) diff --git a/tests/microbenchmarks/bench_moe/routing/native_logits.py b/tests/microbenchmarks/bench_moe/routing/native_logits.py index 44a64cce10cd..ff2c57e250a6 100644 --- a/tests/microbenchmarks/bench_moe/routing/native_logits.py +++ b/tests/microbenchmarks/bench_moe/routing/native_logits.py @@ -31,8 +31,8 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE -from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm.tools.layer_wise_benchmarks.runner import make_forward_impl_check from .builders import RoutingPlan diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index e6fa7e74a994..00dd942fa468 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -24,13 +24,13 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEDeployment, MoEProblem, canonical_activation, canonical_quant, ) -from tensorrt_llm._torch.modules.fused_moe.impl_environment import collect_moe_environment +from tensorrt_llm._torch.moe.fused_moe.impl_environment import collect_moe_environment from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm._utils import local_mpi_size from tensorrt_llm.models.modeling_utils import QuantAlgo diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index f0275e72f23b..45c0cb8e7d8c 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -26,7 +26,7 @@ from dataclasses import asdict, dataclass, field from typing import Any, Dict, List, Literal, Optional, Tuple -from tensorrt_llm._torch.modules.fused_moe.routing import ( +from tensorrt_llm._torch.moe.fused_moe.routing import ( DeepSeekV3MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, diff --git a/tests/microbenchmarks/bench_moe/timing/autotune.py b/tests/microbenchmarks/bench_moe/timing/autotune.py index bdcf947f101b..6ffa7caf5ae7 100644 --- a/tests/microbenchmarks/bench_moe/timing/autotune.py +++ b/tests/microbenchmarks/bench_moe/timing/autotune.py @@ -65,7 +65,7 @@ def _run_autotune( tuner.repeat = saved_repeat tuner.stream_delay_micro_secs = saved_stream_delay if tactic_autotune: - from tensorrt_llm._torch.custom_ops.cute_dsl_megamoe_custom_op import ( + from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( release_megamoe_profiling_scratch, ) diff --git a/tests/microbenchmarks/bench_moe_comm.py b/tests/microbenchmarks/bench_moe_comm.py index a43d0cb7020c..b2ba7fcbcd30 100644 --- a/tests/microbenchmarks/bench_moe_comm.py +++ b/tests/microbenchmarks/bench_moe_comm.py @@ -64,9 +64,9 @@ import tensorrt_llm as tllm from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import CutlassFusedMoE, MoE -from tensorrt_llm._torch.modules.fused_moe.communication import Communication, CommunicationFactory -from tensorrt_llm._torch.modules.fused_moe.routing import DefaultMoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe import CutlassFusedMoE, MoE +from tensorrt_llm._torch.moe.fused_moe.communication import Communication, CommunicationFactory +from tensorrt_llm._torch.moe.fused_moe.routing import DefaultMoeRoutingMethod from tensorrt_llm._utils import local_mpi_rank, mpi_allgather, mpi_barrier, mpi_rank, mpi_world_size from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index 987d093aac0d..d11d0c41aeaf 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -258,7 +258,7 @@ def test_dense_mlp_keys_unchanged_by_preprocess(self): def test_is_special_instance_module_for_moe(self): from unittest.mock import MagicMock - from tensorrt_llm._torch.modules.fused_moe.interface import MoE + from tensorrt_llm._torch.moe.fused_moe.interface import MoE mock_moe = MagicMock(spec=MoE) mock_moe.__class__ = MoE diff --git a/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py index 4e420de68f33..8c9a0b4fcef0 100644 --- a/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py @@ -21,8 +21,8 @@ from torch import nn from tensorrt_llm._torch.models.modeling_laguna import LagunaHfWeightMapper -from tensorrt_llm._torch.modules.fused_moe.impl_base import MoEImplBase -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.impl_base import MoEImplBase +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEDeployment, MoEEligibility, MoEProblem, diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 981dc469eca5..fadd1f0b9cc4 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -51,11 +51,11 @@ is_minimax_m3_vl_config, ) from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl_v2 -from tensorrt_llm._torch.modules.fused_moe.routing import ( +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.moe.fused_moe.routing import ( MiniMaxM2MoeRoutingMethod, MiniMaxM3MoeRoutingMethod, ) -from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.llmapi import MiniMaxM3SparseAttentionConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py index fa522e5a71d0..73b36332846d 100644 --- a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py +++ b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py @@ -183,8 +183,8 @@ def _build_moe_block(moe_backend, exclude_modules, layer_idx, *, sm, quant_confi """ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextSparseMoeBlock - from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoEEnvironment - from tensorrt_llm._torch.modules.fused_moe.impl_environment import override_moe_environment + from tensorrt_llm._torch.moe.fused_moe.impl_contract import MoEEnvironment + from tensorrt_llm._torch.moe.fused_moe.impl_environment import override_moe_environment model_config = ModelConfig( pretrained_config=SimpleNamespace( @@ -213,9 +213,7 @@ def _build_moe_block(moe_backend, exclude_modules, layer_idx, *, sm, quant_confi # ``fused_moe/__init__`` re-exports create_moe as a function, which shadows # the submodule of the same name, so ``import ...create_moe as m`` binds the # function. Both forms below resolve the submodule through sys.modules. - from tensorrt_llm._torch.modules.fused_moe.create_moe import ( - resolve_moe_cls as real_resolve_moe_cls, - ) + from tensorrt_llm._torch.moe.fused_moe.create_moe import resolve_moe_cls as real_resolve_moe_cls # Observe the resolution create_moe actually performs instead of # reproducing its argument list here: a copy drifts, and a copy that omits @@ -230,7 +228,7 @@ def _capture(model_config, **kwargs): with ( override_moe_environment(MoEEnvironment(sm=sm)), patch( - "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", + "tensorrt_llm._torch.moe.fused_moe.create_moe.resolve_moe_cls", _capture, ), pytest.raises(_StopBlockInit), diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index c5f4b260160f..bd8b11d12171 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -32,8 +32,8 @@ CompletionFlagReadTimeout, reject_rank_mask_cuda_graph_capture, ) -from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth -from tensorrt_llm._torch.modules.fused_moe.wide_ep_ft import get_wide_ep_ft_options +from tensorrt_llm._torch.moe.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.moe.fused_moe.wide_ep_ft import get_wide_ep_ft_options class FakeCompletionFlagReader: diff --git a/tests/unittest/_torch/modules/test_awq_quantization.py b/tests/unittest/_torch/modules/test_awq_quantization.py index 9ed9967923b0..f93926b8b0d1 100644 --- a/tests/unittest/_torch/modules/test_awq_quantization.py +++ b/tests/unittest/_torch/modules/test_awq_quantization.py @@ -6,9 +6,9 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import DefaultMoeRoutingMethod, create_moe -from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ConfigurableMoE from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.moe.fused_moe import DefaultMoeRoutingMethod, create_moe +from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 2d840a6e188d..2a0e519fbb70 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -281,7 +281,7 @@ def fail_fp4_quantize(*args, **kwargs): with ( patch( - "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + "tensorrt_llm._torch.moe.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", side_effect=fake_dequant, ), patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True), diff --git a/tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py b/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py similarity index 93% rename from tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py rename to tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py index 7c451fad4d11..9a27486aec4d 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py +++ b/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py @@ -19,10 +19,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM -from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ( - _BACKEND_SYNC_ATTRS, - ConfigurableMoE, -) +from tensorrt_llm._torch.moe.fused_moe.configurable_moe import _BACKEND_SYNC_ATTRS, ConfigurableMoE from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -52,11 +49,11 @@ def _create_backend( backend = Mock() with ( patch( - "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", + "tensorrt_llm._torch.moe.fused_moe.create_moe.resolve_moe_cls", return_value=Mock(), ), patch( - "tensorrt_llm._torch.modules.fused_moe.create_moe.create_moe_backend", + "tensorrt_llm._torch.moe.fused_moe.create_moe.create_moe_backend", return_value=backend, ), ): diff --git a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py b/tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py similarity index 98% rename from tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py rename to tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py index 42face75bdc8..f2def3940f58 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py +++ b/tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py @@ -28,7 +28,7 @@ Run as:: - pytest tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py -v + pytest tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py -v """ from __future__ import annotations @@ -39,7 +39,7 @@ import torch import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm.*) -from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import ( fused_expand_group_quant_fp8, masked_index_copy_group_quant_fp8, preprocess_after_permute, diff --git a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py b/tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py similarity index 97% rename from tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py rename to tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py index 7b224cdd44a9..f9282a40336e 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py +++ b/tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py @@ -24,7 +24,7 @@ Run as:: - pytest tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py -v + pytest tests/unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py -v """ from __future__ import annotations @@ -35,7 +35,7 @@ import torch import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm.*) -from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import ( preprocess_after_permute, triton_fused_gather_finalize, triton_masked_index_gather, diff --git a/tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py similarity index 96% rename from tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py rename to tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py index 6db886fe4dbe..e7a6157efa3c 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py +++ b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py @@ -4,7 +4,7 @@ Unit tests for ``run_triton_fp8_block_scale_moe``. Run as a pytest: - pytest tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py -v + pytest tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py -v """ from __future__ import annotations @@ -15,10 +15,10 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton_fp8_block_scale import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_triton_fp8_block_scale import ( run_triton_fp8_block_scale_moe, ) -from tensorrt_llm._torch.modules.fused_moe.interface import ActivationType +from tensorrt_llm._torch.moe.fused_moe.interface import ActivationType from tensorrt_llm._utils import get_sm_version from tests.unittest._torch.helpers import calc_diff, per_block_cast_to_fp8_e8m0 diff --git a/tests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.py b/tests/unittest/_torch/moe/fused_moe/test_triton_mxfp4_swizzle.py similarity index 97% rename from tests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.py rename to tests/unittest/_torch/moe/fused_moe/test_triton_mxfp4_swizzle.py index 76acebbd1b71..40dcca947857 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.py +++ b/tests/unittest/_torch/moe/fused_moe/test_triton_mxfp4_swizzle.py @@ -6,7 +6,7 @@ from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details.layout import HopperMXScaleLayout, HopperMXValueLayout -from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_triton import ( convert_layout_expert_chunked, update_weight_stride, ) diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py b/tests/unittest/_torch/moe/kimi_k3_ref_moe/_moe_kernels.py similarity index 99% rename from tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py rename to tests/unittest/_torch/moe/kimi_k3_ref_moe/_moe_kernels.py index 9f95757f955e..f7d3fba397b0 100644 --- a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py +++ b/tests/unittest/_torch/moe/kimi_k3_ref_moe/_moe_kernels.py @@ -146,7 +146,7 @@ def pack_routed_expert_weights( * ``gemm2_weights_scale``: ``uint8 [E, H_pad128, I_pad // 32]`` — shuffled + block-scale interleaved. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( maybe_pad_for_mxfp4, trtllmgen_maybe_get_cached_w2_permute_indices, trtllmgen_maybe_get_cached_w3_w1_permute_indices, diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py b/tests/unittest/_torch/moe/kimi_k3_ref_moe/_mxfp4.py similarity index 100% rename from tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py rename to tests/unittest/_torch/moe/kimi_k3_ref_moe/_mxfp4.py diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py b/tests/unittest/_torch/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py similarity index 100% rename from tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py rename to tests/unittest/_torch/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py b/tests/unittest/_torch/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py similarity index 99% rename from tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py rename to tests/unittest/_torch/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py index 1e3f58e5938c..f8346d11f189 100644 --- a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py +++ b/tests/unittest/_torch/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py @@ -48,18 +48,18 @@ from typing import Any, List, Optional, Tuple import torch -from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import ( +from _torch.moe.kimi_k3_ref_moe._moe_kernels import ( assert_native_situ_supported, invoke_native_situ_moe, make_situ_alpha_beta, pack_routed_expert_weights, ) -from _torch.modules.moe.kimi_k3_ref_moe._mxfp4 import ( +from _torch.moe.kimi_k3_ref_moe._mxfp4 import ( DEFAULT_GROUP_SIZE, dequantize_last_dim_mxfp4, quantize_last_dim_mxfp4, ) -from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP, NonSituActivation +from _torch.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP, NonSituActivation from torch import nn from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3RMSNorm diff --git a/tests/unittest/_torch/modules/moe/moe_test_utils.py b/tests/unittest/_torch/moe/moe_test_utils.py similarity index 98% rename from tests/unittest/_torch/modules/moe/moe_test_utils.py rename to tests/unittest/_torch/moe/moe_test_utils.py index b6447fbdf90a..7fbd836a2f85 100644 --- a/tests/unittest/_torch/modules/moe/moe_test_utils.py +++ b/tests/unittest/_torch/moe/moe_test_utils.py @@ -26,24 +26,24 @@ import torch from tensorrt_llm._torch.autotuner import AutoTuner -from tensorrt_llm._torch.modules.fused_moe import ( +from tensorrt_llm._torch.moe.fused_moe import ( CuteDslFusedMoE, CutlassFusedMoE, MarlinFusedMoE, TRTLLMGenFusedMoE, ) -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEDeployment, MoEProblem, canonical_quant, ) -from tensorrt_llm._torch.modules.fused_moe.impl_environment import collect_moe_environment -from tensorrt_llm._torch.modules.fused_moe.interface import MoE -from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm -from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_cute_dsl import ( +from tensorrt_llm._torch.moe.fused_moe.impl_environment import collect_moe_environment +from tensorrt_llm._torch.moe.fused_moe.interface import MoE +from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm +from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_cute_dsl import ( is_megamoe_cute_dsl_runtime_available, ) from tensorrt_llm._torch.utils import ActivationType, is_gated_activation @@ -259,7 +259,7 @@ def should_skip_trtllm( # tritondevel container (single GPU is sufficient). # 2. cd tests/unittest && \ # python3 -m pytest -v -s \ - # _torch/modules/moe/test_moe_backend.py::test_moe_backend \ + # _torch/moe/test_moe_backend.py::test_moe_backend \ # -k "TRTLLM" -p no:randomly # 3. Expect: # FAILED ...alpha=1.702_beta=1.0_limit=7.0-e128_k4_h2880_i2880 @@ -339,7 +339,7 @@ def should_skip_trtllm( # - Renormalize / RenormalizeNaive / Default (softmax-based) # See: cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu if routing_method_cls is not None: - from tensorrt_llm._torch.modules.fused_moe import ( + from tensorrt_llm._torch.moe.fused_moe import ( DeepSeekV3MoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, ) @@ -540,7 +540,7 @@ def should_skip_trtllm( # The remaining failure surface is # W4A8_MXFP4_FP8 + TRTLLM-Gen + ``swiglu_gptoss_style=True`` (any top_k, # any model shape, single- or multi-GPU). The element-wise reproducer at - # ``tests/unittest/_torch/modules/moe/test_w4a8_mxfp4_fp8_divergence_repro.py`` + # ``tests/unittest/_torch/moe/test_w4a8_mxfp4_fp8_divergence_repro.py`` # holds the model config / weights / input identical and only toggles # the SwiGLU shape; on ``e60_k4_h2048_i1408 seq=1`` it shows: # * default SwiGLU (alpha=1, beta=0, limit=inf): @@ -674,7 +674,7 @@ def should_skip_cutedsl( # h=2048 fails 8-17%, h=7168 fails 24-35%. Observed: e60(9.4%), # e64(16.5%), e256(34.6%), e384(30.9%) at threshold 3%. if routing_method_cls is not None: - from tensorrt_llm._torch.modules.fused_moe import Llama4RenormalizeMoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe import Llama4RenormalizeMoeRoutingMethod if ( quant_algo == QuantAlgo.NVFP4 @@ -1104,7 +1104,7 @@ def should_skip_routing_method( if routing_method_cls is None or model_config is None: return None - from tensorrt_llm._torch.modules.fused_moe import DeepSeekV3MoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe import DeepSeekV3MoeRoutingMethod # DeepSeekV3 routing: num_experts must be divisible by n_group for the # view operation in noaux_tc (routing.py:298). n_group = max(1, num_experts // 2), @@ -1480,7 +1480,7 @@ def should_skip_to_accelerate_ci( # --- Rule 1: Large e256-class model restrictions --- if is_large_model: if routing_method_cls is not None: - from tensorrt_llm._torch.modules.fused_moe import DeepSeekV3MoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe import DeepSeekV3MoeRoutingMethod if routing_method_cls != DeepSeekV3MoeRoutingMethod: routing_name = routing_method_cls.__name__ @@ -1564,7 +1564,7 @@ def iter_base_test_configs( dtype, backend_type, quant_algo, routing_method_cls, skip_reason, base_test_id) """ if routing_methods is None: - from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe import RenormalizeMoeRoutingMethod routing_methods = [RenormalizeMoeRoutingMethod] diff --git a/tests/unittest/_torch/multi_gpu/test_moe_a2a.py b/tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py similarity index 99% rename from tests/unittest/_torch/multi_gpu/test_moe_a2a.py rename to tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py index 4440c89cc363..b420c8d5c235 100644 --- a/tests/unittest/_torch/multi_gpu/test_moe_a2a.py +++ b/tests/unittest/_torch/moe/multi_gpu/test_moe_a2a.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,7 +24,8 @@ import tensorrt_llm as tllm from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.distributed import MoeAlltoAll +from tensorrt_llm._torch.moe.fused_moe.communication.moe_alltoall import \ + MoeAlltoAll from tensorrt_llm.mapping import Mapping cloudpickle.register_pickle_by_value(sys.modules[__name__]) diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/moe/quantize_utils.py similarity index 99% rename from tests/unittest/_torch/modules/moe/quantize_utils.py rename to tests/unittest/_torch/moe/quantize_utils.py index 7ed5ace32426..b6d7441ef2f4 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/moe/quantize_utils.py @@ -24,15 +24,15 @@ per_block_cast_to_fp8_e8m0, per_token_cast_to_fp8_e8m0, ) -from _torch.modules.moe.moe_test_utils import skip_if_insufficient_gpu_memory +from _torch.moe.moe_test_utils import skip_if_insufficient_gpu_memory from utils.util import check_accuracy from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import BaseMoeRoutingMethod -from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.modules.mxfp8_utils import quant_bf16_to_mxfp8 +from tensorrt_llm._torch.moe.fused_moe import BaseMoeRoutingMethod +from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.utils import ActivationType, is_gated_activation, relu2 from tensorrt_llm._utils import get_sm_version from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/moe/test_communication_factory.py similarity index 97% rename from tests/unittest/_torch/modules/moe/test_communication_factory.py rename to tests/unittest/_torch/moe/test_communication_factory.py index 3919ea233af4..590007be6fcc 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/moe/test_communication_factory.py @@ -21,12 +21,12 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_moe import nccl_ep_utils -from tensorrt_llm._torch.modules.fused_moe.communication import communication_factory -from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( +from tensorrt_llm._torch.moe.fused_moe import nccl_ep_utils +from tensorrt_llm._torch.moe.fused_moe.communication import communication_factory +from tensorrt_llm._torch.moe.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) -from tensorrt_llm._torch.modules.fused_moe.communication.nccl_ep import NcclEP +from tensorrt_llm._torch.moe.fused_moe.communication.nccl_ep import NcclEP def _make_model_config( diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py similarity index 96% rename from tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py rename to tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 1479d4546474..672210ae2c2e 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -22,17 +22,17 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoEDeployment, MoEEnvironment, MoEProblem, MoERejectReason, canonical_quant, ) -from tensorrt_llm._torch.modules.fused_moe.impl_environment import MoEDep -from tensorrt_llm._torch.modules.fused_moe.quantization import ( +from tensorrt_llm._torch.moe.fused_moe.impl_environment import MoEDep +from tensorrt_llm._torch.moe.fused_moe.quantization import ( NVFP4CuteDslB12xFusedMoEMethod, NVFP4CutlassFusedMoEMethod, ) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py b/tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py similarity index 100% rename from tests/unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py rename to tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py diff --git a/tests/unittest/_torch/modules/test_ep_group_health.py b/tests/unittest/_torch/moe/test_ep_group_health.py similarity index 99% rename from tests/unittest/_torch/modules/test_ep_group_health.py rename to tests/unittest/_torch/moe/test_ep_group_health.py index a834c2481a63..1266b5846471 100644 --- a/tests/unittest/_torch/modules/test_ep_group_health.py +++ b/tests/unittest/_torch/moe/test_ep_group_health.py @@ -23,7 +23,7 @@ import pytest -from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( +from tensorrt_llm._torch.moe.fused_moe.ep_group_health import ( EP_MASK_NUM_WORDS, EPGroupHealth, EPGroupHealthSnapshot, diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/moe/test_fused_moe.py similarity index 97% rename from tests/unittest/_torch/modules/test_fused_moe.py rename to tests/unittest/_torch/moe/test_fused_moe.py index 5f9d23016243..db3b348ea02f 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/moe/test_fused_moe.py @@ -23,19 +23,18 @@ from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import \ - CuteDslFusedMoE -from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE +from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode # isort and yapf will fight against each other here, so we disable isort # isort: off -from tensorrt_llm._torch.modules.fused_moe import (BaseMoeRoutingMethod, - CutlassFusedMoE, - TRTLLMGenFusedMoE, - DefaultMoeRoutingMethod, - RenormalizeMoeRoutingMethod, - TritonFusedMoE, create_moe) -from tensorrt_llm._torch.modules.fused_moe.quantization import \ +from tensorrt_llm._torch.moe.fused_moe import (BaseMoeRoutingMethod, + CutlassFusedMoE, + TRTLLMGenFusedMoE, + DefaultMoeRoutingMethod, + RenormalizeMoeRoutingMethod, + TritonFusedMoE, create_moe) +from tensorrt_llm._torch.moe.fused_moe.quantization import \ NVFP4CutlassFusedMoEMethod # isort: on from tensorrt_llm._torch.modules.gated_mlp import GatedMLP @@ -44,7 +43,7 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig # NOTE: Most tests in this file are deprecated and skipped. They are now covered by the -# unified MoE test framework in tests/unittest/_torch/modules/moe/test_moe_backend.py +# unified MoE test framework in tests/unittest/_torch/moe/test_moe_backend.py # and test_moe_module.py. Add new MoE tests there instead of here. cloudpickle.register_pickle_by_value(sys.modules[__name__]) @@ -80,7 +79,7 @@ def round_up(x, alignment): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.parametrize( "moe_backend, dtype, experts, routing_cls, bias", @@ -203,7 +202,7 @@ def test_fused_moe(moe_backend, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="needs 4 GPUs to run this test") @@ -227,7 +226,7 @@ def test_fused_moe_multi_gpu(moe_cls, ep_size): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_hopper @pytest.mark.parametrize( @@ -419,7 +418,7 @@ def set_tensor_value_4(x, num_row, num_cols): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize( @@ -608,7 +607,7 @@ def grouped_gemm(a: torch.Tensor, b: torch.Tensor, a_sf: torch.Tensor, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize( @@ -746,7 +745,7 @@ def test_fused_moe_fp8_blockwise_cute_dsl(dtype, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_no_hopper @pytest.mark.parametrize( @@ -884,7 +883,7 @@ def test_fused_moe_fp8_blockwise_cutlass(dtype, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_no_hopper @pytest.mark.skipif(torch.cuda.device_count() < 4, @@ -921,7 +920,7 @@ def test_fused_moe_fp8_blockwise_cutlass_multi_gpu(ep_size, routing_method, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.skipif(torch.cuda.device_count() < 4, @@ -958,7 +957,7 @@ def test_fused_moe_fp8_blockwise_cute_dsl_multi_gpu(ep_size, routing_method, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -975,7 +974,7 @@ def test_fused_moe_nvfp4(dtype, moe_backend, finalize_fusion): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize("hidden_size, intermediate_size", [(2880, 2880)]) @@ -1244,7 +1243,7 @@ def run_fused_moe_nvfp4(dtype, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize( @@ -1385,7 +1384,7 @@ def test_fused_moe_w4a8_nvfp4_fp8(moe_backend): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_neither_ada_nor_hopper_unittest @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -1664,7 +1663,7 @@ def process_layer( @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @skip_pre_blackwell @pytest.mark.parametrize( @@ -1919,7 +1918,7 @@ def prepare_weights(num_experts: int, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("hidden_size", [768, 2880]) @@ -2251,7 +2250,7 @@ def mxfp4_to_fp32(tensor, scales): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("weight_dtype", [torch.int8]) diff --git a/tests/unittest/_torch/modules/test_fused_shared_expert.py b/tests/unittest/_torch/moe/test_fused_shared_expert.py similarity index 99% rename from tests/unittest/_torch/modules/test_fused_shared_expert.py rename to tests/unittest/_torch/moe/test_fused_shared_expert.py index e9c3ebed5bae..6f2b51ffa3a9 100644 --- a/tests/unittest/_torch/modules/test_fused_shared_expert.py +++ b/tests/unittest/_torch/moe/test_fused_shared_expert.py @@ -23,7 +23,7 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_shared_expert import fused_sigmoid_gate_mul_add +from tensorrt_llm._torch.moe.fused_shared_expert import fused_sigmoid_gate_mul_add pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py similarity index 98% rename from tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py rename to tests/unittest/_torch/moe/test_kimi_k3_mlp.py index e7db11ec0145..3973dd8e7d52 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py @@ -21,7 +21,7 @@ import pytest import torch -from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP +from _torch.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP from torch import nn from tensorrt_llm._torch.modules.gated_mlp import GatedMLP @@ -206,7 +206,7 @@ def test_kimi_k3_shared_expert_parallel_construction( from tensorrt_llm._torch import distributed from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_kimi_linear - from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE + from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py b/tests/unittest/_torch/moe/test_kimi_k3_moe_gate.py similarity index 97% rename from tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py rename to tests/unittest/_torch/moe/test_kimi_k3_moe_gate.py index 3f284709f187..5b8b52bc1d3b 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_moe_gate.py @@ -6,7 +6,7 @@ import pytest import torch -from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3ReferenceMoEGate +from _torch.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3ReferenceMoEGate from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_and_mul.py similarity index 98% rename from tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py rename to tests/unittest/_torch/moe/test_kimi_k3_situ_and_mul.py index 8366e87677c8..a70307e9716c 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_and_mul.py @@ -16,7 +16,7 @@ import pytest import torch -from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP +from _torch.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP from tensorrt_llm._torch.modules.situ import SituAndMul diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py similarity index 98% rename from tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py rename to tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 3a2c979feb09..927fefbff87e 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -36,21 +36,19 @@ import pytest import torch import torch.distributed as dist -from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import ( +from _torch.moe.kimi_k3_ref_moe._moe_kernels import ( is_native_situ_supported, make_situ_alpha_beta, padded_fused_shapes, ) -from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3SparseMoeBlock +from _torch.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3SparseMoeBlock from utils.util import check_accuracy import tensorrt_llm._torch.models.modeling_kimi_linear as modeling_kimi_linear from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3MoERuntime -from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory -from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm import ( - _MEGA_MOE_SYMM_BUFFER_CACHE, -) +from tensorrt_llm._torch.moe.fused_moe.communication import CommunicationFactory +from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import _MEGA_MOE_SYMM_BUFFER_CACHE from tensorrt_llm._torch.utils import ActType_TrtllmGen from tensorrt_llm._utils import get_free_port, get_sm_version from tensorrt_llm.mapping import Mapping @@ -422,7 +420,7 @@ def test_fc1_swap_mutation_breaks_accuracy(): fused, ref = _make_block_pair(config, device) # Rebuild the fused buffers with w1/w3 swapped. - from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import pack_routed_expert_weights + from _torch.moe.kimi_k3_ref_moe._moe_kernels import pack_routed_expert_weights swapped = pack_routed_expert_weights( w1_packed=fused.expert_bank.w3_packed, @@ -451,7 +449,7 @@ def test_swiglu_act_mutation_breaks_accuracy(): config = _K3Config() fused, ref = _make_block_pair(config, device) - import _torch.modules.moe.kimi_k3_ref_moe._moe_kernels as mk + import _torch.moe.kimi_k3_ref_moe._moe_kernels as mk torch.manual_seed(17) x = torch.randn(1, 64, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 @@ -462,7 +460,7 @@ def swiglu_invoke(**kwargs): kwargs["act_type"] = int(ActType_TrtllmGen.SwiGlu) return orig(**kwargs) - from _torch.modules.moe.kimi_k3_ref_moe import kimi_k3_moe_block + from _torch.moe.kimi_k3_ref_moe import kimi_k3_moe_block kimi_k3_moe_block.invoke_native_situ_moe = swiglu_invoke try: @@ -684,7 +682,7 @@ def _make_routed_moe( """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping.""" from transformers.configuration_utils import PretrainedConfig - from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe + from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE, create_moe from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig pretrained_config = PretrainedConfig() @@ -1840,7 +1838,7 @@ def test_megamoe_streamed_coverage_survives_per_expert_drain() -> None: Pure bookkeeping, so it needs no GPU and no EP rendezvous. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4MegaMoECuteDslMethod n_slots = 4 method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) @@ -1913,13 +1911,11 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: """ import inspect - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( NVFP4CutlassFusedMoEMethod, NVFP4FusedMoEMethod, ) - from tensorrt_llm._torch.modules.fused_moe.quantization import ( - NVFP4MegaMoECuteDslMethod as MegaMoE, - ) + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4MegaMoECuteDslMethod as MegaMoE assert not issubclass(MegaMoE, NVFP4CutlassFusedMoEMethod), ( "MegaMoE is a sibling of the Cutlass method, not a child; if this ever " @@ -1970,9 +1966,7 @@ def test_mega_format_transform_is_slot_blockwise() -> None: Includes an uneven final chunk (a prime slot count against the chunk size), which is the case a divisible-only test would miss. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import ( - NVFP4MegaMoECuteDslMethod as MegaMoE, - ) + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4MegaMoECuteDslMethod as MegaMoE method = MegaMoE.__new__(MegaMoE) diff --git a/tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py similarity index 95% rename from tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py rename to tests/unittest/_torch/moe/test_megamoe_streaming_load.py index b61dc8ce204a..87b5d3d9c4df 100644 --- a/tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py +++ b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py @@ -26,8 +26,8 @@ from torch import nn # noqa: E402 if TYPE_CHECKING: - from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode - from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4MegaMoECuteDslMethod pytestmark = pytest.mark.skipif( not torch.cuda.is_available(), @@ -42,8 +42,8 @@ def _load_classes() -> tuple[type[MoEWeightLoadingMode], type[NVFP4MegaMoECuteDslMethod]]: - from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode - from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4MegaMoECuteDslMethod return MoEWeightLoadingMode, NVFP4MegaMoECuteDslMethod diff --git a/tests/unittest/_torch/modules/moe/test_moe_a2a_cft.py b/tests/unittest/_torch/moe/test_moe_a2a_cft.py similarity index 85% rename from tests/unittest/_torch/modules/moe/test_moe_a2a_cft.py rename to tests/unittest/_torch/moe/test_moe_a2a_cft.py index b27ca73c418f..46dbff86eb55 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_a2a_cft.py +++ b/tests/unittest/_torch/moe/test_moe_a2a_cft.py @@ -15,9 +15,13 @@ import pytest -from tensorrt_llm._torch.distributed.moe_alltoall import get_force_cft as get_force_cft_standalone -from tensorrt_llm._torch.distributed.moe_alltoall import should_use_cft as should_use_cft_standalone -from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import ( +from tensorrt_llm._torch.moe.fused_moe.communication.moe_alltoall import ( + get_force_cft as get_force_cft_standalone, +) +from tensorrt_llm._torch.moe.fused_moe.communication.moe_alltoall import ( + should_use_cft as should_use_cft_standalone, +) +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import ( FORCE_CFT_ENV, get_force_cft, should_use_cft, diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py similarity index 97% rename from tests/unittest/_torch/modules/moe/test_moe_backend.py rename to tests/unittest/_torch/moe/test_moe_backend.py index 910506b40a72..c03fe19da7b0 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -24,7 +24,7 @@ import pytest import torch import torch.distributed as dist -from _torch.modules.moe.moe_test_utils import ( +from _torch.moe.moe_test_utils import ( IS_CI_MODE, MoeBackendType, MoeModelConfig, @@ -36,21 +36,21 @@ skip_if_insufficient_gpu_memory, supports_autotuner_capture, ) -from _torch.modules.moe.quantize_utils import get_test_quant_params +from _torch.moe.quantize_utils import get_test_quant_params from transformers.configuration_utils import PretrainedConfig from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import _select_explicit_fallback_tactic from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import ( +from tensorrt_llm._torch.moe.fused_moe import ( DeepSeekV3MoeRoutingMethod, RenormalizeMoeRoutingMethod, ) -from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE -from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( +from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe_backend +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_marlin import MarlinFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( MoECommPlan, MoEDeployment, MoEEnvironment, @@ -58,18 +58,14 @@ MoERejectReason, MoERunContext, ) -from tensorrt_llm._torch.modules.fused_moe.impl_environment import ( +from tensorrt_llm._torch.moe.fused_moe.impl_environment import ( collect_moe_environment, override_moe_environment, ) -from tensorrt_llm._torch.modules.fused_moe.interface import ( - MoE, - MoESchedulerKind, - MoEWeightLoadingMode, -) -from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm -from tensorrt_llm._torch.modules.fused_moe.moe_resolution import impl_class_for, resolve_moe_impl -from tensorrt_llm._torch.modules.fused_moe.quantization import ( +from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoESchedulerKind, MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm +from tensorrt_llm._torch.moe.fused_moe.moe_resolution import impl_class_for, resolve_moe_impl +from tensorrt_llm._torch.moe.fused_moe.quantization import ( FusedMoEMethodBase, NVFP4FusedMoEMethod, NVFP4MarlinFusedMoEMethod, @@ -99,7 +95,7 @@ def test_import_deep_gemm_rejects_pre_situ_mega_moe_api(monkeypatch): import tensorrt_llm - import tensorrt_llm._torch.modules.fused_moe.quantization as quantization_module + import tensorrt_llm._torch.moe.fused_moe.quantization as quantization_module def fp8_fp4_mega_moe(): pass @@ -191,7 +187,7 @@ def test_kimi_fused_route_quant_skips_prequantized_input(monkeypatch) -> None: """An upstream fused down projection owns quantization on this path.""" monkeypatch.delenv("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", raising=False) monkeypatch.setattr( - "tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen.get_sm_version", + "tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen.get_sm_version", MagicMock(side_effect=AssertionError("SM probe must be short-circuited")), ) backend = TRTLLMGenFusedMoE.__new__(TRTLLMGenFusedMoE) @@ -384,7 +380,7 @@ def setup_quant_scales(self, module): def test_configurable_moe_post_load_weights_uses_backend_staged_hooks(): - from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ConfigurableMoE + from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE class HookTestConfigurableMoE(ConfigurableMoE): def quantize_input(self, x, **kwargs): @@ -412,7 +408,7 @@ def run_moe(self, ctx, *, workspace=None): def test_configurable_moe_load_weights_invalidates_wrapper_transform_guard(): - from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ConfigurableMoE + from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE configurable_moe = ConfigurableMoE.__new__(ConfigurableMoE) torch.nn.Module.__init__(configurable_moe) @@ -886,7 +882,7 @@ def test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config(): def test_create_moe_forwards_megamoe_activation_options(monkeypatch): - create_moe_module = importlib.import_module("tensorrt_llm._torch.modules.fused_moe.create_moe") + create_moe_module = importlib.import_module("tensorrt_llm._torch.moe.fused_moe.create_moe") configurable_moe = MagicMock(return_value=object()) monkeypatch.setattr(create_moe_module, "ConfigurableMoE", configurable_moe) monkeypatch.setattr( @@ -943,7 +939,7 @@ def test_megamoe_init_rejects_uneven_num_slots_with_value_error(): def test_megamoe_post_load_rejects_uneven_num_slots_with_value_error(monkeypatch): - import tensorrt_llm._torch.modules.fused_moe.quantization as quantization_module + import tensorrt_llm._torch.moe.fused_moe.quantization as quantization_module class DummyModule: _weights_loaded = True @@ -1456,9 +1452,7 @@ def test_moe_backend( # Clear class-level permute indices cache between parametrized test cases # to work around a B200-specific kernel bug (tactic [32,5] illegal memory access) - from tensorrt_llm._torch.modules.fused_moe.quantization import ( - NVFP4TRTLLMGenFusedMoEBaseMethod, - ) + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEBaseMethod NVFP4TRTLLMGenFusedMoEBaseMethod._cache_permute_indices.clear() diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/moe/test_moe_comm.py similarity index 99% rename from tests/unittest/_torch/modules/moe/test_moe_comm.py rename to tests/unittest/_torch/moe/test_moe_comm.py index 4c3bbb84ce79..7ac9a0796669 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/moe/test_moe_comm.py @@ -63,20 +63,20 @@ import tensorrt_llm as tllm from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( +from tensorrt_llm._torch.moe.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) -from tensorrt_llm._torch.modules.fused_moe.communication.deep_ep import DeepEP -from tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency -from tensorrt_llm._torch.modules.fused_moe.communication.nccl_ep import NcclEP -from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import NVLinkOneSided -from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided -from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided_flashinfer import ( +from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep import DeepEP +from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency +from tensorrt_llm._torch.moe.fused_moe.communication.nccl_ep import NcclEP +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import NVLinkOneSided +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_two_sided_flashinfer import ( NVLinkTwoSidedFlashinfer, ) -from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import deep_ep_installed -from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth -from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import is_nccl_ep_installed +from tensorrt_llm._torch.moe.fused_moe.deep_ep_utils import deep_ep_installed +from tensorrt_llm._torch.moe.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import is_nccl_ep_installed from tensorrt_llm.deep_ep.buffer import Buffer from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/modules/test_moe_host_sharer.py b/tests/unittest/_torch/moe/test_moe_host_sharer.py similarity index 98% rename from tests/unittest/_torch/modules/test_moe_host_sharer.py rename to tests/unittest/_torch/moe/test_moe_host_sharer.py index a725572373e2..428c2632d4ac 100644 --- a/tests/unittest/_torch/modules/test_moe_host_sharer.py +++ b/tests/unittest/_torch/moe/test_moe_host_sharer.py @@ -5,7 +5,7 @@ import torch from mpi4py import MPI -from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import \ +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import \ HostMoeTensorSharer pytestmark = pytest.mark.cpu_only @@ -211,6 +211,6 @@ def tensor_callback(expert_id, tensor_name, tensor_data): if __name__ == "__main__": # This file should be run with mpirun, for example: - # mpirun -np 2 python -m unittest tests/unittest/_torch/modules/test_moe_host_sharer.py + # mpirun -np 2 python -m unittest tests/unittest/_torch/moe/test_moe_host_sharer.py # Run tests using unittest unittest.main() diff --git a/tests/unittest/_torch/modules/test_moe_load_balancer.py b/tests/unittest/_torch/moe/test_moe_load_balancer.py similarity index 99% rename from tests/unittest/_torch/modules/test_moe_load_balancer.py rename to tests/unittest/_torch/moe/test_moe_load_balancer.py index 816c26a0d88d..aa5101e1be76 100644 --- a/tests/unittest/_torch/modules/test_moe_load_balancer.py +++ b/tests/unittest/_torch/moe/test_moe_load_balancer.py @@ -7,7 +7,7 @@ from transformers import PretrainedConfig from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import ( +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, SingleLayerMoeLoadBalancer, get_moe_load_balancer, maybe_create_moe_load_balancer, moe_load_balancer_add_single_layer) diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/moe/test_moe_module.py similarity index 99% rename from tests/unittest/_torch/modules/moe/test_moe_module.py rename to tests/unittest/_torch/moe/test_moe_module.py index 7626ae0ab4fb..0ecd83776600 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/moe/test_moe_module.py @@ -41,7 +41,7 @@ import pytest import torch import torch.distributed as dist -from _torch.modules.moe.moe_test_utils import ( +from _torch.moe.moe_test_utils import ( IS_CI_MODE, MoeBackendType, MoeModelConfig, @@ -63,7 +63,7 @@ skip_trtllm_bf16_on_sm103, supports_autotuner_capture, ) -from _torch.modules.moe.quantize_utils import get_test_quant_params +from _torch.moe.quantize_utils import get_test_quant_params from mpi4py import MPI from mpi4py.futures import MPIPoolExecutor from transformers.configuration_utils import PretrainedConfig @@ -72,7 +72,7 @@ from tensorrt_llm._mnnvl_utils import MnnvlMemory from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import ( +from tensorrt_llm._torch.moe.fused_moe import ( DeepSeekV3MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, @@ -83,13 +83,13 @@ SigmoidRenormMoeRoutingMethod, create_moe, ) -from tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency -from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode -from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import ( +from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency +from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, ) -from tensorrt_llm._torch.modules.fused_moe.quantization import ( +from tensorrt_llm._torch.moe.fused_moe.quantization import ( DeepSeekFP8BlockScalesFusedMoEMethod, DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm, FP8QDQFusedMoEMethod, @@ -792,7 +792,7 @@ def _reset_moe_comm_state(): import gc as _gc try: - from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import ( + from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import ( NVLinkOneSided as _NVOS, ) except ImportError: diff --git a/tests/unittest/_torch/modules/test_moe_routing.py b/tests/unittest/_torch/moe/test_moe_routing.py similarity index 98% rename from tests/unittest/_torch/modules/test_moe_routing.py rename to tests/unittest/_torch/moe/test_moe_routing.py index 1dfcb96e7756..84c50764acbb 100644 --- a/tests/unittest/_torch/modules/test_moe_routing.py +++ b/tests/unittest/_torch/moe/test_moe_routing.py @@ -10,15 +10,15 @@ from transformers.configuration_utils import PretrainedConfig from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe import ( +from tensorrt_llm._torch.moe.fused_moe import ( BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, DeepSeekV4MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, LoadBalancedMoeRoutingMethod, MiniMaxM2MoeRoutingMethod, RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod, create_load_balanced_logits, create_moe) -from tensorrt_llm._torch.modules.fused_moe import routing as moe_routing -from tensorrt_llm._torch.modules.fused_moe.routing import \ +from tensorrt_llm._torch.moe.fused_moe import routing as moe_routing +from tensorrt_llm._torch.moe.fused_moe.routing import \ get_cached_perfect_router_logits from tensorrt_llm._utils import mpi_rank from tensorrt_llm.mapping import Mapping @@ -558,7 +558,7 @@ def test_perfect_router_get_cached_logits_e2e(routing_name, shape): # ----------------------------------------------------------------- # Unified routing-method builder. Covers every RoutingMethodType in -# tensorrt_llm._torch.modules.fused_moe.routing: +# tensorrt_llm._torch.moe.fused_moe.routing: # Default, Renormalize, RenormalizeNaive, Llama4, MiniMax2, DeepSeekV3. # Bias-aware routers seed a non-zero bias so tests can verify the # perfect-router path resets it to zeros. @@ -748,7 +748,7 @@ def _perfect_router_worker(parallel_mode, routing_name, num_tokens, dtype, # Local import inside the worker so module-level _PERFECT_ROUTER_LOGITS_CACHE # is the copy living in the child process. Clear any stale state. - from tensorrt_llm._torch.modules.fused_moe import routing as moe_routing + from tensorrt_llm._torch.moe.fused_moe import routing as moe_routing moe_routing._PERFECT_ROUTER_LOGITS_CACHE.clear() spec = _PERFECT_ROUTER_ROUTING_SPECS[routing_name] @@ -903,7 +903,7 @@ def _reset_perfect_router_comm_state() -> None: """ import gc - from tensorrt_llm._torch.modules.fused_moe import routing as moe_routing + from tensorrt_llm._torch.moe.fused_moe import routing as moe_routing # Wait for any in-flight GPU work first. This teardown runs from a # ``finally`` block, so on the error path a worker may raise while kernels @@ -914,7 +914,7 @@ def _reset_perfect_router_comm_state() -> None: moe_routing._PERFECT_ROUTER_LOGITS_CACHE.clear() try: - from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import \ + from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import \ NVLinkOneSided as _NVOS except ImportError: _NVOS = None diff --git a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py index 22a0f61b66b1..fd5603b41bbc 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py +++ b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py @@ -24,7 +24,7 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.peft.lora.cuda_graph_lora_params import CudaGraphLoraParams from tensorrt_llm._torch.peft.lora.layer import LoraModuleType diff --git a/tests/unittest/_torch/peft/test_moe_lora_extract.py b/tests/unittest/_torch/peft/test_moe_lora_extract.py index 26abb8bb9f90..a359cafa06fc 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_extract.py +++ b/tests/unittest/_torch/peft/test_moe_lora_extract.py @@ -21,7 +21,7 @@ # These imports are pure-Python; skip cleanly if the package layout changes. -fused_moe_cutlass = pytest.importorskip("tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass") +fused_moe_cutlass = pytest.importorskip("tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass") lora_layer = pytest.importorskip("tensorrt_llm._torch.peft.lora.layer") CutlassFusedMoE = fused_moe_cutlass.CutlassFusedMoE diff --git a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py index 69663873bf75..99c0b7bff096 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py +++ b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py @@ -30,7 +30,7 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.peft.lora.moe_layout import make_per_expert_lora, reference_swiglu_moe_lora from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig diff --git a/tests/unittest/_torch/peft/test_moe_lora_model_path.py b/tests/unittest/_torch/peft/test_moe_lora_model_path.py index 391af7e5be06..bed5faff8c44 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_model_path.py +++ b/tests/unittest/_torch/peft/test_moe_lora_model_path.py @@ -23,10 +23,10 @@ import torch from tensorrt_llm._torch.models.modeling_mixtral import MixtralMoE -from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ConfigurableMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE -from tensorrt_llm._torch.modules.fused_moe.moe_scheduler import ExternalCommMoEScheduler +from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE +from tensorrt_llm._torch.moe.fused_moe.moe_scheduler import ExternalCommMoEScheduler from tensorrt_llm._torch.peft.lora.layer import LoraModuleType pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/test_mnnvl_utils.py b/tests/unittest/_torch/test_mnnvl_utils.py index eeed206b6840..7cbd30c14e67 100644 --- a/tests/unittest/_torch/test_mnnvl_utils.py +++ b/tests/unittest/_torch/test_mnnvl_utils.py @@ -18,7 +18,7 @@ import pynvml from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency +from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency def setup_function() -> None: @@ -205,15 +205,15 @@ def test_supports_mnnvl_accepts_full_fabric( @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", True, ) @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.get_sm_version", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.get_sm_version", return_value=90, ) @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", return_value=0, ) @patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=True) @@ -224,15 +224,15 @@ def test_deep_ep_low_latency_rejects_split_topology( @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", True, ) @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.get_sm_version", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.get_sm_version", return_value=100, ) @patch( - "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", + "tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", return_value=0, ) @patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=False) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 7952f4fa9fee..118ac8963a48 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -4,8 +4,8 @@ from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import cute_dsl_nvfp4_grouped_gemm_ref -from tensorrt_llm._torch.modules.fused_moe.quantization import interleave_linear_and_gate +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import cute_dsl_nvfp4_grouped_gemm_ref +from tensorrt_llm._torch.moe.fused_moe.quantization import interleave_linear_and_gate from tensorrt_llm._torch.utils import ( ActivationType, is_gated_activation, diff --git a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py index 4fb55ed93d76..d285fc199dac 100644 --- a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py +++ b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py @@ -28,7 +28,7 @@ import torch.nn.functional as F from utils.util import skip_pre_blackwell -from tensorrt_llm._torch.modules.fused_moe.quantization import interleave_linear_and_gate +from tensorrt_llm._torch.moe.fused_moe.quantization import interleave_linear_and_gate from tensorrt_llm._torch.utils import swizzle_sf, unswizzle_sf from tensorrt_llm.math_utils import pad_up diff --git a/tests/unittest/_torch/thop/parallel/test_moe_densegemm.py b/tests/unittest/_torch/thop/parallel/test_moe_densegemm.py index c5d917096ab8..783c6db7fb6e 100644 --- a/tests/unittest/_torch/thop/parallel/test_moe_densegemm.py +++ b/tests/unittest/_torch/thop/parallel/test_moe_densegemm.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch.modules.fused_moe.quantization import interleave_linear_and_gate +from tensorrt_llm._torch.moe.fused_moe.quantization import interleave_linear_and_gate from tensorrt_llm._torch.utils import swizzle_sf, unswizzle_sf from tensorrt_llm._utils import get_sm_version from tensorrt_llm.math_utils import pad_up diff --git a/tests/unittest/_torch/thop/serial/test_moe.py b/tests/unittest/_torch/thop/serial/test_moe.py index e50c3557432e..d10e05d6fc10 100644 --- a/tests/unittest/_torch/thop/serial/test_moe.py +++ b/tests/unittest/_torch/thop/serial/test_moe.py @@ -24,7 +24,7 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # NOTE: Some tests in this file are deprecated and skipped. They are now covered by the -# unified MoE test framework in tests/unittest/_torch/modules/moe/test_moe_backend.py +# unified MoE test framework in tests/unittest/_torch/moe/test_moe_backend.py # and test_moe_module.py. Add new MoE tests there instead of here. from enum import Enum @@ -32,7 +32,7 @@ from utils.util import getSMVersion from tensorrt_llm._torch.autotuner import AutoTuner, autotune -from tensorrt_llm._torch.modules.fused_moe import RoutingMethodType +from tensorrt_llm._torch.moe.fused_moe import RoutingMethodType from tensorrt_llm._torch.utils import next_positive_power_of_2 from tensorrt_llm.quantization.utils.fp4_utils import ( reorder_rows_for_gated_act_gemm, shuffle_matrix_a, shuffle_matrix_sf_a) @@ -82,7 +82,7 @@ def test_situ_runner_has_valid_configs(num_tokens, get_tactics): tactic list here means the whole path is unreachable. Numerics and launch evidence for both families live in - tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py. + tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py. """ assert get_tactics( num_tokens), f"No valid SiTu tactic for num_tokens={num_tokens}" @@ -932,7 +932,7 @@ def are_groups_valid(top_k_groups, n_groups): @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.skipif( getSMVersion() < 100 or getSMVersion() >= 110, @@ -1070,7 +1070,7 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.skipif( getSMVersion() < 100 or getSMVersion() >= 110, @@ -2009,7 +2009,7 @@ def run_moe_fp8_fp4_test(self, num_tokens: int, hidden_size: int, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.skipif( getSMVersion() < 100 or getSMVersion() >= 110, @@ -2238,7 +2238,7 @@ def test_moe_fp8_per_tensor_scale(num_tokens, hidden_size, intermediate_size, @pytest.mark.skip( reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." + "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." ) @pytest.mark.skipif( getSMVersion() != 100, diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_ad_moe_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_ad_moe_op.py index 1bcd958afcc1..c05b1d2e24cd 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_ad_moe_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_ad_moe_op.py @@ -20,7 +20,7 @@ import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale -from tensorrt_llm._torch.modules.fused_moe import MoE # noqa: F401 +from tensorrt_llm._torch.moe.fused_moe import MoE # noqa: F401 def setup_moe_test(dtype, num_experts): diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py index 5c63535aed15..8e4d01ad9dc4 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py @@ -148,7 +148,7 @@ def test_fc1_bias_is_shuffled_with_same_row_permutation_as_fc1_weights(): padded (not shuffled), causing the trtllm-gen kernel to add the wrong bias to each output row. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( trtllmgen_maybe_get_cached_w3_w1_permute_indices, ) @@ -203,7 +203,7 @@ def test_fc1_bias_is_shuffled_with_same_row_permutation_as_fc1_weights(): def test_fc2_bias_is_shuffled_with_same_row_permutation_as_fc2_weights(): """Regression: fc2 bias must follow the (non-gated) TMA row permute used by w2.""" - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( trtllmgen_maybe_get_cached_w2_permute_indices, ) @@ -254,7 +254,7 @@ def test_prep_against_pt_reference_loader_byte_identical(): load_expert_w2_weight_scale_mxfp4}`` is the gold standard the AD prep helper must mirror. Any divergence here is a kernel-layout bug. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( _get_weight_alignment, maybe_pad_for_mxfp4, trtllmgen_maybe_get_cached_w2_permute_indices, diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_trtllm_gen_diag.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_trtllm_gen_diag.py index b980ca954071..196806291bbc 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_trtllm_gen_diag.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_trtllm_gen_diag.py @@ -88,7 +88,7 @@ def test_compare_module_vs_ad_weight_processing(): """Compare processed weights/scales from AD fusion with module processing.""" import torch - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( trtllmgen_maybe_get_cached_w3_w1_permute_indices, ) from tensorrt_llm.quantization.utils.fp4_utils import ( @@ -192,11 +192,11 @@ def test_direct_kernel_call(): """Call fp4_block_scale_moe_runner directly with module-processed weights.""" import torch - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( trtllmgen_maybe_get_cached_w2_permute_indices, trtllmgen_maybe_get_cached_w3_w1_permute_indices, ) - from tensorrt_llm._torch.modules.fused_moe.routing import RoutingMethodType + from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType device = "cuda" dtype = torch.bfloat16 @@ -393,7 +393,7 @@ def test_verify_actual_ad_tensors(): """Extract and verify the actual tensors the AD fusion produces.""" import torch - from tensorrt_llm._torch.modules.fused_moe.quantization import ( + from tensorrt_llm._torch.moe.fused_moe.quantization import ( trtllmgen_maybe_get_cached_w2_permute_indices, trtllmgen_maybe_get_cached_w3_w1_permute_indices, )