Repository navigation
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
September 10, 2026 15:14
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
xuantengh
force-pushed
the
xuantengh/dsa_gqa_decode
branch
from
September 16, 2026 06:18
82bcafd to
2a546a7
Compare
xuantengh
force-pushed
the
xuantengh/dsa_gqa_decode
branch
from
October 18, 2026 13:03
2a546a7 to
7d99ae0
Compare
xuantengh
changed the title
xuantengh
force-pushed
the
xuantengh/dsa_gqa_decode
branch
from
October 22, 2026 11:49
22951fe to
30bd7e9
Compare
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
xuantengh
force-pushed
the
xuantengh/dsa_gqa_decode
branch
from
October 22, 2026 12:56
ebcedd4 to
30bd7e9
Compare
xuantengh
marked this pull request as ready for review
October 22, 2026 12:57
xuantengh
changed the title
Contributor
Author
|
/ok to test 1a3ec8f |
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Replace the CuTe decode kernels with two Triton kernels: - Indexer scores stream page-contiguous key tiles on tensor cores, and each request's programs stride over its row only up to its length, so the launch is fixed by batch and width while the work follows the real context lengths. - Sparse attention handles all query heads of a KV group in one program, reading each selected key and value once, splits the top-k slots across programs, and merges the slices (flash decoding). Select the top-k with DeepSelect when it is installed; it reads each row only up to its length, so the indexer skips the positions past it. Without DeepSelect, decode falls back to torch.topk over full rows and logs a warning. Per-step integers are not specialized, so each kernel compiles once per model. One decode layer at batch 64, 32k context, top-k 2048 drops from 1.34 ms to 0.16 ms on GB300, and no longer grows with the CUDA-graph score width. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
…d CUDA graphs Treat every query token of a step, decode or prefill, as a row that routes to the top-k of its causal prefix. The dynamic context builds the rows (request, causal length) and the indexer's score blocks (consecutive rows of one request) on the CPU in initialize_attention_state() and copies them with the bookkeeping, padding included, so the attention layer reads static buffers with no host syncs. The indexer kernel scores each block's rows against shared key tiles on tensor cores, and the sparse-attention kernel maps each row to its request. Decode, chunked prefill, prefix-cache continuations, and mixed batches now share one path whose launch sizes depend only on the padded batch, so every step type can run as a CUDA graph; the serving kit captures them all by default (CUDA_GRAPHS=decode keeps decode-only graphs). Inference no longer needs the training sparse-attention package. On GB300 with the midtrained top-k 2048 checkpoint, a mixed step with a 2048-token prefill chunk and 64 decodes per rank drops from about 130 ms (eager, per-request prefill) to about 68 ms (graphed); decode-only steps are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Score and select rows in windows whose score buffer fits SCORE_BUFFER_BYTES (1 GiB): a 2048-token prefill chunk at 1M context needed 8 GiB, now 1 GiB. The window count depends only on the padded rows and the width, so CUDA graphs replay it. The indexer kernel takes a row window; calls that cover every row compile without the window checks, so decode and shorter contexts run exactly as before (prefill at 1M context costs about 25% more). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
When a step's context bound reaches 64K keys, each row now thresholds at a sampled score, a CuTe DSL filter scores every key on tensor cores and keeps only those reaching the threshold (warp MMAs for decode, tcgen05 for prefill chunks), and DeepSelect picks the top k among the candidates. Rows the threshold misses take an exact radix fallback, so the result is still the exact top k. Sampling and both selections read only as far as each row's keys reach, so graphed steps cost what their real lengths need. The context builds 128-row filter blocks alongside the score blocks. Serving the midtrain checkpoint (top-k 2048) with a 262K bound and all-step CUDA graphs, 32 requests of 200K tokens: mixed steps 71.6 ms vs 76.1 ms, decode steps 8.49 ms vs 8.99 ms; the probe passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
--dsa-inference-backend flashinfer runs DSA attention through FlashInfer's DSA kernels, as vLLM deploys them, on Megatron's paged caches: - decode rows: dense BF16 indexer scores, exact top-k (gvr_2), trtllm-gen token-sparse attention; - prefill rows, in 128-row tiles: a threshold from sampled scores, one fused candidate sweep with an exact in-launch repair, the exact top-k among the candidates, then FA4 over each 8-token group's union of selected keys with a per-token bitmask. The indexer cache is viewed as FlashInfer's 64-key pages and the K/V caches as page-size-1 slots, so the caches are unchanged. The dynamic context builds the step's plan with the DSA row metadata and copies it with the bookkeeping; every launch size follows the padded batch, so decode, prefill and mixed steps all run as CUDA graphs. One layer, GB300, top-k 2048, graphed (Triton backend -> FlashInfer): decode batch 64 at 16K 131 -> 96 us, at 1M 2683 -> 2562 us; 32K prompt 13137 -> 6949 us; 2048-token chunk on 1M keys 3342 -> 2489 us. Serving the midtrain checkpoint (262K bound, all-step graphs, 32 x 200K requests): mixed steps 70.3 ms (71.6 ms Triton), decode 8.26 ms (8.49 ms); the probe passes. MMLU-Pro 71.49 and MATH-500 pass@1 79.85, against 72.32 and 79.85 on the reference eval branch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Replace handwritten inline PTX selection with a CuTe DSL histogram fast path and an exact Triton radix fallback. Keep Gluon scoring and Triton attention, with exact chunk merging for long contexts and top-k 512/1024/2048. Switch the inference extra from DeepSelect to CuTe DSL and retain the training-kernel, model-logprob, and CUDA Graph regression contracts. Validation: 102 targeted tests, seven determinism manifest checks, changed kernel coverage, and changed-file lint pass. Same-node checkpoint serving probes pass; decode latency versus FlashInfer is +3.0% at 32K/B8, +3.2% at 32K/B64, and +5.9% at 200K/B8. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Remove the inference-only CLI mapping that changed the simplified GQA indexer's rotary configuration for existing checkpoints. Keep upstream training defaults and explicitly supplied config settings unchanged. Replace the argument-propagation assertion with a numerical checkpoint reload regression covering non-default model RoPE settings. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Restore upstream VLM checkpoint detection and argument tests, removing the standalone VLM regression test to keep this PR focused on DSA decode. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Restore the upstream Blackwell FP32 probability reduction and remove the test that required rounding probabilities before the value reduction. Keep decode kernels and parity tolerances unchanged, and update the parity test comment to describe the original training behavior. Validation: 36 kernel parity cases and 4 model log-probability cases passed. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Move training/decode parity and CuTe top-k replay coverage into the DSA inference test module, and update the kernel manifest to the new path. Preserve all test bodies, parameterizations, and numerical tolerances. Validation: all 55 inference cases pass in the ordinary test environment without determinism-package initialization or a pinned autotune policy. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Replace the Triton radix selector with torch.topk. Capture fallback work inside a CUDA Graph conditional node so replay repairs rejected rows only when needed. Preserve valid negative-infinity scores ahead of padding and keep the existing cutlass.utils shared-memory allocator API. Cover fast-path/fallback transitions in one graph, ragged rows, infinities, and long-context chunk merging. Restore upstream checkpoint loading and backend documentation to keep this PR focused on DSA inference. Validation: 56 inference cases and 9 older-NGC cases passed, including four model log-probability comparisons. Changed-file pre-commit checks pass; the all-files check retains unrelated baseline failures. At batch 64 and top-k 2048, conditional-node overhead increased the no- fallback top-k stage from 16.5 to 24.6 us at 32K context and from 69.6 to 88.1 us at 200K context. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Move CUTLASS DSL from the inference extra to the default Linux runtime dependencies so plain installations include the DSA top-k dependency. Pin CUTLASS DSL and its runtime wheels to the compatible 4.8.1 release. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
xuantengh
force-pushed
the
xuantengh/dsa_gqa_decode
branch
from
November 3, 2026 15:38
1c6e0cc to
647a518
Compare
This branch is waiting to be deployed
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.
What does this PR do?
This PR adds native dynamic inference for simplified DSA over GQA, including paged KV caches, prefill, single-token decode, and mixed batches.
The Blackwell decode path combines Triton/Gluon indexer scoring, CuTe DSL exact topk selection, and Triton sparse attention. Exceptional topk cases use torch.topk, dispatched through a CUDA graph conditional node. The implementation also supports graph replay with changing token positions and active lengths, and manages decode workspaces across KV-cache suspend/resume. CUTLASS DSL 4.8.1 is included in the default Linux dependencies.
Performance
We also conducted the performance comparison between this PR (Triton/Gluon indexer + sparse attention with CuTe DSL topk against FlashInfer SDSA not public available yet):
Accurancy
Checkpoint accuracy evaluation: