[DSA] Trim top-k v2 output modes and tighten its PDL waits (#35041)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
co-authored by
Claude Opus 5
Claude Opus 5
parent
1c82955861
commit
746418a1ec
@@ -28,6 +28,11 @@ namespace sglang {
|
||||
namespace impl = device::topk;
|
||||
using impl::TopKProblem;
|
||||
|
||||
enum class TopKMode {
|
||||
INDICES, ///< raw selected indices into `out`; `page_table` unused
|
||||
PAGE_TABLE, ///< page-table-transformed indices into `out`
|
||||
};
|
||||
|
||||
using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read
|
||||
using Register4 = impl::TopKRegister<4>; // <= 16384, register-resident, 1 read
|
||||
using Streaming = impl::TopKStreaming;
|
||||
@@ -64,7 +69,6 @@ struct TopKLaunchParams {
|
||||
const int32_t* __restrict__ seq_lens;
|
||||
const int32_t* __restrict__ page_table;
|
||||
int32_t* __restrict__ page_indices;
|
||||
int32_t* __restrict__ raw_indices; // optional raw (pre-transform) indices output; nullptr if unused
|
||||
const PlanItem* __restrict__ metadata; // [0]=GlobalMetadata, [1+i]=PlanItem
|
||||
int64_t score_stride;
|
||||
int64_t page_table_stride;
|
||||
@@ -89,7 +93,6 @@ struct TopKLaunchParams {
|
||||
return TopKProblem{
|
||||
.in = scores + batch_id * score_stride,
|
||||
.out = page_indices + batch_id * k,
|
||||
.raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
|
||||
.page_table = page_table + batch_id * page_table_stride,
|
||||
.topk = topk,
|
||||
.seq_len = seq_len,
|
||||
@@ -133,12 +136,17 @@ SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
|
||||
}
|
||||
}
|
||||
|
||||
template <bool kPDL>
|
||||
template <bool kPDL, TopKMode kMode>
|
||||
SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
|
||||
device::PDLWaitPrimary<kPDL>();
|
||||
device::PDLTriggerSecondary<kPDL>();
|
||||
for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
|
||||
problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
|
||||
const auto idx = tx < problem.seq_len ? static_cast<int32_t>(tx) : -1;
|
||||
if constexpr (kMode == TopKMode::INDICES) {
|
||||
problem.emit(tx, idx);
|
||||
} else {
|
||||
problem.transform_output(tx, idx);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -160,17 +168,25 @@ SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
|
||||
* - Level 2: max_seq_len <= cluster_floor -> trivial + register<4> + streaming
|
||||
* - Level 3: max_seq_len > cluster_floor -> + epilogue process of cluster path
|
||||
*/
|
||||
template <bool kPDL, int kLevel>
|
||||
template <bool kPDL, int kLevel, TopKMode kMode>
|
||||
TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
|
||||
device::enable_smem_spilling();
|
||||
auto problem = params.problem(blockIdx.x);
|
||||
constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
|
||||
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
|
||||
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
|
||||
__shared__ int32_t topk_indices[kMaxTopK];
|
||||
problem.out = topk_indices;
|
||||
|
||||
constexpr bool kHandleCluster = (kLevel == 3);
|
||||
// Only the cluster path consumes the cluster kernel's output, so only it waits
|
||||
// on that kernel (kPDLFinal). Every other path waits at most on the indexer
|
||||
// (kPDLEarly) and must not be held on an SM slot until the long-running
|
||||
// persistent pool retires -- that would serialize the short items behind it.
|
||||
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
|
||||
constexpr bool kPDLFinal = kPDL && kHandleCluster;
|
||||
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
|
||||
if (problem.seq_len <= problem.topk) return trivial_transform<kPDLEarly, kMode>(problem);
|
||||
|
||||
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
|
||||
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
|
||||
if constexpr (kNeedStaging) problem.out = s_topk_indices;
|
||||
|
||||
// non-trivial path: dispatch based on level and seq_len
|
||||
const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
|
||||
if constexpr (kLevel == 0) {
|
||||
@@ -181,34 +197,35 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams param
|
||||
Register4::forward<kPDL>(problem, &smem); // max_seq_len <= 16384 guarantees seq <= 16384
|
||||
} else {
|
||||
static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
|
||||
// if using cluster, we can delay the PDL wait
|
||||
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
|
||||
constexpr bool kPDLFinal = kPDL && kHandleCluster;
|
||||
if (problem.seq_len <= kReg4MaxSeqLen) {
|
||||
Register4::forward<kPDLEarly>(problem, &smem);
|
||||
} else if (problem.seq_len <= cluster_threshold) {
|
||||
Streaming::forward<kPDLEarly>(problem, &smem);
|
||||
} else { // cluster path do nothing here
|
||||
} else {
|
||||
// Cluster path: the pool already selected into our output row; the only
|
||||
// work left is the epilogue, so this is the one path that waits for it.
|
||||
problem.out = params.get_output_ptr(blockIdx.x);
|
||||
device::PDLWaitPrimary<kPDLFinal>();
|
||||
}
|
||||
device::PDLWaitPrimary<kPDLFinal>();
|
||||
}
|
||||
|
||||
// page-table transform pass (gathers kept out of the hot scatter loop),
|
||||
// then trigger the dependent kernel only after the full output is written.
|
||||
device::PDLTriggerSecondary<kPDL>();
|
||||
__syncthreads();
|
||||
problem_transform(problem, params.get_output_ptr(blockIdx.x));
|
||||
if constexpr (kNeedStaging) {
|
||||
__syncthreads();
|
||||
problem_transform(problem, params.get_output_ptr(blockIdx.x));
|
||||
}
|
||||
}
|
||||
|
||||
template <bool kPDL>
|
||||
template <bool kPDL, TopKMode kMode>
|
||||
CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
|
||||
device::enable_smem_spilling();
|
||||
auto problem = params.problem(blockIdx.x);
|
||||
__shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
|
||||
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
|
||||
__shared__ int32_t topk_indices[kMaxTopK];
|
||||
problem.out = topk_indices;
|
||||
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL, kMode>(problem);
|
||||
|
||||
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
|
||||
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
|
||||
if constexpr (kNeedStaging) problem.out = s_topk_indices;
|
||||
|
||||
// randomly elect one worker rank to avoid workload imbalance
|
||||
const auto worker_rank = blockIdx.x % kClusterSize;
|
||||
@@ -217,30 +234,34 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLau
|
||||
if (problem.seq_len <= kReg4MaxSeqLen) {
|
||||
if (blockIdx.y != worker_rank) return;
|
||||
Register4::forward<kPDL>(problem, &smem);
|
||||
device::PDLWaitPrimary<kPDL>();
|
||||
__syncthreads();
|
||||
} else if (problem.seq_len <= params.cluster_floor) {
|
||||
if (blockIdx.y != worker_rank) return;
|
||||
Streaming::forward<kPDL>(problem, &smem);
|
||||
device::PDLWaitPrimary<kPDL>();
|
||||
__syncthreads();
|
||||
} else {
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
|
||||
Cluster::forward<kPDL>(problem, &smem); // write to peer's output shared memory
|
||||
device::PDLWaitPrimary<kPDL>();
|
||||
cluster.sync();
|
||||
if (blockIdx.y != worker_rank) return;
|
||||
if constexpr (kNeedStaging) {
|
||||
problem.out = cluster.map_shared_rank(s_topk_indices, worker_rank);
|
||||
}
|
||||
Cluster::forward<kPDL>(problem, &smem);
|
||||
if constexpr (kNeedStaging) {
|
||||
cluster.sync();
|
||||
if (blockIdx.y != worker_rank) return;
|
||||
}
|
||||
}
|
||||
|
||||
// Only the elected worker reaches here, and it mapped `topk_indices` to
|
||||
// itself, so `problem.out` is this block's own buffer. Stating that keeps the
|
||||
// shared::cluster address out of the load problem_transform issues -- which is
|
||||
// load-bearing, not an optimization: without it cicc segfaults on CUDA 13.1+
|
||||
// for sm_90a (issue #32830, previously worked around by copying `problem` in
|
||||
// #32910). Verified: dropping this line reproduces the crash on 13.1/13.2/13.3.
|
||||
__builtin_assume(problem.out == topk_indices);
|
||||
problem_transform(problem, params.get_output_ptr(blockIdx.x));
|
||||
device::PDLTriggerSecondary<kPDL>();
|
||||
if constexpr (kNeedStaging) {
|
||||
// Only the elected worker reaches here, and it mapped `topk_indices` to
|
||||
// itself, so `problem.out` is this block's own buffer. Stating that keeps the
|
||||
// shared::cluster address out of the load problem_transform issues -- which is
|
||||
// load-bearing, not an optimization: without it cicc segfaults on CUDA 13.1+
|
||||
// for sm_90a (issue #32830, previously worked around by copying `problem` in
|
||||
// #32910). Verified: dropping this line reproduces the crash on 13.1/13.2/13.3.
|
||||
__builtin_assume(problem.out == s_topk_indices);
|
||||
problem_transform(problem, params.get_output_ptr(blockIdx.x));
|
||||
}
|
||||
}
|
||||
|
||||
// --- Plan: choose cluster_threshold from the seq_len distribution -----------
|
||||
@@ -360,11 +381,10 @@ struct TopKKernel {
|
||||
static void transform(
|
||||
const tvm::ffi::TensorView scores,
|
||||
const tvm::ffi::TensorView seq_lens,
|
||||
const tvm::ffi::TensorView page_table,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> page_table,
|
||||
const tvm::ffi::TensorView page_indices,
|
||||
const uint32_t page_size,
|
||||
const tvm::ffi::TensorView metadata,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
|
||||
const tvm::ffi::TensorView metadata) {
|
||||
using namespace host;
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
|
||||
@@ -384,11 +404,19 @@ struct TopKKernel {
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(seq_lens);
|
||||
TensorMatcher({B, -1}) // page_table
|
||||
.with_strides({P, 1})
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(page_table);
|
||||
// Absent means "no page transform": `page_indices` then receives the raw
|
||||
// selected indices and nothing dereferences a page table.
|
||||
const int32_t* page_table_ptr = nullptr;
|
||||
int64_t page_table_stride = 0;
|
||||
if (page_table.has_value()) {
|
||||
TensorMatcher({B, -1}) // page_table
|
||||
.with_strides({P, 1})
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(page_table.value());
|
||||
page_table_ptr = static_cast<const int32_t*>(page_table.value().data_ptr());
|
||||
page_table_stride = P.unwrap();
|
||||
}
|
||||
TensorMatcher({B, K}) // page_indices
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
@@ -398,12 +426,6 @@ struct TopKKernel {
|
||||
.with_device(device_)
|
||||
.verify(metadata);
|
||||
|
||||
int32_t* raw_indices_ptr = nullptr;
|
||||
if (raw_indices.has_value()) {
|
||||
TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
|
||||
raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
|
||||
}
|
||||
|
||||
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)");
|
||||
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
|
||||
@@ -425,12 +447,11 @@ struct TopKKernel {
|
||||
const auto params = TopKLaunchParams{
|
||||
.scores = static_cast<const float*>(scores.data_ptr()),
|
||||
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
|
||||
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
|
||||
.page_table = page_table_ptr,
|
||||
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
|
||||
.raw_indices = raw_indices_ptr,
|
||||
.metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
|
||||
.score_stride = S.unwrap(),
|
||||
.page_table_stride = P.unwrap(),
|
||||
.page_table_stride = page_table_stride,
|
||||
.topk = topk,
|
||||
.page_bits = page_bits,
|
||||
.cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
|
||||
@@ -438,33 +459,44 @@ struct TopKKernel {
|
||||
|
||||
const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
|
||||
constexpr bool kUsePDL = true;
|
||||
if (use_cluster) {
|
||||
if (batch_size <= kNumPersistentClusters) {
|
||||
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
|
||||
.launch(topk_small_batch_kernel<kUsePDL>, params);
|
||||
} else {
|
||||
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
|
||||
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
|
||||
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
|
||||
const auto mode = page_table.has_value() ? TopKMode::PAGE_TABLE : TopKMode::INDICES;
|
||||
const auto dispatch = [&]<typename F>(F&& f) {
|
||||
switch (mode) {
|
||||
case TopKMode::INDICES:
|
||||
return f.template operator()<TopKMode::INDICES>();
|
||||
default:
|
||||
return f.template operator()<TopKMode::PAGE_TABLE>();
|
||||
}
|
||||
};
|
||||
dispatch([&]<TopKMode kMode>() {
|
||||
if (use_cluster) {
|
||||
if (batch_size <= kNumPersistentClusters) {
|
||||
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
|
||||
.launch(topk_small_batch_kernel<kUsePDL, kMode>, params);
|
||||
} else {
|
||||
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
|
||||
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
|
||||
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3, kMode>, params);
|
||||
}
|
||||
} else if (max_seq_len <= kReg2MaxSeqLen) {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0, kMode>, params);
|
||||
} else if (max_seq_len <= kReg4MaxSeqLen) {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1, kMode>, params);
|
||||
} else {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2, kMode>, params);
|
||||
}
|
||||
} else if (max_seq_len <= kReg2MaxSeqLen) {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
|
||||
} else if (max_seq_len <= kReg4MaxSeqLen) {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
|
||||
} else {
|
||||
LaunchKernel(batch_size, kBlockSize, device)
|
||||
.config({.use_pdl = kUsePDL})
|
||||
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -168,25 +168,19 @@ SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint3
|
||||
|
||||
/// One batch element's worth of work. `emit(pos, raw_idx)` writes the selected raw
|
||||
/// index to output slot `pos`; `transform_output` then applies the page-table
|
||||
/// transform in a separate pass (and records the raw index in `raw_out` if set).
|
||||
/// transform in a separate pass.
|
||||
struct TopKProblem {
|
||||
const float* __restrict__ in;
|
||||
int32_t* __restrict__ out; // page_indices [topk]
|
||||
int32_t* __restrict__ raw_out; // optional raw (pre-transform) indices [topk]; nullptr if unused
|
||||
int32_t* __restrict__ out; // page_indices [topk]
|
||||
const int32_t* __restrict__ page_table;
|
||||
uint32_t topk;
|
||||
uint32_t seq_len;
|
||||
uint32_t page_bits;
|
||||
|
||||
// Write the raw selected index; the page-table transform is applied afterwards
|
||||
// by transform_output() in a separate, pipelined pass. Keeping the per-element
|
||||
// page_table gather off the atomic-serialized scatter loop is measurably faster
|
||||
// for both short and long context.
|
||||
SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
|
||||
out[pos] = static_cast<int32_t>(raw_idx);
|
||||
}
|
||||
SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
|
||||
if (raw_out != nullptr) raw_out[t] = raw;
|
||||
out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -89,13 +89,21 @@ def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Ten
|
||||
def topk_transform_512_v2(
|
||||
scores: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
page_tables: torch.Tensor,
|
||||
page_tables: Optional[torch.Tensor],
|
||||
out_page_indices: torch.Tensor,
|
||||
page_size: int,
|
||||
metadata: torch.Tensor,
|
||||
out_raw_indices: Optional[torch.Tensor] = None,
|
||||
) -> None:
|
||||
"""Fused top-k + page-table transform (DeepSeek-V4 top-k v2 kernel).
|
||||
"""Fused top-k + optional page-table transform (DeepSeek-V4 top-k v2 kernel).
|
||||
|
||||
Two output modes, chosen by whether ``page_tables`` is given and resolved to
|
||||
a device-side template parameter, so an unused page-table gather is compiled
|
||||
out rather than skipped at runtime:
|
||||
|
||||
* ``page_tables=None`` -- ``out_page_indices`` receives the raw selected
|
||||
indices and no page table is read.
|
||||
* ``page_tables`` given -- ``out_page_indices`` receives the page-table
|
||||
transform of them.
|
||||
|
||||
IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE, and
|
||||
``metadata`` must come from :func:`plan_topk_v2` over the same ``seq_lens``
|
||||
@@ -114,5 +122,4 @@ def topk_transform_512_v2(
|
||||
out_page_indices,
|
||||
page_size,
|
||||
metadata,
|
||||
out_raw_indices,
|
||||
)
|
||||
|
||||
@@ -289,8 +289,6 @@ def _topk_transform_v2_paged(
|
||||
assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}"
|
||||
|
||||
page_table = attn_metadata.real_page_table
|
||||
assert page_table.dtype == torch.int32
|
||||
lengths_i32 = lengths.to(torch.int32)
|
||||
|
||||
# The plan is preprocessed once per forward (DSAMetadata.topk_v2_plan,
|
||||
# refreshed in-place under CUDA graph) and reused across layers. A missing or
|
||||
@@ -302,8 +300,8 @@ def _topk_transform_v2_paged(
|
||||
), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)"
|
||||
|
||||
page_size = attn_metadata.page_size
|
||||
out = logits.new_full((num_rows, topk), -1, dtype=torch.int32)
|
||||
topk_transform_512_v2(logits, lengths_i32, page_table, out, page_size, plan)
|
||||
out = logits.new_empty((num_rows, topk), dtype=torch.int32)
|
||||
topk_transform_512_v2(logits, lengths, page_table, out, page_size, plan)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
@@ -141,28 +141,41 @@ def _reference(scores, seq_lens, k):
|
||||
return ref
|
||||
|
||||
|
||||
def _plan(seq_lens):
|
||||
"""Plan, then break stream adjacency with the transform launch.
|
||||
|
||||
The transform kernel prefetches the plan metadata BEFORE its PDL wait, which
|
||||
is only legal while the plan kernel is not the immediately preceding kernel
|
||||
in the stream -- in production the plan is built during per-forward metadata
|
||||
prep, a whole model forward earlier. Launching the two back to back would
|
||||
read the plan through a programmatic dependency that guarantees no memory
|
||||
visibility, so mirror the production separation instead.
|
||||
"""
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
torch.cuda.synchronize()
|
||||
return metadata
|
||||
|
||||
|
||||
def _run(scores, seq_lens, page_table, inv_cpu, k):
|
||||
batch = scores.shape[0]
|
||||
metadata = _plan(seq_lens)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata)
|
||||
torch.cuda.synchronize()
|
||||
out_cpu = out.cpu().tolist()
|
||||
return [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
|
||||
|
||||
|
||||
def _run_raw(scores, seq_lens, page_table, k):
|
||||
"""Run the kernel and return its optional raw (pre-transform) top-k index
|
||||
output per row, dropping -1 padding -- the selected positions themselves,
|
||||
NOT the page-table transform of them."""
|
||||
def _run_raw(scores, seq_lens, k):
|
||||
"""Run with no page table and return the selected indices per row, dropping
|
||||
-1 padding -- the selected positions themselves, NOT a page transform."""
|
||||
batch = scores.shape[0]
|
||||
metadata = _plan(seq_lens)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
raw = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
|
||||
topk_transform_512_v2(scores, seq_lens, None, out, PAGE_SIZE, metadata)
|
||||
torch.cuda.synchronize()
|
||||
raw_cpu = raw.cpu().tolist()
|
||||
return [[v for v in raw_cpu[i] if v != -1] for i in range(batch)]
|
||||
out_cpu = out.cpu().tolist()
|
||||
return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_mode", ["identity", "perm"])
|
||||
@@ -229,68 +242,25 @@ def test_topk_v2_ragged(batch: int, shape: str, k: int, per_row_pt: bool) -> Non
|
||||
_assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, lengths.cpu(), k)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_mode", ["identity", "perm"])
|
||||
@pytest.mark.parametrize(
|
||||
"batch,seq",
|
||||
[
|
||||
(8, 256), # trivial
|
||||
(8, 4096), # register
|
||||
(4, 131072), # fused small-batch cluster
|
||||
(64, 131072), # persistent cluster + main<3> epilogue
|
||||
(256, 131072), # non-cluster streaming
|
||||
],
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_topk_v2_raw_indices(batch: int, seq: int, page_mode: str) -> None:
|
||||
"""The optional raw-index output must be the pre-transform position of each
|
||||
transformed output slot (out[j] == page_to_indices(raw[j])), and -1 aligns."""
|
||||
k = 512
|
||||
torch.manual_seed(batch * 131 + seq)
|
||||
device = "cuda"
|
||||
width = (seq + 3) & ~3
|
||||
scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
|
||||
seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
|
||||
num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
page_table, inv_cpu = _make_page_table(batch, num_pages, page_mode, device)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=device)
|
||||
raw = torch.full((batch, k), -1, dtype=torch.int32, device=device)
|
||||
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out_cpu, raw_cpu = out.cpu().tolist(), raw.cpu().tolist()
|
||||
for i in range(batch):
|
||||
for j in range(k):
|
||||
o, r = out_cpu[i][j], raw_cpu[i][j]
|
||||
if o == -1:
|
||||
assert r == -1, f"b={i} j={j}: out=-1 but raw={r}"
|
||||
else:
|
||||
inv = (int(inv_cpu[i][o >> PAGE_BITS]) << PAGE_BITS) | (o & PAGE_MASK)
|
||||
assert r == inv, f"b={i} j={j}: raw={r} != inverse(out)={inv}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k", [512, 1024, 2048])
|
||||
@pytest.mark.parametrize("batch,seq", FIXED_CONFIGS)
|
||||
@torch.inference_mode()
|
||||
def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None:
|
||||
"""Validate the raw (pre-transform) index output DIRECTLY against torch.topk.
|
||||
"""Validate the selected indices DIRECTLY against torch.topk.
|
||||
|
||||
Unlike ``test_topk_v2`` -- which checks the page-transformed output and inverts
|
||||
it through the page table -- this exercises the selected indices themselves, so
|
||||
it isolates the top-k selection from the page-table transform. A permuted page
|
||||
table is used so raw != out, catching any bug that leaks transformed page
|
||||
indices into the raw buffer. Covers every dispatch template/boundary.
|
||||
Runs the no-page-table mode, so the output is the selected positions
|
||||
themselves. Unlike ``test_topk_v2`` -- which checks the page-transformed
|
||||
output and inverts it through the page table -- this isolates the top-k
|
||||
selection from the transform, and it is the only coverage of that mode.
|
||||
Covers every dispatch template/boundary.
|
||||
"""
|
||||
torch.manual_seed(batch * 100003 + seq * 7 + k + 1)
|
||||
device = "cuda"
|
||||
width = (seq + 3) & ~3
|
||||
scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
|
||||
seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
|
||||
num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
page_table, _ = _make_page_table(batch, num_pages, "perm", device)
|
||||
|
||||
our_raw = _run_raw(scores, seq_lens, page_table, k)
|
||||
our_raw = _run_raw(scores, seq_lens, k)
|
||||
ref_raw = _reference(scores, seq_lens, k)
|
||||
_assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user