trtllm-moe-develop

name	trtllm-moe-develop
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, 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 asks whether a MoE design follows the current architecture or whether a MoE refactor is reasonable.
license	Apache-2.0
metadata	{"author":"NVIDIA Corporation"}

TensorRT-LLM MoE Code Quality

Use this skill to keep MoE changes aligned with the current TensorRT-LLM MoE architecture. Favor module roles, API boundaries, and testability over local style cleanup.

Required Context

Before proposing or editing MoE code, read:

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

Also inspect these files when the area is relevant:

• Forward execution/chunking: inspect moe_scheduler.py, configurable_moe.py, interface.py, backend run_moe/quantize_input paths, and communication code.
• MegaMoE/fused communication: inspect moe_scheduler.py, mega_moe/, configurable_moe.py, quantization.py, and communication code.
• Communication: tensorrt_llm/_torch/modules/fused_moe/communication/base.py and communication_factory.py.
• Quantization and weights: tensorrt_llm/_torch/modules/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 quantize_utils.py when adding backend, quantization, skip, or parameter coverage.

For module-specific work, read references/moe-canonical-code-examples.md after the guide and load only the relevant section. Each design gate or review should cite at least one concrete code example with file:line evidence.

Working With MOE_DEVELOPER_GUIDE.md

Treat MOE_DEVELOPER_GUIDE.md as the in-repo source of truth for MoE architecture. Treat this skill as the agent workflow layer that tells Codex how to apply that source of truth while designing, editing, or reviewing code.

Use the guide this way:

• Start from the guide sections that match the requested change: Architecture, File Map, Backend Capability Matrix, execution-flow/EPLB constraints, Canonical Examples, and Anti-Patterns.
• Use guide content to fill the design gate: owner boundary, main API, reference pattern, and test plan.
• Do not duplicate fast-changing matrices or backend support tables in this skill; prefer the guide as the current reference.
• If a code change adds a backend, quantization method, communication strategy, fused-communication behavior, EPLB behavior, or test convention, check whether the guide also needs an update.
• If guide and code disagree, inspect code and tests, mention the mismatch, and either update the guide as part of the change or report it as follow-up.

Guide-update checklist:

• File map changed: update File Map.
• Backend or quant support changed: update Backend Capability Matrix.
• New backend/communication/forward-execution pattern: update Canonical Examples.
• New forbidden pattern or ownership rule: update Anti-Patterns.
• Test convention changed: update Tests.

Core Principle

Preserve these owner boundaries:

• ConfigurableMoE is the assembler/orchestrator. It wires backend, communication, EPLB, weight lifecycle delegation, and shared wrapper bookkeeping.
• Backends declare capabilities, run MoE computation, and own the MoE module's weight lifecycle boundary. They expose and implement create_weights, load_weights, post_load_weights, process_weights_after_loading, and pre_reload_weights as needed, select any FusedMoEMethod, and make those hooks compatible with ConfigurableMoE deferred weight creation and reload flows. Backends may delegate quantization-specific tensor layout, loading, post-load transforms, and scale setup to a quantization method, but backend lifecycle hooks remain the public owner of weight handling. New ConfigurableMoE-compatible backends should expose quantize_input and run_moe, not forward or forward_impl, unless the user explicitly asks for legacy standalone behavior. For an active ConfigurableMoE-compatible backend, run_moe must be a concrete implementation, not an empty stub, and ConfigurableMoE or its scheduler should call backend.run_moe(...) as the compute entrypoint. Backend-specific alternatives such as run_with_prequant are acceptable only as private helpers called from run_moe, not as public wrapper/scheduler targets that bypass the common contract. The current MoE interface still covers both legacy standalone MoE modules and newer backends; treat legacy forward methods as transitional until a dedicated MoEBackend interface exists. Backends should not become orchestration or external-communication state machines.
• Quantization methods are backend-selected implementation helpers for quantization-specific weight tensor layout, loading details, post-load transforms, scale setup, and EPLB fix-up registration. They do not replace backend ownership of the weight lifecycle API.
• Communication strategies own external cross-rank dispatch/combine.
• MoEScheduler owns forward-time policy: padding/truncation, chunking, dispatch/quantize ordering, EPLB hook ordering, zero-token chunk behavior, external-vs-fused communication workflow, and backend run_moe invocation. ConfigurableMoE constructs the scheduler from backend.scheduler_kind and delegates to it; schedulers may read wrapper state and call wrapper helpers but must not own lifecycle, weight loading, DWDP record, repeat_idx advancement, or communication lifetime. The only sanctioned scheduler mutation of moe.comm is through determine_communication_method fallback.
• Shared test helpers own backend/quantization matrices and skip logic. Updating one test file while leaving moe_test_utils.py or quantize_utils.py stale is usually incomplete.
• Tests should exercise the boundary that changed: backend, module, communication, routing, EPLB, or multi-GPU behavior.

