[AMD] Enable GLM DSA prefill top-k to the v2 kernel (#37889)
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
Eric.Chin.AMD
co-authored by
Thomas Wang
Thomas Wang
parent
5a6a1bb883
commit
3c71bb018a
@@ -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<uint32_t>(row_to_batch[bx]);
|
||||
return {page_table + static_cast<int64_t>(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 <bool kPDL>
|
||||
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<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
|
||||
__shared__ int32_t s_topk_indices[kMaxTopK];
|
||||
|
||||
const auto bx = blockIdx.x;
|
||||
const auto seq_len = static_cast<uint32_t>(params.seq_lens[bx]);
|
||||
const auto row_start = static_cast<uint32_t>(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<int64_t>(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<kPDL, TopKMode::PAGE_TABLE>(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<kPDL>();
|
||||
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<int32_t>(rem));
|
||||
problem.input_start = broadcast(rem);
|
||||
|
||||
if (problem.seq_len <= Register2::kMaxSeqLen) {
|
||||
Register2::forward<kPDL>(problem, &smem);
|
||||
} else if (problem.seq_len <= Register4::kMaxSeqLen) {
|
||||
Register4::forward<kPDL>(problem, &smem);
|
||||
} else {
|
||||
Streaming::forward<kPDL>(problem, &smem);
|
||||
}
|
||||
device::PDLTriggerSecondary<kPDL>();
|
||||
__syncthreads();
|
||||
paged_transform<TopKMode::PAGE_TABLE>(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<kUsePDL>, 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<tvm::ffi::TensorView> 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<kDLGPU>();
|
||||
|
||||
TensorMatcher({B, L}) // score
|
||||
.with_strides({S, 1})
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(scores);
|
||||
TensorMatcher({B}) // seq_lens
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(seq_lens);
|
||||
TensorMatcher({B}) // row_starts
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(row_starts);
|
||||
TensorMatcher({R, -1}) // page_table
|
||||
.with_strides({-1, 1})
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(page_table);
|
||||
TensorMatcher({B, K}) // page_indices
|
||||
.with_dtype<int32_t>()
|
||||
.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<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(row_to_batch.value());
|
||||
row_to_batch_ptr = static_cast<const int32_t*>(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<uint32_t>(K.unwrap());
|
||||
RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
|
||||
|
||||
const auto params = TopKPackedParams{
|
||||
.scores = static_cast<float*>(scores.data_ptr()),
|
||||
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
|
||||
.row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
|
||||
.row_to_batch = row_to_batch_ptr,
|
||||
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
|
||||
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
|
||||
.score_stride = S.unwrap(),
|
||||
.page_table_stride = page_table.stride(0),
|
||||
.topk = topk,
|
||||
.page_bits = static_cast<uint32_t>(std::countr_zero(page_size)),
|
||||
};
|
||||
LaunchKernel(static_cast<uint32_t>(B.unwrap()), kBlockSize, device_.unwrap())
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_packed_kernel<kUsePDL>, params);
|
||||
}
|
||||
#endif // USE_ROCM
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]))
|
||||
|
||||
Reference in New Issue
Block a user