docs: close v0.6.13 doc-check gaps + fix(moe) misleading topk_indices ICHECK message - #3546
Merged
Merged
Conversation
Resolve the doc-check findings new in v0.6.13rc1 vs v0.6.12, touching only docstrings and docs/api/*.rst (no new @flashinfer_api decorators). MISSING (API exported but not listed in .rst): - fused_moe.rst: add a "Multi-LoRA MoE (BGMV)" section listing bgmv_moe / bgmv_moe_shrink / bgmv_moe_expand. - quantization.rst: list nvfp4_quantize_per_token_cute_dsl under its canonical flashinfer.quantization.kernels.nvfp4_quantize currentmodule (the package-level re-export is guarded by is_cute_dsl_available() and is not importable at docs-build time without nvidia-cutlass-dsl). STALE (.rst listed a symbol with no corresponding @flashinfer_api fn): Since adding decorators is out of scope, clear these by rst cleanup. - cute_dsl.rst: drop the "Availability" autosummary for is_cute_dsl_available (a bool capability probe, not a runtime API); it is still referenced in prose. - comm.rst: drop trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant from the trtllm_mnnvl_ar autosummary (its non-quant sibling is also undecorated and stays as-is to keep the change minimal). Docstring completeness: - nvfp4_quantize_per_token_cute_dsl: expand the one-line summary into a full NumPy-style docstring with Parameters/Returns, mirroring the sibling nvfp4_quantize_cute_dsl and documenting the per-token scale output. (mla_rope_quantize_fp8 was already backfilled by flashinfer-ai#3456.) AI-assisted.
The dtype check for topk_indices in NoAuxTc requires int32 (encode_dlpack_dtype(...) == int32_code), but the assertion message was copy-pasted from the topk_values check and wrongly read "must have the same dtype as scores", which misleads debugging (scores are fp32/bf16, not int32). Change the message to "topk_indices must be int32 dtype", matching the nearby routing_replay_out int16 check style. Message-only change; the runtime condition is unchanged. AI-assisted.
kangbintNV
requested review from
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
June 9, 2026 02:18
Contributor
📝 WalkthroughWalkthroughThis PR updates API documentation across multiple modules (comm, cute_dsl, fused_moe, quantization) and corrects a kernel validation error message. The changes include removing stale documentation entries, documenting new BGMV MoE operators, expanding NVFP4 quantization kernel documentation, and fixing a topk_indices dtype validation message in the NoAuxTc kernel. ChangesKernel and API Documentation Updates
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsStopped waiting for pipeline failures after 30000ms. One of your pipelines takes longer than our 30000ms fetch window to run, so review may not consider pipeline-failure results for inline comments if any failures occurred after the fetch window. Increase the timeout if you want to wait longer or run a Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request updates error messaging in the fused MoE CUDA kernels, expands documentation for Multi-LoRA MoE (BGMV) and per-token NVFP4 quantization, and cleans up unused API references in the documentation. Specifically, a detailed docstring was added for nvfp4_quantize_per_token_cute_dsl. There are no review comments, so I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
kangbintNV
marked this pull request as draft
June 9, 2026 02:33
Re-add two autosummary entries that the earlier docs commit removed, because both are genuinely user-facing and should stay discoverable in the rendered docs: - docs/api/comm.rst: trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant, a real tensor-in/tensor-out MNNVL API with a full docstring. - docs/api/cute_dsl.rst: is_cute_dsl_available (Availability section), a public helper in flashinfer.cute_dsl.__all__ that users call to guard CuTe-DSL code paths. Keeping them listed means doc-check still reports both as STALE (neither carries @flashinfer_api), which is the accepted trade-off under the "no new decorators" constraint: prefer documenting the API over silencing the checker. Their undecorated siblings are treated the same. AI-assisted.
kangbintNV
marked this pull request as ready for review
June 9, 2026 03:03
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two small, low-risk changes surfaced by the v0.6.13 doc-check pass. Touches docs/docstrings only, plus a one-line C++ assertion-message fix (no runtime behavior change).
docs: close v0.6.13-only doc-check gaps (MISSING / STALE / docstring) issues/3538
Resolves the doc-check findings that are new in v0.6.13 vs v0.6.12, using docs + docstrings only (no new
@flashinfer_apidecorators)..rst):docs/api/fused_moe.rst: add a "Multi-LoRA MoE (BGMV)" section listingbgmv_moe/bgmv_moe_shrink/bgmv_moe_expand.docs/api/quantization.rst: listnvfp4_quantize_per_token_cute_dslunder its canonicalflashinfer.quantization.kernels.nvfp4_quantizecurrentmodule (the package-level re-export is guarded byis_cute_dsl_available()and is not importable at docs-build time withoutnvidia-cutlass-dsl).nvfp4_quantize_per_token_cute_dsl: expand the one-line summary into a full NumPy-style docstring with Parameters/Returns, mirroring the siblingnvfp4_quantize_cute_dsland documenting the per-token scale output.fix(moe): correct misleading
topk_indicesdtype ICHECK messageIn
csrc/fused_moe/noAuxTcKernels.cu, the dtype check fortopk_indicesrequires int32 (encode_dlpack_dtype(...) == int32_code), but the assertion message was copy-pasted from thetopk_valuescheck and wrongly read "must have the same dtype as scores" (scores are fp32/bf16, not int32), which misleads debugging. Message changed to "topk_indices must be int32 dtype", matching the nearbyrouting_replay_outint16 check style. Message-only change; the runtime condition is unchanged.Test plan
docs/build_docs.sh) succeeds; new autosummary entries render and removed STALE entries no longer produce dangling references.flashinfer_document_checkshows 0 MISSING and the two v0.6.13 STALE items (is_cute_dsl_available,trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant) gone; Docstring Completeness has 0 findings.AI-assisted.
Summary by CodeRabbit
New Features
bgmv_moe,bgmv_moe_shrink, andbgmv_moe_expand.Documentation