Repository navigation
[None][fix] Drop the stale choices list on --moe-backend-for-prefill - #17612
Conversation
--moe-backend and --moe-backend-for-prefill are handed to the same
Runner(moe_backend=...), which puts both into MoeConfig(backend=...).
MoeConfig.backend is a pydantic Literal over eleven values and is the
authoritative list; it rejects anything else with a clear validation error.
--moe-backend has no argparse choices and relies on that. This one carried a
second, hardcoded list, and it drifted: it accepted CUTLASS and DEEPGEMM only,
excluding TRTLLM and CUTEDSL -- both of which runner.py itself names as
supported in replace_routing_method_ctx, and TRTLLM being what the DeepSeek-V4
configs actually use.
The effect is that passing a backend the library supports fails in argparse
before anything runs:
run.py: error: argument --moe-backend-for-prefill: invalid choice: 'TRTLLM'
(choose from 'CUTLASS', 'DEEPGEMM')
Removing the list makes the two flags symmetric and leaves validation with
MoeConfig, which cannot drift from itself.
Not changed here, but worth flagging separately: this argument defaults to
CUTLASS when unset, and CUTLASS fp8 block-scale GEMM is Hopper-only, so the
default path is not usable on Blackwell. Changing a default affects existing
users and belongs in its own change.
Signed-off-by: Ruodi Lu <ruodil@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan includes up to 12 reviews per rolling hour; 10 remain after this review. WalkthroughThe benchmark CLI now accepts arbitrary values for ChangesMoE backend validation
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change removes a stale CLI restriction so supported MoE backends can be passed through to existing validation; no actionable merge-blocking risk remains beyond normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run --disable-fail-fast |
|
PR_Github #65795 [ run ] triggered by Bot. Commit: |
|
PR_Github #65795 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Removing the choices list is right: both flags feed the same Runner(moe_backend=...) → MoeConfig.backend, and only one of them had a duplicate list to drift. Two small things: the comment block is 6 lines for a 1-line change and will itself go stale (it enumerates counts and backend names); a single line pointing at MoeConfig.backend as the authority would age better. And the PR description lists WIDEEP among the backends replace_routing_method_ctx accepts — the check at tensorrt_llm/tools/layer_wise_benchmarks/runner.py:692 only allows CUTEDSL/CUTLASS/DEEPGEMM/TRTLLM.
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
…ble claim The six-line block enumerated backend names and a count, which is exactly the kind of thing that goes stale: it said "eleven values" while MoeConfig.backend has had twelve since NVIDIA#16190, and it attributed replace_routing_method_ctx to "this file's own runner.py" when that function lives in tensorrt_llm/tools/layer_wise_benchmarks/runner.py. Keep only the part that cannot drift -- MoeConfig.backend is the authority, so a second list here can only diverge from it. Signed-off-by: Ruodi Lu <ruodil@users.noreply.github.com>
|
Review comments addressed in Why the last run was red, and why it wasn't this PR.
Both builds passed ( Why skip rather than re-run. The net diff against |
|
/bot skip --comment "argparse-only change in examples/layer_wise_benchmarks/run.py; prior #53501 failures (GPT-OSS Eagle3 unit tests, pybind11 FetchContent clone, BERT quickstart timeout) are unrelated" |
|
PR_Github #66665 [ skip ] triggered by Bot. Commit: |
|
PR_Github #66665 [ skip ] completed with state |
What
--moe-backend-for-prefillcarried an argparsechoiceslist that had drifted out of sync with the values the library actually supports, so passing a supported backend failed before anything ran:Why the list should not be there
--moe-backendand--moe-backend-for-prefillare handed to the sameRunner(moe_backend=...), which puts both intoMoeConfig(backend=...):examples/layer_wise_benchmarks/run.py:184—Runner(..., moe_backend=args.moe_backend, ...)examples/layer_wise_benchmarks/run.py:254—Runner(..., moe_backend=args.moe_backend_for_prefill, ...)tensorrt_llm/tools/layer_wise_benchmarks/runner.py:411—llm_args = TorchLlmArgs(..., moe_config=MoeConfig(backend=moe_backend, ...))MoeConfig.backendis a pydanticLiteraland is the authoritative list.--moe-backendhas nochoicesand relies on exactly that. The second list could only ever drift from it, and it did:MoeConfig.backend(authoritative,llm_args.py:1420)replace_routing_method_ctx(tensorrt_llm/tools/layer_wise_benchmarks/runner.py:692)run.py --moe-backend-for-prefill(before)TRTLLM and CUTEDSL are excluded by the CLI even though
replace_routing_method_ctxnames both as supported. TRTLLM is what the DeepSeek-V4 configs actually use, which is how this surfaced.Removing the list makes the two flags symmetric and leaves validation with
MoeConfig, which cannot drift from itself.Flagged, not changed
This argument defaults to
CUTLASSwhen unset (run.py:137-138). CUTLASS fp8 block-scale GEMM is Hopper-only, so the default path is not usable on Blackwell — a caller on B200 must pass the flag, and until this change could not pass the backend its model uses. Changing a default affects existing users, so it belongs in its own change rather than being folded in here.Test Coverage
Argparse-level change only. No behaviour change for any value that was already accepted, and no new value is accepted that
MoeConfigwould not accept.I could not execute the validation path locally — this machine has neither
torchnorpydanticinstalled. The claim thatMoeConfigrejects invalid values rests on reading the type chain (class MoeConfig(StrictBaseModel)→StrictBaseModel(pydantic.BaseModel)→backend: Literal[...]), which is standard pydantic semantics, rather than on a run. Worth a reviewer confirming that an obviously bogus value still fails cleanly — just later, and with a pydantic error instead of an argparse one.Review follow-ups
Two corrections from @brnguyen2's review, both verified against
main:replace_routing_method_ctxaccepts. It does not — the check attensorrt_llm/tools/layer_wise_benchmarks/runner.py:692allows only CUTEDSL/CUTLASS/DEEPGEMM/TRTLLM. Fixed above.replace_routing_method_ctxlives in "this file's own runner.py". There is norunner.pyunderexamples/layer_wise_benchmarks/; the function is intensorrt_llm/tools/layer_wise_benchmarks/runner.py:688.The comment block also said
MoeConfig.backendhad "eleven values" — it has had twelve since #16190 (2026-07-29), so the enumeration was already stale on the day it was written. It has been collapsed to the one claim that cannot drift:PR Checklist
[None][fix] <description>Dev Engineer Review
choices=["CUTLASS", "DEEPGEMM"]restriction from--moe-backend-for-prefill.--moe-backendand delegates backend validation toRunnerandMoeConfig(backend=...).TRTLLMandCUTEDSLcan now reach library-level validation.CUTLASS.torchandpydanticwere unavailable.QA Engineer Review
No test changes.