Skip to content

Fix DSV4 C4 top-k v2 dual output in one pass - #38016

Closed
AMD-yanfeiwang wants to merge 3 commits into
sgl-project:mainfrom
AMD-yanfeiwang:fix/dsv4-c4-topk-v2-raw-rocm
Closed

AMD-yanfeiwang wants to merge 3 commits into
sgl-project:mainfrom
AMD-yanfeiwang:fix/dsv4-c4-topk-v2-raw-rocm

Conversation

@AMD-yanfeiwang

@AMD-yanfeiwang AMD-yanfeiwang commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Motivation

Fixes #37892.

DSV4 paged prefill produces two representations of the same C4 selection:

  • request-local raw C4 positions for sparse prefill / capture
  • page-table-transformed KV slots for attention

The generic DSV4 backend allocates c4_sparse_raw_indices for paged prefill, but the C4 dispatch only selected top-k v2 when raw_indices is None. Those prefills therefore fell back to the legacy v1 transform.

ROCm's ordinary non-CP HIP backend does not request this raw buffer, but CP prefill, indexer capture, and HiSparse decode can. #36684 explicitly enabled top-k v2 for DeepSeek-V4 on HIP, yet each of those raw-output paths still falls back to the AOT v1 operator. The exact GB300 crash is architecture-specific; the raw-output dispatch and performance issue is shared.

#35041 removed v2's old optional raw output because no production caller used
it and page-only calls paid a nullable-pointer branch. Paged prefill now is that production caller.

This takes the same dispatch direction as #37991, but emits both outputs in one v2 invocation rather than repeating the complete top-k selection.

Changes

  • Add a compile-time PAGE_TABLE_AND_INDICES v2 mode. It writes aligned raw and paged outputs from the same selected indices.
  • Keep raw-output stores compiled out of the existing page-only and raw-only modes, preserving [DSA] Trim top-k v2 output modes and tighten its PDL waits #35041's page-only hot-path property.
  • Route the C4 sgl-kernel backend to v2 even when raw indices are requested.
  • Reuse a correctly sized view of cached metadata for HIP FP4 row chunks. ROCm compiles out the cluster planner, so allocating a no-op plan per layer/chunk is unnecessary.
  • Keep XPU dual output on its existing single-selection operator.
  • Cover zero-length DP rows, ties, ragged lengths, k={512,1024,2048}, and the trivial/Register2/Register4/streaming/CUDA small-cluster/persistent-cluster dispatch boundaries, plus XPU wrapper dispatch.

Correctness

MI355X, ROCm 7.2, PyTorch 2.11:

  • test_topk_v2.py: 297 passed
  • focused C4 dispatch regression: 1 passed / 13 deselected
  • fused raw and paged outputs are exactly aligned, including tie-heavy rows
  • fused output selected sets match both two-pass v2 and legacy v1
  • DeepSeek-V4-Pro DCP8 raw-output prefill produced identical output token IDs before and after the change

Two independent v2 calls select the same sets but do not preserve the same output ordering between the raw and paged calls. The fused mode derives both buffers from one selection and keeps them slot-aligned.

Performance

MI355X, ROCm 7.2, PyTorch 2.11, median of 8 interleaved samples. L is the C4 sequence length.

  • B=8192, L=2048, k=512:
    • v1 dual output: 0.12080 ms
    • two-pass v2: 0.13874 ms
    • fused dual-output v2: 0.07249 ms
    • fused is 1.67x faster than v1 and 1.91x faster than two-pass v2
  • B=1024, L=8192, k=512:
    • v1: 0.03314 ms; two-pass v2: 0.02958 ms; fused v2: 0.01562 ms
    • fused is 2.12x faster than v1 and 1.89x faster than two-pass v2
  • B=8192, L=17920, k=512 (about 70K uncompressed context):
    • v1: 0.41592 ms; two-pass v2: 0.36814 ms; fused v2: 0.17933 ms
    • fused is 2.32x faster than v1 and 2.05x faster than two-pass v2
  • B=1024, L=17920, k=1024:
    • v1: 0.04971 ms; two-pass v2: 0.05312 ms; fused v2: 0.02790 ms
    • fused is 1.78x faster than v1 and 1.90x faster than two-pass v2

The fused mode is within -0.6% to +2.7% of page-only v2 across those shapes, so the second output is nearly free compared with repeating selection.

DeepSeek-V4-Pro TP8/DCP8 (which requests raw C4 indices), chunked prefill 16384, concurrency 2, one output token, five repeats:

  • 32K input median TTFT: 1726.21 -> 1729.66 ms (+0.20%)
  • 70K input median TTFT: 3968.36 -> 3975.41 ms (+0.18%)

There is no measurable end-to-end improvement in this configuration; these deltas are within system noise. The kernel-level result is the performance claim.

Testing

python3 -m py_compile \
  python/sglang/kernels/ops/attention/dsv4/topk.py \
  python/sglang/srt/layers/attention/dsv4/indexer.py \
  test/registered/kernels/ops/attention/test_topk_v2.py \
  test/registered/kernels/ops/attention/test_dsa_indexer.py

python3 -m pytest \
  test/registered/kernels/ops/attention/test_topk_v2.py -q

python3 -m pytest \
  test/registered/kernels/ops/attention/test_dsa_indexer.py \
  -k c4_raw_indices_keep_sgl_backend_on_topk_v2 -q

Pre-commit checks passed for isort, ruff, ruff-format, codespell, and clang-format.


CI States

Latest PR Test (Base): ❌ Run #33961904069
Latest PR Test (Extra): ❌ Run #33961903803
Latest PR Test (AMD ROCm 7.2): ❌ Run #33961903978

@alphabetc1

Copy link
Copy Markdown
Collaborator

I think this is a better implementation. Let's move forward with it.

@DarkSharpness

Copy link
Copy Markdown
Collaborator

maybe duplicate of #33672

@AMD-yanfeiwang

AMD-yanfeiwang commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

@DarkSharpness Agreed, thanks for pointing this out. Looking forward to #33672 landing. Once it is merged, I’ll close this PR and follow up with a separate PR for the ROCm-specific optimizations and additional test coverage from this work.

@AMD-yanfeiwang

Copy link
Copy Markdown
Contributor Author

Solved by #33672

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

3 participants