A refactor is good only if it keeps these roles clearer than before.

Module Blocks

ConfigurableMoE: Assembler

Role:

• Compose backend, communication strategy, EPLB, and wrapper-level lifecycle.
• Keep forward_impl focused on wrapper-level work: resolve output dtype, delegate execution, record DWDP, advance repeat_idx once.
• Own backend construction/sync and validation, not backend-specific forward policy.

Main APIs / references:

• configurable_moe.py: ConfigurableMoE.__init__, backend construction, communication strategy creation/bypass, scheduler construction, forward_impl, validate_backend.
• MOE_DEVELOPER_GUIDE.md: ConfigurableMoE orchestrator and file map.

Checklist:

• New behavior still leaves ConfigurableMoE as an assembler.
• No new backend-specific fast path in forward_impl unless it is a temporary compatibility bridge with a clear follow-up.
• forward_impl or an extracted scheduler should invoke backend computation through backend.run_moe(...); direct calls to backend-specific compute entrypoints such as run_with_prequant are red flags unless the change is explicitly a short-lived adapter and run_moe remains the real implementation.
• Shared wrapper state such as repeat_idx, DWDP record, backend attr sync, and communication lifetime stays in one place.
• Scheduler creation happens after backend, communication, chunking streams, validation, and optional DWDP setup are initialized, because schedulers read that wrapper state.
• forward_impl should not accumulate chunking, routing, communication, EPLB, or fused-kernel branches; that policy belongs in MoEScheduler.

MoE Scheduler: Forward Execution Strategy

Role:

• Own per-forward execution policy for ConfigurableMoE: padding/truncation, chunking, dispatch ordering, adaptive pre/post quant dispatch, EPLB wait/stat update/route/CPU-stage hook ordering, zero-token chunk behavior, and backend run_moe invocation.
• Select external-vs-fused communication behavior through MoESchedulerKind, not through wrapper isinstance checks.
• Read wrapper state and call wrapper helpers, but do not own module lifecycle, backend construction, weight lifecycle, communication object lifetime, DWDP record, or repeat_idx advancement.

Main APIs / references:

• moe_scheduler.py: MoEScheduler, ExternalCommMoEScheduler, FusedCommMoEScheduler, create_moe_scheduler.
• interface.py: MoESchedulerKind and backend scheduler_kind.
• configurable_moe.py: scheduler construction and thin forward_impl delegation.
• communication/base.py: supports_post_quant_dispatch, prepare_dispatch, dispatch, and combine contracts used by ExternalCommMoEScheduler.

Checklist:

• New forward policy goes in moe_scheduler.py, not in ConfigurableMoE or a backend, unless it is truly backend-local compute inside run_moe.
• ExternalCommMoEScheduler owns host-side dispatch/combine, communication fallback, optional multi-stream chunk overlap, padding/truncation, and external communication EPLB statistic paths.
• FusedCommMoEScheduler owns fused-kernel lockstep: ADP stripping, per-rank-consistent chunk count, zero-token launches, no external dispatch/combine, and ignore_allreduce=False EPLB statistic update.
• Schedulers call backend.quantize_input(...) and backend.run_moe(...); they must not call backend-specific alternate compute helpers that bypass run_moe.
• Schedulers must not advance repeat_idx, run DWDP record/prefetch, create or destroy communication strategies, or call weight lifecycle hooks.
• If backend-specific kwargs are needed, keep them centralized and narrow inside scheduler helper code, with comments explaining why the common run_moe contract is insufficient for that backend.
• Add/update module-level tests for changed scheduler behavior, especially chunking, zero-token chunks, DP padding/truncation, EPLB hook order, and fused-communication lockstep.

MoE Backend

Role:

• Pure MoE computation and backend-specific capability/config validation.
• Own module-level weight handling and lifecycle delegation through create_weights, load_weights, post_load_weights, process_weights_after_loading, and pre_reload_weights.
• Own quantize_input and run_moe shape/kernel contracts. run_moe must launch the backend compute path for every active ConfigurableMoE-compatible backend. Do not leave it as NotImplementedError while the wrapper calls an alternate method such as run_with_prequant.
• Do not implement forward or forward_impl for new ConfigurableMoE-compatible backends unless the user explicitly requests legacy standalone behavior; if required, document why the normal backend contract is insufficient.
• Declare whether the backend's cross-rank exchange is external to the kernel or fused inside the kernel.

Main APIs / references:

• interface.py: MoE, scheduler_kind, can_implement, _supports_load_balancer, validate_configurable_moe when present, and weight lifecycle hooks (create_weights, load_weights, post_load_weights, process_weights_after_loading, pre_reload_weights).
• fused_moe_cutlass.py: reference backend using external communication.
• mega_moe/: reference area for a fused-communication backend.
• create_moe.py: backend selection and fallback path.

Checklist:

• can_implement() returns clear (False, reason) for unsupported quant, dtype, shape, or hardware.
• Backend weight lifecycle hooks are implemented or explicitly rejected with a narrow error; create_weights() is safe under ConfigurableMoE deferred weight creation, load_weights() honors or rejects allow_partial_loading, and post_load_weights() / process_weights_after_loading() / pre_reload_weights() keep transformed weights and reload metadata coherent.
• The backend selects and stores the quantization method before delegating layout-specific weight registration/loading/transforms; callers should not need to reach into quantization.py directly.
• run_moe is implemented and is the method reached by ConfigurableMoE or the scheduler. If a helper like run_with_prequant exists for performance or naming compatibility, it is called from run_moe, not directly from wrapper policy code.
• Cross-rank exchange ownership is explicit via scheduler_kind and not hidden behind wrapper isinstance checks. Backends with kernel-fused exchange declare MoESchedulerKind.FUSED_COMM; normal backends use EXTERNAL_COMM.
• Backend-specific wrapper constraints go in a validation hook or an equivalent narrow contract, not in scattered forward branches.
• Weight handling remains backend API scope even when the actual tensor layout is implemented by a FusedMoEMethod.
• Do not add external host communication logic to a backend, except for a true fused-communication backend whose kernel owns the exchange.
• New backend tests belong in test_moe_backend.py.
• Existing legacy forward methods can be read for compatibility context, but they are not the default pattern for new backend work.

Imported Kernel ABI Checklist

When importing or wrapping an upstream kernel, derive the TRT-LLM adapter contract from the lowest-level kernel consumer. Comments, docs, design notes, and parameter names are useful hints, but they are not proof of the runtime ABI.

• Derive weight shape and layout from the kernel entrypoint, make_layout, TMA, MMA/GEMM transforms, and stride usage. Record required tensor shape, stride, physical storage layout, and boundary view layout.
• Derive alpha and scale semantics from kernel consumption points. Trace where alpha, norm constants, block scales, activation scales, and weight scales are loaded and multiplied before deciding how upper layers compute or pack them. Treat weight bytes, block scales/SF, and global alpha/norm constants as separate contracts.
• Design the upper-layer adapter from the kernel ABI upward. Map each kernel input/output to an adapter responsibility: storage tensor, view/transposition, dtype reinterpretation, padding, scale packing, workspace ownership, synchronization, and output reduction. Validate parity with upstream invocation dumps, not just final output.

Quantization And Weights

Role:

• Weight handling is backend scope at the module/API boundary: the backend exposes the lifecycle hooks, owns when they are called, and is accountable for reload/EPLB consistency.
• Quantization-specific tensor creation, loading details, post-load transforms, quant scales, and EPLB weight fix-ups should live in quantization.py as a backend-selected FusedMoEMethod implementation when they are specific to a quantization layout.
• When adding new weight handling, first look for a reusable existing quant method or base class before creating a new one, then make the backend select and invoke it through the lifecycle hooks.

Main APIs / references:

• quantization.py: FusedMoEMethodBase, create_weights, load_weights, post_load_weights, setup_quant_scales, eplb_support_status, supports_online_eplb, need_load_shared_weights.
• Existing quant methods in quantization.py are the reference patterns.

Checklist:

• New backend weight handling is surfaced through backend lifecycle hooks; new quantization-specific tensor layouts are represented by a backend-selected quantization method, not ad hoc caller or wrapper code.
• Existing quant method/layout is reused when the tensor layout and scale semantics match.
• create_weights() registers module parameters with the correct slot, expert, hidden, intermediate, and scale layout.
• load_weights() handles supported loading modes and rejects unsupported ones clearly. Preserve the EPLB split: common MoE FC weights/biases (w3_w1_weight, w2_weight, and bias tensors when present) use the shared FusedMoEMethodBase.load_weights() / post_load_weights() path, where need_load_shared_weights(module) gates CPU shared staging and registration.
• Quantization methods add only their quantization-specific EPLB registrations for scales, alphas, transformed weights, or layout-specific views that are not covered by the base FC weight path. Those extra tensors must also be gated by need_load_shared_weights(module) before loading, transforming, or registering shared copies. If a specialized method cannot reuse the base FC path because its raw parameter layout is incompatible, the design must call out that exception and preserve equivalent base semantics explicitly.
• post_load_weights() performs transforms, shared-weight setup, and scale setup in the quantization method only for tensors outside the base FC path; base FC weight registration should still flow through the base class whenever possible.
• setup_quant_scales() is updated when a quant mode exposes scales consumed by backend, communication, or forward-execution paths.
• EPLB support status is explicit: SUPPORTED, NOT_SUPPORTED, or NOT_VERIFIED.

EPLB

Role:

• EPLB is cross-cutting. A correct change may need updates in interface, quantization, forward execution, communication, and tests.
• Do not treat EPLB as only a backend flag.

Main APIs / references:

• interface.py: _supports_load_balancer, _add_raw_shared_weights_for_unmap, _using_load_balancer, _using_dynamic_load_balancer, validation hooks.
• quantization.py: eplb_support_status, need_load_shared_weights, register_all_parameter_slot_and_to_fix_weight_fns, setup_quant_scales, post_load_weights.
• Current forward-execution code: statistic update, route, ignore_allreduce, per-chunk first/last hook ordering.
• test_moe_module.py: EPLB params and generate_*_eplb_test_params.

Checklist:

• Backend reports whether load balancing is supported.
• Quantization method declares online EPLB status.
• EPLB weight registration is split into two layers:
  1. Common MoE FC weights/biases are handled by FusedMoEMethodBase using need_load_shared_weights(module) in its shared-load/register flow.
  2. Quantization-specific scales, alphas, transformed weights, or layout views are handled by the concrete quantization method and must add their own need_load_shared_weights(module) gated shared-load/register logic.
• Shared quant-specific tensors needed by EPLB are registered in the quantization method, including any fix-up functions for transformed weights.
• Forward execution collects routing statistics and chooses ignore_allreduce correctly for the communication path.
• EPLB hook order is preserved around routing, run_moe, and CPU weight migration.
• num_slots, num_experts, ep_size, and slot-vs-expert IDs are not mixed.
• Add or update concrete EPLB tests in test_moe_module.py, including the backend/comm/quant combination that changed.

CPU shared-staging buffer family (EPLB migration)

Dynamic EPLB needs host-resident copies of per-expert tensors so that MoeLoadBalancer can migrate experts between ranks via host shared memory. Each per-expert nn.Parameter on the module has a parallel CPU staging buffer; all of them are passed to register_all_parameter_slot_and_to_fix_weight_fns once loading finishes. Any new per-expert Parameter MUST add its own staging buffer and migration hook, or the shared-load path will either write out of bounds or silently corrupt routed slots (NVBug 6130334 / PR #13856).

Full family in the NVFP4 path (quantization.py):

GPU nn.Parameter on module	CPU shared staging buffer	Sized by
w3_w1_weight (packed FP4)	module.local_shared_w3_w1_tensors	len(local_shared_load_expert_ids)
w2_weight (packed FP4)	module.local_shared_w2_tensors	same
w3_w1_bias / w2_bias (if bias=True)	module.local_shared_w3_w1_bias_tensors / module.local_shared_w2_bias_tensors	same
w3_w1_weight_scale / w2_weight_scale (block scales)	module.local_shared_w3_w1_scale_tensors / module.local_shared_w2_scale_tensors	same
fc31_alpha / fc2_alpha (per-expert fp32 scalar)	shared_fc31_alpha / shared_fc2_alpha (local variables in process_weights_after_loading)	num_shared = len(tmp_shared_weight_scale_2)
fc31_weight_scale_2 / fc2_weight_scale_2 (per-expert fp32 scalar, gated by force_dynamic_quantization)	shared_fc31_weight_scale_2 / shared_fc2_weight_scale_2 (local variables)	same

Key index-space distinction:

• expert_size_per_partition = num_slots / ep_size is the routed-slot count on this rank; sizes the on-GPU module Parameters.
• num_shared = len(local_shared_load_expert_ids) = num_experts / shared_size, where shared_size = shared_mpi_comm.Get_size() is the same-node MPI rank count (from MPI_COMM_TYPE_SHARED split); sizes the CPU staging buffers.
• On multi-node setups shared_size < ep_size is legal and makes num_shared > expert_size_per_partition. Any code that writes into a routed-sized Parameter using a staging-space index will go out of bounds.
• On single-node setups shared_size == ep_size is enforced by the assert shared_size == local_size in MoeLoadBalancer._setup_mpi_comm, so single-node unit tests cannot exercise the num_shared > expert_size_per_partition failure mode through parameter tuning alone. A regression test for staging-index correctness must either (a) invoke the reconcile/migration function directly with a crafted staging dict, or (b) run on a real multi-node Slurm environment.

Naming convention quirk: bulk weights and block-scales use module.local_shared_*_tensors (attribute on module, deleted after register); per-expert scalars (alphas, weight_scale_2) use shared_* (function-local). Both are equally valid migration sources -- the distinction is historical.

Checklist for adding a new per-expert Parameter to an EPLB-supporting quantization method:

1. Register the on-module nn.Parameter sized expert_size_per_partition in create_weights().
2. In whichever loader fills it, also fill a tmp_shared_*_weight_scale_X dict keyed by enumerate(local_shared_load_expert_ids) during the need_load_shared_weights(module) branch.
3. In process_weights_after_loading() (or the equivalent finalize step), allocate a CPU shared_* buffer sized num_shared and fill it from the temp dict. Pass it as an explicit destination to reconcile/compute helpers -- do NOT write into the on-module .data[expert_idx] from the shared path, since expert_idx is in staging space and the on-module Parameter is in routed space.
4. Add the staging buffer to the weight_fns dict handed to register_all_parameter_slot_and_to_fix_weight_fns({...}) so migration can find it.
5. If the reconcile/compute helper is shared between routed and staging paths, its signature must take the destination tensor as a parameter (not read module.<param>.data directly), so the same body serves both index spaces.

Red flags:

• A new per-expert Parameter registered in create_weights() but never added to any weight_fns migration dict -- it will be stale after the first EPLB migration.
• A reconcile/compute function that both reads tmp_shared_* and writes module.<per_expert_param>.data[expert_idx] -- the staging-space index can exceed the routed-space bound (multi-node) or silently overwrite routed slots (single-node).
• Asymmetric gating: one of fc31_* / fc2_* pair registered but its twin not (or one added to weight_fns but not the other) -- migration will leave half the state stale.

Communication

Role:

• External communication strategies implement dispatch/combine and expose what ordering they support relative to quantization.
• Backends whose kernel owns cross-rank exchange should bypass external communication strategies rather than being forced through the factory.

Main APIs / references:

• communication/base.py: Communication, is_platform_supported, is_workload_feasible, supports_post_quant_dispatch, prepare_dispatch, dispatch, combine.
• communication/communication_factory.py: strategy selection.
• Existing strategies: nvlink_one_sided.py, nvlink_two_sided.py, deep_ep.py, allgather_reducescatter.py.

Checklist:

• Strategy selection and forced method behavior are handled through the factory.
• supports_post_quant_dispatch() is correct for the payload layout.
• prepare_dispatch() is used only for metadata/statistics that must happen before dispatch.
• dispatch() and combine() maintain enough internal state for the pair to be correct.
• EPLB statistics gathered by the communication strategy are fed back to the load balancer through the forward-execution path.
• Add/update test_moe_comm.py or module-level tests when changing strategy behavior.

Forward Execution And Chunking

Role:

• Treat moe_scheduler.py as the current owner of forward-time policy. Use this section as the detailed checklist for scheduler changes and for reviews that suspect policy has leaked back into the wrapper or backend.
• Keep lifecycle outside this policy: backend construction, weight loading, communication strategy lifetime, DWDP record, and repeat_idx advancement remain wrapper-level concerns.

Main APIs / references:

• moe_scheduler.py: scheduler ABC, external/fused scheduler implementations, chunk helpers, EPLB hook order, and backend kwargs construction.
• configurable_moe.py: scheduler construction and wrapper lifecycle after scheduler return.
• Current communication interfaces and backend run_moe/quantize_input contracts.
• Existing tests that exercise module forward, multi-GPU EP, EPLB, and communication behavior.

Checklist:

• The wrapper advances repeat_idx once per forward_impl; schedulers must not mutate it independently.
• External-communication scheduler respects padding, chunking, communication fallback, quantize/dispatch order, EPLB hooks, and output truncation.
• Fused-communication path does not call external Communication.dispatch or combine.
• Per-chunk EPLB first/last-call behavior is preserved.
• Multi-stream overlap is used only on paths that support it.
• Add module or focused forward-path tests for new policy, especially chunking and zero-token behavior.

Routing And Factory

Role:

• Routing methods map router logits to expert or slot selections.
• Factory/config code selects a backend based on requested backend, quantization, hardware capability, and model config.

Main APIs / references:

• routing.py: routing method implementations.
• create_moe.py: get_moe_cls, create_moe_backend, create_moe.
• moe_test_utils.py: backend enum, backend class map, skip logic.

Checklist:

• Routing output dtype/shape matches backend and forward-execution expectations.
• Unsupported backend/quant/model combinations fall back or skip with clear reasons.
• Test skip logic mirrors backend can_implement() instead of hiding bugs with broad skips.

Test Matrix And Helpers

Role:

• Keep backend, quantization, model-shape, routing, communication, and CI/local test matrices centralized and consistent across backend-level and module-level tests.
• Keep skip reasons aligned with production capability checks such as can_implement() instead of hiding failures with broad local skips.

Main APIs / references:

• tests/unittest/_torch/modules/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 generation and quant-parameter setup.
• test_moe_backend.py: backend interface tests for quantize_input and run_moe.
• test_moe_module.py: ConfigurableMoE integration matrix, multi-GPU, and EPLB coverage.
• test_moe_comm.py: communication dispatch/combine coverage.

Checklist:

• New backend is added to MoeBackendType, get_backend_class, backend/module matrices, and skip logic.
• New quantization method is added to test quant parameters and EPLB support checks when applicable.
• New unsupported combination returns a precise skip reason tied to production capability checks.
• CI subset and local exhaustive matrix stay intentionally different and are documented in the test helpers.
• Legacy tests such as test_fused_moe.py are used only for compatibility; new ConfigurableMoE behavior belongs in test_moe_backend.py, test_moe_module.py, or focused comm/routing/load-balancer tests.

Design Gate

Before editing, write a short gate:

## MoE Design Gate
- Change area: <ConfigurableMoE / MoEScheduler-forward-execution / backend / quantization-weights / EPLB / communication / routing-factory / test-matrix / tests>
- Owner boundary: <where the behavior belongs and why>
- Main API touched: <method/class names>
- Reference pattern: <existing file/class/function from references/moe-canonical-code-examples.md, with file:line evidence>
- Guide sections used: <MOE_DEVELOPER_GUIDE.md sections>
- Guide update needed: <yes/no; which section if yes>
- Refactor needed: <yes/no; one reason tied to architecture, not style>
- Test plan: <backend/module/comm/routing/EPLB/multi-GPU tests>

If the owner boundary is unclear, inspect more code before editing.

Refactor Rubric

Recommend a refactor when it:

• Moves behavior to the correct owner boundary.
• Simplifies ConfigurableMoE while preserving its assembler role.
• Clarifies backend ownership of the weight lifecycle and quantization-method delegation for weights/scales.
• Makes backend capabilities and unsupported combinations explicit.
• Separates external-communication and fused-communication policies cleanly in MoEScheduler rather than wrapper/backend branches.
• Makes EPLB support testable across interface, quantization, forward execution, and module tests.
• Updates shared test matrices/helpers when backend, quantization, or skip semantics change.
• Reduces duplicate dispatch/chunking/EPLB ordering logic by centralizing forward-time policy in moe_scheduler.py without changing performance-critical semantics.

Reject or question a refactor when it:

• Adds backend-specific forward branches to ConfigurableMoE instead of selecting behavior through MoESchedulerKind / MoEScheduler.
• Moves weight layout logic out of quantization methods without a strong reason.
• Hides hardware or quantization constraints behind vague abstractions.
• Changes communication/EPLB ordering without tests.
• Adds one-off skips in individual tests instead of shared capability/skip helpers.
• Touches legacy MoE paths for new features when the ConfigurableMoE path should be used.

Review Output

For reviews, lead with findings and concrete references:

## Findings
- [High] <file:line> <architecture, correctness, or testability issue>
- [Medium] <file:line> <maintainability or boundary issue>
- [Low] <file:line> <local cleanup>

## Architecture Fit
- ConfigurableMoE remains assembler: <yes/no>
- Owner boundaries respected: <yes/no>
- Scheduler boundary respected: <yes/no; forward policy in `moe_scheduler.py`, lifecycle in wrapper, compute in backend>
- Refactor recommended: <yes/no + reason>

## Guide Alignment
- Sections checked: <MOE_DEVELOPER_GUIDE.md sections>
- Guide update needed: <yes/no + section>

## Checklist Coverage
- Weights/quantization: <covered/gap>
- EPLB: <covered/gap>
- Communication: <covered/gap>
- MoEScheduler/forward execution: <covered/gap>
- Backend: <covered/gap>
- Forward execution/chunking details: <covered/gap>
- Test matrix/helpers: <covered/gap>
- Tests: <covered/gap>

If there are no findings, say so and list remaining test or performance risk.

Test Selection

Prefer the unified MoE tests:

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

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