[None][fix] Drop the stale choices list on --moe-backend-for-prefill - #17612
[None][fix] Drop the stale choices list on --moe-backend-for-prefill#17612ruodil wants to merge 1 commit into
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)
WalkthroughThe benchmark CLI now accepts arbitrary values for ChangesMoE backend validation
Estimated code review effort: 2 (Simple) | ~10 minutes Mergeability Score: ⚪ Minimal · up to This localized change removes a stale CLI restriction so supported backend values can reach existing validation; no actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the 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.
| parser.add_argument("--max-num-tokens", type=int) | ||
| parser.add_argument("--moe-backend", type=str) | ||
| parser.add_argument("--moe-backend-for-prefill", type=str, choices=["CUTLASS", "DEEPGEMM"]) | ||
| # No choices= here, for the same reason --moe-backend above has none: both values reach |
There was a problem hiding this comment.
The comment says "this file's own runner.py", but replace_routing_method_ctx lives in tensorrt_llm/tools/layer_wise_benchmarks/runner.py:688, not in this directory. Suggest collapsing the whole block to something that can't go stale, e.g.:
# No choices= (same as --moe-backend above): both flags feed MoeConfig.backend,
# whose Literal is the authoritative list. A second list here can only drift.4 participants
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=...):run.py:182—Runner(..., moe_backend=args.moe_backend, ...)run.py:252—Runner(..., moe_backend=args.moe_backend_for_prefill, ...)runner.py—llm_args = TorchLlmArgs(..., moe_config=MoeConfig(backend=moe_backend, ...))MoeConfig.backendis a pydanticLiteralover eleven values and 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)runner.pyreplace_routing_method_ctxrun.py --moe-backend-for-prefill(before)TRTLLM and CUTEDSL are excluded by the CLI even though
runner.pyin this same directory names 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:135-136). 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.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.QA Engineer Review
No test changes.