diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index 99ed41ce7..24059de22 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh @@ -250,6 +250,103 @@ TOPK_KERNEL void topk_ragged_kernel(const __grid_constant__ TopKRaggedParams par // PDL trigger secondary at the end the block typically has no use, so ignore it } +#ifdef USE_ROCM +// Only the ROCm DSA prefill emits this layout today, so CUDA/XPU builds stay +// unchanged. Nothing below is AMD-specific; the guard can be dropped later. + +/** + * \brief Parameters of the packed (DSA extend) layout. + * + * Same addressing as the ragged layout -- every row's window lives inside one + * batch-global score buffer starting at `row_starts[i]` -- but the selected + * columns are mapped through a page table before they are written out. Prefill + * expands one request into many query-token rows, so several score rows share a + * page-table row; `row_to_batch[i]` says which one. + */ +struct TopKPackedParams { + // NOTE: may write. The head of the window is masked in place, see the kernel. + float* __restrict__ scores; + const int32_t* __restrict__ seq_lens; // per-row window length + const int32_t* __restrict__ row_starts; // per-row score column offset + const int32_t* __restrict__ row_to_batch; // per-row page-table row; null => identity + const int32_t* __restrict__ page_table; + int32_t* __restrict__ page_indices; + int64_t score_stride; + int64_t page_table_stride; + uint32_t topk; + uint32_t page_bits; + + SGL_DEVICE PageTransform get_transform(uint32_t bx) const { + const auto table_row = row_to_batch == nullptr ? bx : static_cast(row_to_batch[bx]); + return {page_table + static_cast(table_row) * page_table_stride, page_bits, nullptr}; + } +}; + +/** + * \brief Top-k over packed rows, emitting page-table indices. + * \tparam kPDL whether to use PDL to synchronize with the indexer kernel + * + * Dispatch mirrors `topk_ragged_kernel`: both are prefill kernels, so the level + * is picked per row at runtime and only the register and streaming + * implementations are instantiated (no plan, no cluster path). + */ +template +TOPK_KERNEL void topk_packed_kernel(const __grid_constant__ TopKPackedParams params) { + device::enable_smem_spilling(); + constexpr uint32_t kVecSize = impl::TopKStreaming::kVecSize; + __shared__ impl::MaxSmem smem; + __shared__ int32_t s_topk_indices[kMaxTopK]; + + const auto bx = blockIdx.x; + const auto seq_len = static_cast(params.seq_lens[bx]); + const auto row_start = static_cast(params.row_starts[bx]); + const auto topk = params.topk; + const auto transform = params.get_transform(bx); + const auto out = params.page_indices + bx * static_cast(topk); + const auto score = params.scores + bx * params.score_stride; + + auto problem = TopKProblem{ + .in = score + row_start, + .out = out, + .topk = topk, + .seq_len = seq_len, + }; + if (seq_len <= topk) { + return trivial_transform(problem, transform); + } + + // Round the window down to a `kVecSize` boundary and mask the <= 3 columns + // that pulls in, with the same bias / input_start as `topk_ragged_kernel`. + const auto rem = row_start % kVecSize; + if (rem != 0) { + // The mask has to land after the indexer has retired + // Otherwise it may be accidentally overwritten by DG upstream + device::PDLWaitPrimary(); + static_assert(kVecSize <= kBlockSize, "not enough threads "); + if (const auto tx = threadIdx.x; tx < rem) { + score[row_start - rem + tx] = impl::padding_value(); + } + } + using device::topk::broadcast; + problem.in -= rem; + problem.out = s_topk_indices; // write into stage buffer in smem first + problem.seq_len = seq_len + rem; + problem.bias = broadcast(-static_cast(rem)); + problem.input_start = broadcast(rem); + + if (problem.seq_len <= Register2::kMaxSeqLen) { + Register2::forward(problem, &smem); + } else if (problem.seq_len <= Register4::kMaxSeqLen) { + Register4::forward(problem, &smem); + } else { + Streaming::forward(problem, &smem); + } + device::PDLTriggerSecondary(); + __syncthreads(); + paged_transform(problem, out, transform); +} +#endif // USE_ROCM + /** * \brief Main kernel for the short items and epilogue of long items. * \tparam kPDL whether to use PDL to synchronize with the cluster kernel (if any) @@ -786,6 +883,94 @@ struct TopKKernel { .config({.use_pdl = kUsePDL}) .launch(topk_ragged_kernel, params); } + +#ifdef USE_ROCM // see the packed kernel above + /** + * \brief Packed (DSA extend prefill) variant of `transform_paged`: per-row + * window inside one batch-global score buffer, page-table output, no plan. + * + * `scores` is written in place exactly like `transform_ragged` does, so rows + * must not overlap and the buffer must have no consumer after this call. + * + * `row_to_batch` absent means the page table is indexed by score row; present, + * it maps each score row onto the table row of the request it belongs to, + * which is what prefill needs (one request expands into many query rows). + */ + static void transform_packed( + const tvm::ffi::TensorView scores, + const tvm::ffi::TensorView seq_lens, + const tvm::ffi::TensorView row_starts, + const tvm::ffi::TensorView page_table, + const tvm::ffi::TensorView page_indices, + const uint32_t page_size, + const tvm::ffi::Optional row_to_batch) { + using namespace host; + auto B = SymbolicSize{"batch_size"}; + auto L = SymbolicSize{"max_seq_len"}; + auto S = SymbolicSize{"score_stride"}; + auto R = SymbolicSize{"page_table_rows"}; + auto K = SymbolicSize{"topk"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + TensorMatcher({B, L}) // score + .with_strides({S, 1}) + .with_dtype() + .with_device(device_) + .verify(scores); + TensorMatcher({B}) // seq_lens + .with_dtype() + .with_device(device_) + .verify(seq_lens); + TensorMatcher({B}) // row_starts + .with_dtype() + .with_device(device_) + .verify(row_starts); + TensorMatcher({R, -1}) // page_table + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(page_table); + TensorMatcher({B, K}) // page_indices + .with_dtype() + .with_device(device_) + .verify(page_indices); + const int32_t* row_to_batch_ptr = nullptr; + if (row_to_batch.has_value()) { + TensorMatcher({B}) // row_to_batch + .with_dtype() + .with_device(device_) + .verify(row_to_batch.value()); + row_to_batch_ptr = static_cast(row_to_batch.value().data_ptr()); + } else { + RuntimeCheck(R.unwrap() == B.unwrap(), "page_table must have one row per score row unless row_to_batch is given"); + } + + RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2"); + RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)"); + // The kernel masks the head of each window in place, so overlapping rows + // would clobber each other. + RuntimeCheck(S.unwrap() >= L.unwrap(), "scores rows must not overlap"); + const auto topk = static_cast(K.unwrap()); + RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]"); + + const auto params = TopKPackedParams{ + .scores = static_cast(scores.data_ptr()), + .seq_lens = static_cast(seq_lens.data_ptr()), + .row_starts = static_cast(row_starts.data_ptr()), + .row_to_batch = row_to_batch_ptr, + .page_table = static_cast(page_table.data_ptr()), + .page_indices = static_cast(page_indices.data_ptr()), + .score_stride = S.unwrap(), + .page_table_stride = page_table.stride(0), + .topk = topk, + .page_bits = static_cast(std::countr_zero(page_size)), + }; + LaunchKernel(static_cast(B.unwrap()), kBlockSize, device_.unwrap()) + .config({.use_pdl = kUsePDL}) + .launch(topk_packed_kernel, params); + } +#endif // USE_ROCM }; } // namespace sglang diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index d40391655..74128f2b0 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -49,16 +49,20 @@ def _jit_topk_v2_module(): f"-DSGL_TOPK_V2_MAX_C{cluster_size}_OCC{occupancy}={max_active_clusters}" ) kernel = f"TopKKernel<{args}>" + wrappers = [ + ("topk_transform_paged", f"{kernel}::transform_paged"), + ("topk_transform_ragged", f"{kernel}::transform_ragged"), + ("topk_plan", f"{kernel}::plan"), + ] + if is_hip_runtime(): + # transform_packed only exists under USE_ROCM, see topk_v2.cuh + wrappers.append(("topk_transform_packed", f"{kernel}::transform_packed")) return load_jit( make_name("topk_v2"), *args, extra_cuda_cflags=extra_cuda_cflags, cuda_files=["deepseek_v4/topk_v2.cuh"], - cuda_wrappers=[ - ("topk_transform_paged", f"{kernel}::transform_paged"), - ("topk_transform_ragged", f"{kernel}::transform_ragged"), - ("topk_plan", f"{kernel}::plan"), - ], + cuda_wrappers=wrappers, ) @@ -212,6 +216,9 @@ def topk_transform_paged_v2( * Both outputs given -- ``out_page_indices`` receives the page-table transform and ``out_raw_indices`` receives the selected raw indices. + For the packed (DSA extend prefill) layout see + :func:`topk_transform_packed_v2`. + NOTE: every entry of `seq_lens` must be NON-NEGATIVE, and `metadata` must come from :func:`plan_topk_v2` over the same `seq_lens` values. A length of 0 is the valid way to express "no tokens": the row takes the @@ -247,3 +254,50 @@ def topk_transform_paged_v2( metadata, out_raw_indices, ) + + +def topk_transform_packed_v2( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + *, + row_starts: torch.Tensor, + row_to_batch: Optional[torch.Tensor] = None, +) -> None: + """Packed (DSA extend prefill) fused top-k + page-table transform. + + Row ``i`` selects the top-k of ``scores[i, ks : ks + seq_lens[i]]`` + (``ks = row_starts[i]``) and writes the page-table transform of the selected + row-local positions into ``out_page_indices``, ``-1`` padded. Prefill expands + one request into many query-token rows, so ``row_to_batch[i]`` (optional, + ``(rows,)`` int32) names the ``page_tables`` row of the request row ``i`` + belongs to; omitting it indexes the table by score row. ``row_to_batch`` is + not range-checked. + + This is :func:`topk_transform_ragged_v2` with a page-table output instead of + an additive offset. Like ragged, it dispatches the implementation per row at + runtime, so it needs no plan and no :func:`plan_topk_v2` metadata. + + NOTE: ``scores`` is MODIFIED IN PLACE -- the <= 3 columns ahead of each row's + window that the 16-byte-aligned read base pulls in are masked out. They are + invalid for that row and the buffer must have no other consumer, so do not + pass a view with overlapping rows. + ``seq_lens`` entries must be NON-NEGATIVE, as for the paged entry point. + + ROCm only: the kernel is compiled under ``USE_ROCM`` so that CUDA and XPU + builds are untouched. Nothing in it is AMD-specific -- no non-ROCm caller + produces this layout today. + """ + assert is_hip_runtime(), "topk_transform_packed_v2 is compiled under USE_ROCM only" + module = _jit_topk_v2_module() + module.topk_transform_packed( + scores, + seq_lens, + row_starts, + page_tables, + out_page_indices, + page_size, + row_to_batch, + ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py index 2ff9256f9..7802f9590 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -7,10 +7,13 @@ import torch from sglang.srt.environ import envs from sglang.srt.runtime_context import get_exec, get_spec +from sglang.srt.utils import is_hip if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner +_is_hip = is_hip() + _FLASHINFER_TIE_BREAK_VALUES = { "small": 1, "large": 2, @@ -148,6 +151,39 @@ class DSATopKBackend(Enum): logits, lengths, topk, topk_indices_offset, row_starts ) + # Packed PAGED extend (GLM DSA prefill), ROCm-only: CUDA gets the same + # fusion from RAGGED above. Unsupported shapes fall back, not raise. + # The row -> request map is `token_to_batch_idx` for a whole-forward call + # and the chunk's own `batch_idx_list` when the indexer split the logits. + if batch_idx_list is None: + row_to_batch = attn_metadata.token_to_batch_idx + elif isinstance(batch_idx_list, torch.Tensor): + row_to_batch = batch_idx_list + else: + # The prefill-CP list selects requests, not rows: leave it on legacy. + row_to_batch = None + if ( + _is_hip + and self.should_use_topk_v2() + and topk_transform_method == TopkTransformMethod.PAGED + and 0 < topk <= 2048 + and lengths.shape[0] == logits.shape[0] + and logits.dtype == torch.float32 + and logits.stride(1) == 1 + and logits.stride(0) % 4 == 0 + and row_starts is not None + and row_to_batch is not None + and row_to_batch.shape[0] == logits.shape[0] + ): + return _topk_transform_v2_packed( + logits, + lengths, + topk, + attn_metadata, + row_starts=row_starts, + row_to_batch=row_to_batch, + ) + # The legacy transforms below read attn_metadata.page_table_1 (page_size=1), # which is always present here: the fold only drops it for the decode case # dispatched to v2 above. @@ -291,6 +327,9 @@ def _topk_transform_v2_paged( typically 64) yields the same physical slots as gathering the page_size=1 table, without materializing that wide table. + For DSA extend's packed batch-global scores see + :func:`_topk_transform_v2_packed`. + This is a committed contract, not a best-effort path: ``topk_transform`` routes here only for the decode-shaped PAGED case, and the fused-decode CUDA graph drops the page_size=1 table for exactly this case (see @@ -341,6 +380,54 @@ def _topk_transform_v2_paged( return out +def _topk_transform_v2_packed( + logits: torch.Tensor, + lengths: torch.Tensor, + topk: int, + attn_metadata, + row_starts: torch.Tensor, + row_to_batch: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Fused packed-row top-k + page-table transform (DSA extend prefill). + + Same output contract as :func:`_topk_transform_v2_paged` -- ``(num_rows, + topk)`` int32 physical KV slots, ``-1`` padded -- but the scores are packed: + row ``i`` owns the window at ``row_starts[i]`` of one batch-global buffer and + maps through page-table row ``row_to_batch[i]`` (prefill expands one request + into many query-token rows). Selected indices stay row-local. + + Being a prefill-only path it dispatches per row inside the kernel, so unlike + the paged entry point it needs no ``topk_v2_plan``. + + NOTE: ``logits`` is MODIFIED IN PLACE (the <= 3 columns ahead of each window + are masked); the caller must not reuse it. ``lengths`` must be NON-NEGATIVE, + for the same reason as in :func:`_topk_transform_v2_paged`. + """ + from sglang.kernels.ops.attention.dsv4.topk import topk_transform_packed_v2 + + num_rows = logits.shape[0] + assert ( + logits.dtype == torch.float32 + and logits.stride(1) == 1 + and logits.stride(0) % 4 == 0 + ), ( + f"v2 top-k expects fp32 scores with unit row stride and 16B-aligned score_stride, got {logits.dtype=} {logits.stride()=}" + ) + assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}" + + out = logits.new_empty((num_rows, topk), dtype=torch.int32) + topk_transform_packed_v2( + logits, + lengths, + attn_metadata.real_page_table, + out, + attn_metadata.page_size, + row_starts=row_starts.to(torch.int32), + row_to_batch=(None if row_to_batch is None else row_to_batch.to(torch.int32)), + ) + return out + + def _topk_transform_v2_ragged( logits: torch.Tensor, lengths: torch.Tensor, diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 8fdd7ba12..c40c08572 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -747,9 +747,9 @@ class DeepseekSparseAttnBackend( # Preprocess the folded top-k v2 plan once per forward (shared across # layers), at metadata-build time, from the same seqlens the transform # receives as `lengths` (dsa_seqlens_expanded). This must cover EVERY shape - # that dispatches to `_topk_transform_v2_paged` -- decode AND MTP - # target-verify / draft-extend, whose expanded row count is exactly what v2 - # sees -- otherwise the helper's plan-present assertion fires. None only + # that dispatches to `_topk_transform_v2_paged` -- decode, MTP target-verify + # / draft-extend, and packed PAGED extend, whose expanded row count is what + # v2 sees -- otherwise the helper's plan-present assertion fires. None only # when the SGL v2 path is disabled; such metadata is never dispatched to v2. if not self.dsa_topk_backend.should_use_topk_v2(): return None diff --git a/test/registered/kernels/ops/attention/test_topk_v2.py b/test/registered/kernels/ops/attention/test_topk_v2.py index c1fa207b4..3869c4d28 100644 --- a/test/registered/kernels/ops/attention/test_topk_v2.py +++ b/test/registered/kernels/ops/attention/test_topk_v2.py @@ -22,6 +22,9 @@ cluster floor and pool size are per-arch (see topk_v2.cuh), so the (batch, seq) grid below brackets the fixed boundaries (8192/8193, 16384/16385) exactly and spans the arch-dependent ones, across k in {512,1024,2048} and identity/perm page tables. + +``test_topk_v2_packed_rows`` covers the DSA extend layout on top of that: all +requests packed into one score buffer, sharing a table row per request. """ from __future__ import annotations @@ -33,9 +36,11 @@ import torch from sglang.kernels.ops.attention.dsv4.topk import ( plan_topk_v2, + topk_transform_packed_v2, topk_transform_paged_v2, topk_transform_ragged_v2, ) +from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -465,5 +470,143 @@ def test_topk_v2_ragged_no_row_starts(k: int) -> None: assert sorted(explicit[i]) == sorted(implicit[i]), f"row {i} differs" +@pytest.mark.skipif( + not is_hip(), reason="packed layout is compiled under USE_ROCM only" +) +@pytest.mark.parametrize("k", [512, 2048]) +@pytest.mark.parametrize( + "extend_lens", + [ + [7], # one request + [4, 4], # equal row counts + [1, 13, 2], # ragged, including a single-row request + ], +) +@torch.inference_mode() +def test_topk_v2_packed_rows(extend_lens: list[int], k: int) -> None: + """DSA extend layout: batch-global packed scores + shared page-table rows. + + Rows are causal within a request; a distinct page-table permutation per request + catches row/request index mix-ups, and the ragged case leaves most window starts + off the 16-byte load boundary (the production case). + """ + torch.manual_seed(4242 + k + len(extend_lens)) + device = "cuda" + + # Keep every row longer than k so no row takes the trivial path. + prefix = k + 1024 + kv_lens = [prefix + e for e in extend_lens] + k_offsets = [0] + for kv in kv_lens[:-1]: + k_offsets.append(k_offsets[-1] + kv) + total_kv = sum(kv_lens) + + row_starts, lengths, row_to_batch = [], [], [] + for i, e in enumerate(extend_lens): + for local in range(e): + row_starts.append(k_offsets[i]) + lengths.append(kv_lens[i] - e + local + 1) + row_to_batch.append(i) + rows = len(lengths) + + width = (total_kv + 3) & ~3 + scores = torch.randn(rows, width, dtype=torch.float32, device=device)[:, :total_kv] + lengths_t = torch.tensor(lengths, dtype=torch.int32, device=device) + row_starts_t = torch.tensor(row_starts, dtype=torch.int32, device=device) + row_to_batch_t = torch.tensor(row_to_batch, dtype=torch.int32, device=device) + + num_pages = (max(kv_lens) + PAGE_SIZE - 1) // PAGE_SIZE + page_table, inv_cpu = _make_page_table( + len(extend_lens), num_pages, "perm", device, per_row=True + ) + + out = torch.full((rows, k), -1, dtype=torch.int32, device=device) + # The kernel masks in place, so reference values must be read before the call. + scores_cpu = scores.cpu() + topk_transform_packed_v2( + scores, + lengths_t, + page_table, + out, + PAGE_SIZE, + row_starts=row_starts_t, + row_to_batch=row_to_batch_t, + ) + torch.cuda.synchronize() + + out_cpu = out.cpu().tolist() + for r in range(rows): + L, start, req = lengths[r], row_starts[r], row_to_batch[r] + window = scores_cpu[r, start : start + L] + ref = torch.topk(window, k, sorted=False).indices.tolist() + our = _invert(out_cpu[r], inv_cpu[req]) + _assert_topk_close(window.unsqueeze(0), [ref], [our], 1, [L], k) + + +@pytest.mark.skipif( + not is_hip(), reason="packed layout is compiled under USE_ROCM only" +) +@pytest.mark.parametrize("residue", [1, 2, 3]) +@pytest.mark.parametrize("boundary", [8192, 16384]) +@torch.inference_mode() +def test_topk_v2_packed_level_boundary(boundary: int, residue: int) -> None: + """Rows whose length sits on an implementation's max_seq_len boundary. + + The masked head widens the problem by ``residue``, so a row of exactly + ``boundary`` tokens spills past the register implementation sized for it. The + packed kernel picks the implementation per row from the widened length, so + these must still be exact; a compile-time choice made from the un-widened + length would overflow. + """ + torch.manual_seed(boundary + residue) + device = "cuda" + k = 512 + + # Request 0 exists only to push request 1's window off the 16-byte boundary. + kv_lens = [residue, boundary + 1] + lengths = [residue] + [boundary - 1, boundary, boundary + 1] + row_starts = [0] + [residue] * 3 + row_to_batch = [0, 1, 1, 1] + rows = len(lengths) + total_kv = sum(kv_lens) + + width = (total_kv + 3) & ~3 + scores = torch.randn(rows, width, dtype=torch.float32, device=device)[:, :total_kv] + lengths_t = torch.tensor(lengths, dtype=torch.int32, device=device) + row_starts_t = torch.tensor(row_starts, dtype=torch.int32, device=device) + row_to_batch_t = torch.tensor(row_to_batch, dtype=torch.int32, device=device) + + num_pages = (max(kv_lens) + PAGE_SIZE - 1) // PAGE_SIZE + page_table, inv_cpu = _make_page_table( + len(kv_lens), num_pages, "perm", device, per_row=True + ) + + out = torch.full((rows, k), -1, dtype=torch.int32, device=device) + scores_cpu = scores.cpu() + topk_transform_packed_v2( + scores, + lengths_t, + page_table, + out, + PAGE_SIZE, + row_starts=row_starts_t, + row_to_batch=row_to_batch_t, + ) + torch.cuda.synchronize() + + out_cpu = out.cpu().tolist() + for r in range(rows): + L, start, req = lengths[r], row_starts[r], row_to_batch[r] + window = scores_cpu[r, start : start + L] + our = _invert(out_cpu[r], inv_cpu[req]) + if L <= k: + # Trivial path: every position, then -1 padding. + assert sorted(our[:L]) == list(range(L)), f"row {r} trivial output wrong" + assert all(v == -1 for v in out_cpu[r][L:]), f"row {r} padding wrong" + continue + ref = torch.topk(window, k, sorted=False).indices.tolist() + _assert_topk_close(window.unsqueeze(0), [ref], [our], 1, [L], k) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"]))