[DSV4] Support raw-index output in TopK v2 - #33672
Conversation
|
/tag-and-rerun-ci |
22bfab0 to
8787445
Compare
|
/rerun-test test/registered/unit/layers/test_dsv4_nonpaged_indexer.py |
|
Results for ⛔ |
|
@weireweire Is this for InfX? I'm wondering why you using |
weireweire
changed the title
|
@b8zhong it's for infX. Even without |
8787445 to
948829f
Compare
|
/rerun-failed-ci |
|
/tag-and-rerun-ci extra |
Branch-local port of sgl-project#33672 (open upstream). The v2 kernel signature, c4_sparse_raw_indices internal allocation, and capturer buffer reuse already existed on this branch; the missing piece was the dispatch gate in _run_topk_transform plus buffer-priority ordering. - Prioritize core_metadata.c4_sparse_raw_indices over the capture-only temporary buffer (sparse-prefill buffer is canonical; the capturer copies the same result afterward instead of stealing it). - Drop the raw_indices-is-None gate on the TopK v2 path so sparse prefill no longer falls back to TopK v1; pass the raw-output buffer through. - Chunked path keeps the existing topk_metadata regeneration guard. - Test adaptions vs upstream: get_parallel patched (CP group uninitialized in unit tests); fake_v2 signatures take the 7th raw_output positional arg. GB300 microbenchmark (8,192 queries, max C4 2048, topk=1024): v2 kernel 0.0646 ms vs v1 0.1023 ms (-36.9%), larger wins on long prefixes.
3c7ecbb to
c9bd7c6
Compare
|
@weireweire Please fix conflicts |
|
@b8zhong conflict fixed |
DarkSharpness
left a comment
DarkSharpness
left a comment
There was a problem hiding this comment.
LGTM. I just wonder in which cases do we need both page & the topk indices?
|
@DarkSharpness thanks, both sparse prefill and |
Root cause: TopK v2 cannot emit raw and paged indices together. Fix: Add dual-output kernel dispatch and preserve the sparse-prefill output buffer when capture is enabled. Rebase the complete PR patch onto current main. Validation: Full pre-commit checks passed. Stable patch-id matches the original PR patch; git diff --check passed.
362ad8e to
a103b0d
Compare
|
/rerun-failed-ci |
|
/rerun-test registered/rl/test_return_indexer_topk.py registered/kernels/ops/attention/test_topk_v2.py registered/unit/layers/test_dsv4_nonpaged_indexer.py |
|
Results for 🚀 🚀 🚀 |
|
/rerun-failed-ci |
|
could we merge this? |
…) (sgl-project#39098) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> (cherry picked from commit da64c5c)
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca> Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
6 participants
Summary
Allow TopK v2 to populate a raw-index output buffer. DSV4 sparse prefill internally allocates
c4_sparse_raw_indiceseven when--enable-return-indexer-topkis disabled. The previousraw_indices is Nonegate therefore forced the InfX sparse-prefill path to fall back to TopK v1.When
--enable-return-indexer-topkis enabled, the capture-only temporary buffer could additionally take precedence over the functional sparse-prefill buffer. This change makes the sparse-prefill buffer canonical, passes it through TopK v2, and lets the capturer copy the same result afterward.Changes
topk_transform_512_v2.core_metadata.c4_sparse_raw_indices; the capturer reads the same result afterward.No CUDA kernel, ABI, metadata format, default, or HiSparse behavior changes.
Related: #32771 contains the same buffer-ordering change; #33288 owns chunked-path metadata refresh.
Testing
pre-commit run --all-filesA prior GB300 microbenchmark (8,192 queries, max C4 length 2,048,
topk=1024, metadata planning excluded) measured 0.1023 ms for v1 and 0.0646 ms for v2 (36.9% lower kernel latency). Should have bigger influence on long prefix as the indexer will take more time.CI States
Latest PR Test (Base): ⏳ Run #34427908088
Latest PR Test (Extra): ✅ Run #34427908071
Latest PR Test (AMD ROCm 10): ❌ Run #34427908194