Skip to content

[None][feat] add Rubin topology support to GVR V2 prefill - #18764

Draft
longcheng-nv wants to merge 16 commits into
NVIDIA:mainfrom
longcheng-nv:feat/gvr-v2-rubin-prefill
Draft

[None][feat] add Rubin topology support to GVR V2 prefill#18764
longcheng-nv wants to merge 16 commits into
NVIDIA:mainfrom
longcheng-nv:feat/gvr-v2-rubin-prefill

Conversation

@longcheng-nv

Copy link
Copy Markdown
Collaborator

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

  • Carry the self-sampling GVR V2 prefill path from [None][feat] Self-sampling GVR V2 prefill indexer top-K #18702 to SM107.
  • Replace B200-specific launch assumptions with runtime SM topology.
  • Preserve the existing prefill interval, input-format, short-row, slab, and radix-fallback behavior.
  • Retain the inherited default-off decode locality-domain prototype, including its BS=1 full-device fallback and fail-closed capture/device checks.

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.

  • The indexer producer and logits remain global/full-device.
  • The low-level launch is topology/domain aware, but production prefill does not enable dual-domain routing.
  • TMA, PDL, localized logits/K-cache allocation, and producer-side locality are not implemented.
  • The inherited decode prototype is limited to one in-flight model execution per CUDA device.
  • No R200 correctness or performance result is claimed.

Validation

Passed locally on the restacked branch:

  • GVR decode host tests: 11/11.
  • Rubin prefill topology tests: 10/10.
  • Decode route parity with the companion decode branch: 21,504 cases.
  • Decode streaming parity: 43,008 comparisons.
  • Prefill R=1 route/topology invariants: 21,504 cases.
  • Relevant pre-commit formatting/lint, Python compilation, DCO, and whitespace checks.

Full package/GPU pytest collection is blocked in the current environment by the external bindings.so mismatch (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:

  • Run the companion decode R200 matrix over every BS=1-1024 and compressed N={4K,8K,16K,32K,64K,128K,256K} on this complete stacked branch.
  • Validate prefill correctness for K={512,1024,2048}, supported compression ratios, short rows, unaligned [ks, ke) intervals, long rows, and slab-boundary cases.
  • Exercise eager execution and CUDA Graph warmup/replay on strict and balanced Rubin topology.
  • Measure prefill and indexer + Top-K end-to-end latency against the full-device baseline.
  • Confirm that SM100/SM103 behavior and unsupported-path fallbacks remain unchanged.

longcheng-nv and others added 15 commits September 3, 2026 19:35
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant