[None][feat] add Rubin topology support to GVR V2 prefill - #18764
Draft
longcheng-nv wants to merge 16 commits into
Draft
[None][feat] add Rubin topology support to GVR V2 prefill#18764longcheng-nv wants to merge 16 commits into
longcheng-nv wants to merge 16 commits into
Conversation
…ntion config enable_heuristic_topk keeps selecting the GVR family over the exact radix path; a new use_self_sampling_topk config field (default True) selects the hint-free self-sampling engine over the temporal-hint engines. The CUTE_DSL_GVR_V2 enum folds into CUTE_DSL_GVR behind a gvr_self_sampling module flag, TopK.needs_gvr_prior follows the two-level decision, and the retired TRTLLM_GVR_SELF_SAMPLING env only warns. The field threads llm_args -> model_config -> DSAParams/DSAMetadataParams -> indexer and the warmup mirror (whose top_k source also moves off a dead index_topk getattr to sparse_mla_topk). Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Nothing selects the CUDA heuristic once the unified DSL GVR router is
in: delete heuristicTopKDecode.{cu,h} / heuristic_topk.cuh, the
canUseHeuristic dispatch and the GVR SchemeX bounds in indexerTopK.cu
(radix keeps a cached SM-count helper), shrink the indexer_topk_decode
thop schema and its register_fake (pre_idx / heuristic_scratch gone),
drop the CUDA_GVR enum plus module branch, and retire the
heuristic-only distribution / hostile-hint / tie-plateau test arms.
The radix insertion / histogram / split-work tiers are untouched.
C++ changes are not compiled yet: build + CI plus the 886x11 CUDA-v1
vs DSL-v1 paired A/B sign-off gate this draft.
Made-with: Claude Code (Fable 5)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
The self-sampling engine keeps no cross-step state, so the framework no
longer allocates it any: the per-layer gvr_prior_indices arena, the LJF
row-reorder buffer, prefill seeding, the aux-stream write-back, and the
indexer-side prior slice all key on needs_gvr_prior = two-level dispatch
selecting the temporal engine. A shared use_self_sampling_gvr() predicate
in dsa/params.py keeps the indexer's per-layer TopK construction and the
metadata's allocation decision in agreement (live indexers only exist on
cr in {1, 4} layers, matching the metadata's representative ratio). With
the CUDA heuristic gone, the temporal engine requires the CuTe DSL on
SM100/103; enable_heuristic_topk without it falls back to exact radix
with a one-time warning.
Made-with: Claude Code (Fable 5)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… param Promote the emission-assisted block-skip optimization from the TRTLLM_GVR_EMISSION env var to a `use_gvr_emission` sparse-attention config field (default False). It only takes effect on the temporal-hint (V1) GVR path with FP4 paged-MQA logits; the self-sampling (V2) engine derives its bracket from the current row and never uses emission. Threads the field through llm_args -> model_config (V4 + V3.2 rebuilds) -> DSAParams / DSAMetadataParams -> the indexer gate, and adds config-threading unit tests. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…e GVR dispatch fields Adds the two new sparse-attention config fields (use_self_sampling_topk, use_gvr_emission; both bool, captured by value) via scripts/generate_llm_args_golden_manifest.py so test_build_capture_manifest_matches_committed_golden passes again. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Picks up the main-side waive of unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py (NVIDIA#18685) that failed the CPU lanes in the previous CI round. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Picks up the main-side waive of unittest/_torch/visual_gen/test_trtllm_serve_e2e.py TestFlux1/Flux2TextToImage (nvbugs/6720944) that failed the DGX_B200 lane in the previous CI round. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Resolve tests/unittest/_torch/modules/test_top_k.py: keep the new CUTE_DSL_RADIX prefill-dispatch test, drop the removed CUDA_GVR workspace test, and drop the HEAD unsupported-impl test that now targets the main-supported CUTE_DSL_RADIX; keep the CUTE_DSL_GVR unsupported-prefill guard test (this branch has no GVR prefill form). Made-with: Claude Code (Opus 4.8) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Extend the hint-free self-sampling GVR top-K from decode to the DSA prefill phase, replacing the CUDA radix prefill on the layers that already select self-sampling for decode (same enable_heuristic_topk x use_self_sampling_topk dispatch; no new config/API). Kernel: a const_expr prefill mode on GvrMainKernel for per-row [ks, ke) windows. ks/ke ride the unused pre_idx/kv_lens ABI slots (byte-identical signature; distinct compile key). The base rounds down to a 16B boundary and the <=3 lead lanes are masked positionally (no materialized -inf); the clamps tighten to the last in-window float4 so the reads are exact. Output is the local (column - ks) frame with a trailing -1 pad; nv <= k emits identity. All edits are const_expr-gated, so decode/legacy codegen is unchanged. Host: run_prefill (stride(0) for all row counts, no device reads, <=32768-row slabs for gridDim.y, R=1) + warmup_prefill (<=6 engines/k). Module: a CUTE_DSL_GVR prefill branch with an all-short short-circuit and a format-gate radix fallthrough. Indexer selects the engine for prefill iff it does for decode; metadata warms the prefill leg. Compression-ratio agnostic, so V3.2, V4 Flash and V4 Pro share one prefill path. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Stale-base artifact from the merge-from-main baseline, not a Rubin/GVR change; restore so the PR does not drop authorized users. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
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
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.
Dependencies
Built on #18446, #18702, and the companion decode Draft #18763. Keep this PR in Draft until those dependencies are settled, then rebase out the duplicated base commits.
Summary
Current scope
Prefill deliberately remains a full-device launch. The current path has one logical row and one CTA per row, so dispatching it to two domains would add launch/event overhead without creating more row-level parallelism or memory locality.
Validation
Passed locally on the restacked branch:
Full package/GPU pytest collection is blocked in the current environment by the external
bindings.somismatch (global_steady_clock_now). The installed CuTe DSL also lacks the Rubin helper, so SM107 compile, eager, CUDA Graph, correctness, and performance paths have not been exercised.Draft exit criteria
Before marking Ready:
[ks, ke)intervals, long rows, and slab-boundary cases.