iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[TopK] In-kernel canonical emission (MERLIN_TOPK_SORTED_EMIT)

Browse Source 
Emits the selected indices in canonical ascending order inside the kernel so
the Python QSA sort can be skipped. Output-equivalent to the Python sort;
reference model, property tests and GPU arms included. Default off in this
commit.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 20:04:37 +07:00
co-authored by Hinata
parent a4aecf9ec7
commit eeea036c12
5 changed files with 444 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/persistent_topk.cuh
+219 -8
View File
@@ -200,6 +200,17 @@ struct PersistentTopKParams {
// concurrency the atomic exists for.
// Resolved once per process on the host, so it is a launch constant.
bool fault_inject_emit;
// MERLIN_TOPK_SORTED_EMIT: every row leaves
// the kernel in the canonical order MERLIN_QSA_SORT used to impose in
// Python -- ascending index, padding last as -1 -- so the Python sort can be
// dropped. Paths that already emit that order by construction (trivial,
// medium-path restart, deterministic large path) are untouched; the paths that
// place by atomic arrival (decode-buffered, medium-buffered, the legacy
// large tail) re-emit their OWN selected set through sorted_emit_row().
// Selection is never re-derived, so the selected set cannot change.
// Resolved once per process on the host, so it is a launch constant a
// cudagraph bakes in. Default OFF.
bool sorted_emit;
};
// ============================================================================
@@ -272,6 +283,149 @@ enum TopKExactCounter : int {
kNumTopKExactCounters = 11, // total slots: device buffer AND reader output
};
// ============================================================================
// MERLIN_TOPK_SORTED_EMIT -- in-kernel ascending emission of a row's OWN
// selected set. Declared here, above every caller (decode/medium dispatch,
// the legacy large tail, and the standalone fallback kernel), for the
// two-phase-lookup reason given above the counter enum.
//
// CONTRACT. On return, row[0, TopK) is byte-identical to what
// `_qsa_canonicalise` (vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py) makes
// of the row it was handed:
//
// tmp = where(row < 0, INT32_MAX, row); tmp = sort(tmp)
// row = where(tmp == INT32_MAX, -1, tmp)
//
// i.e. the non-negative entries ascending, then -1 for every negative entry
// (and, to the byte, for any entry equal to INT32_MAX -- the Python spelling
// maps those to -1 too). It is a function of the MULTISET of values in the
// row and nothing else, so it cannot change which blocks were selected: the
// selection is whatever the path already wrote, including a truncated or
// racy one.
//
// WHAT WAS TAKEN FROM UPSTREAM #55122 (head b2312b2de7, det_select_row, the
// emission half). The placement idea: a per-index membership flag, a
// BlockScan exclusive prefix over it with BLOCKED ownership (each thread owns
// consecutive indices, so scan order is index order), and a store straight to
// the ascending position -- no atomics on the placement path. WHAT IS NOT
// TAKEN: #55122 re-derives membership with a full-row radix rescan and its
// own lowest-index tie rule, and canonicalises -0.0 to +0.0. Adopting either
// would change the selected set against production (most-recent-wins,
// +0.0 > -0.0), so membership here comes from the path's own emission
// instead: a 32,768-bit shared bitmap, one word per thread. Thread tx owns
// indices [32 tx, 32 tx + 32); its exclusive prefix of popcounts is the
// ascending position of its first set bit.
//
// FAST PATH. Every non-negative entry is < 32,768 (= RADIX_THRESHOLD, so every
// single-CTA path qualifies) and no entry repeats. Cost: one smem word
// cleared and at most TopK/1024 global loads per thread, TopK shared
// atomicOr's with no two on the same bit, one 1,024-wide BlockScan, TopK
// global stores, four block barriers. No global atomics.
//
// FALLBACK. An out-of-range value or a duplicate (bitmap atomicOr returns the
// bit already set) cannot be expressed as a bitmap, so the row is ranked
// instead: rank(i) = #{j : key_j < key_i} + #{j < i : key_j == key_i}, a stable
// counting sort that equals torch.sort on ANY int32 multiset. O(TopK^2 / 1024)
// shared reads per thread; reached only by the legacy large tail (indices past
// 32,768, control arm only) and by rows whose unfilled slots hold stale
// values (buffered overflow with its exact repair switched off).
//
// SHARED MEMORY. Aliases the dynamic window from byte 0. The caller's shared
// state is dead on entry -- the leading barrier makes sure no thread is still
// reading it -- and the trailing barrier keeps the next row's loader off this
// row's bitmap. Size is static_asserted against kSmemMedium, the floor every
// persistent_topk launch allocates.
// ============================================================================
constexpr uint32_t kSortedEmitBitmapWords = static_cast<uint32_t>(kThreadsPerBlock);
constexpr uint32_t kSortedEmitBitmapBits = 32u * kSortedEmitBitmapWords;
static_assert(kSortedEmitBitmapWords == 1024u,
"one bitmap word per thread: the scan below is one item per thread");
static_assert(RADIX_THRESHOLD <= kSortedEmitBitmapBits,
"every single-CTA path's indices must fit the fast-path bitmap");
template <int TopK>
struct SortedEmitSmem {
typename cub::BlockScan<uint32_t, kThreadsPerBlock>::TempStorage scan;
uint32_t bitmap[kSortedEmitBitmapWords];
int32_t vals[TopK];
uint32_t fallback;
};
template <int TopK>
__host__ __device__ constexpr size_t sorted_emit_smem_bytes() {
return (sizeof(SortedEmitSmem<TopK>) + 15) & ~size_t(15);
}
template <int TopK>
__device__ __noinline__ void sorted_emit_row(int32_t* __restrict__ row) {
static_assert(sorted_emit_smem_bytes<TopK>() <= kSmemMedium,
"sorted emission must fit the smallest persistent_topk window");
static_assert(alignof(SortedEmitSmem<TopK>) <= 16,
"SortedEmitSmem over-aligned for the dynamic window");
using BlockScan = cub::BlockScan<uint32_t, kThreadsPerBlock>;
extern __shared__ __align__(16) char sorted_emit_smem[];
auto* s = reinterpret_cast<SortedEmitSmem<TopK>*>(sorted_emit_smem);
const uint32_t tx = threadIdx.x;
// E0: every emission store of the path is issued and no thread still reads
// the path's shared state (block-scope ordering of the global stores too).
__syncthreads();
s->bitmap[tx] = 0u;
if (tx == 0) s->fallback = 0u;
__syncthreads();
// E1: membership. Reads of `row` all happen before E2; writes all after.
for (uint32_t i = tx; i < static_cast<uint32_t>(TopK); i += kThreadsPerBlock) {
const int32_t v = row[i];
s->vals[i] = v;
if (v >= 0) {
const uint32_t u = static_cast<uint32_t>(v);
if (u < kSortedEmitBitmapBits) {
const uint32_t bit = 1u << (u & 31u);
if ((atomicOr(&s->bitmap[u >> 5], bit) & bit) != 0u) {
atomicOr(&s->fallback, 1u); // duplicate value
}
} else {
atomicOr(&s->fallback, 1u); // beyond the bitmap
}
}
}
__syncthreads(); // E2
if (s->fallback == 0u) { // block-uniform: written only before E2
uint32_t word = s->bitmap[tx];
uint32_t pos = 0u;
uint32_t total = 0u;
BlockScan(s->scan).ExclusiveSum(static_cast<uint32_t>(__popc(word)), pos,
total);
while (word != 0u) {
const uint32_t b = static_cast<uint32_t>(__ffs(word) - 1);
row[pos++] = static_cast<int32_t>((tx << 5) | b);
word &= word - 1u;
}
for (uint32_t i = total + tx; i < static_cast<uint32_t>(TopK);
i += kThreadsPerBlock) {
row[i] = -1;
}
} else {
constexpr int32_t kBig = 0x7FFFFFFF;
for (uint32_t i = tx; i < static_cast<uint32_t>(TopK);
i += kThreadsPerBlock) {
const int32_t vi = s->vals[i];
const int32_t ki = vi < 0 ? kBig : vi;
uint32_t rank = 0u;
for (uint32_t j = 0; j < static_cast<uint32_t>(TopK); ++j) {
const int32_t vj = s->vals[j];
const int32_t kj = vj < 0 ? kBig : vj;
rank += (kj < ki || (kj == ki && j < i)) ? 1u : 0u;
}
row[rank] = (ki == kBig) ? -1 : ki;
}
}
// E3: the next row (or the next path in this row) reuses the window.
__syncthreads();
}
// ============================================================================
// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192)
// Uses 11-bit half-precision bins for fine granularity.
@@ -1081,8 +1235,13 @@ __device__ __noinline__ void medium_exact_restart(
// by: DarkSharpness
// which at the same time is an optimized topk kernel copied from tilelang
// kernel
//
// Returns TRUE when the row was handed to medium_exact_restart(), whose
// emission is already ascending (MERLIN_TOPK_SORTED_EMIT skips those rows);
// FALSE when the buffered path placed the row by atomic arrival. The return
// value is block-uniform (every return below is taken by the whole block).
template <int TopK>
__device__ __noinline__ void histogram_256_topk(
__device__ __noinline__ bool histogram_256_topk(
const float* __restrict__ logits, int* __restrict__ output_indices,
int logits_offset, int seq_len, bool exact_medium,
unsigned long long* __restrict__ counters, bool det_ties,
@@ -1181,7 +1340,7 @@ __device__ __noinline__ void histogram_256_topk(
}
}
__syncthreads();
return;
return false;
}
// Past the early return the threshold bin WILL be collected, so the M1
@@ -1262,7 +1421,7 @@ __device__ __noinline__ void histogram_256_topk(
medium_exact_restart<TopK>(logits + logits_offset, output_indices,
static_cast<uint32_t>(seq_len), counters,
effective_tie_recent(det_ties, tie_recent));
return;
return true;
}
__syncthreads();
}
@@ -1361,6 +1520,7 @@ __device__ __noinline__ void histogram_256_topk(
}
__syncthreads();
}
return false;
}
// ============================================================================
@@ -1949,7 +2109,8 @@ __device__ void radix_topk(const float* __restrict__ row_input,
uint32_t ctas_per_group, int& barrier_phase,
uint32_t radix_iter, uint32_t tx, bool det_large,
unsigned long long* counters, int fault_inject_cta,
bool fault_inject_emit, bool tie_recent) {
bool fault_inject_emit, bool tie_recent,
bool sorted_emit) {
const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
? my_chunk_start + chunk_size
: seq_len;
@@ -2186,6 +2347,30 @@ __device__ void radix_topk(const float* __restrict__ row_input,
// the MERLIN_TOPK_DET_LARGE=0 control, so the control now carries this
// repair too, like the histogram-publication barrier (an earlier commit).
__syncthreads();
// MERLIN_TOPK_SORTED_EMIT on the legacy tail (MERLIN_TOPK_DET_LARGE=0, or
// the host's staging-capacity gate): the row was written by EVERY CTA of the
// group, by arrival order, so CTA 0 may only re-emit it after all of them
// have finished. One more inter-CTA barrier -- the same red_release /
// wait_ge pair as every other barrier here, taken by every CTA of the group
// because `sorted_emit` is a launch constant, so barrier_phase stays
// group-uniform. The closing __syncthreads() just above orders each CTA's
// stores before its thread 0's release fence. Indices here run past
// 32,768, so sorted_emit_row() takes its counting-sort fallback: this is the
// control arm, never the production large path (which is the deterministic one and is
// already ascending).
if (sorted_emit) {
if (tx == 0) {
red_release(&state->arrival_counter, 1);
}
wait_ge(&state->arrival_counter,
(barrier_phase + 1) * static_cast<int>(ctas_per_group), tx);
barrier_phase++;
__syncthreads();
if (cta_in_group == 0) {
sorted_emit_row<TopK>(row_output);
}
}
}
// ============================================================================
@@ -2280,15 +2465,25 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
row_input, row_output, seq_len, params.det_ties,
params.tie_recent, params.exact_counters, params.exact_decode);
if (decode_needs_exact) {
// Already ascending (restart emission): no sorted re-emission.
medium_exact_restart<TopK>(
row_input, row_output, seq_len, params.exact_counters,
effective_tie_recent(params.det_ties, params.tie_recent));
} else if (params.sorted_emit) {
// Decode-buffered placement is by atomic arrival.
sorted_emit_row<TopK>(row_output);
}
} else {
histogram_256_topk<TopK>(row_input, row_output, 0, seq_len,
params.exact_medium, params.exact_counters,
params.det_ties, params.tie_recent);
const bool restarted = histogram_256_topk<TopK>(
row_input, row_output, 0, seq_len, params.exact_medium,
params.exact_counters, params.det_ties, params.tie_recent);
if (!restarted && params.sorted_emit) {
// Medium-buffered placement is by atomic arrival.
sorted_emit_row<TopK>(row_output);
}
}
// The trivial branch above writes [0, seq_len) then -1 in index
// order, which is already the canonical order.
}
continue;
}
@@ -2300,11 +2495,27 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx,
params.det_large, params.exact_counters, params.fault_inject_cta,
params.fault_inject_emit,
effective_tie_recent(params.det_ties, params.tie_recent));
effective_tie_recent(params.det_ties, params.tie_recent),
params.sorted_emit);
radix_iter++;
}
}
// MERLIN_TOPK_SORTED_EMIT for the two host-side fallbacks that bypass
// persistent_topk_kernel (FilteredTopK, needs >= 128 KiB opt-in smem, and
// top_k_per_row_decode, needs a cooperative launch that does not fit). Both
// are unreachable on a 3090 at the shapes Flash-Next serves (the fallbacks
// never dispatch there), but with the knob on the Python sort is gone, so the op's
// contract -- canonical order on EVERY dispatch -- is kept here rather than
// assumed. One block per row, launched on the same stream right after the
// fallback, so it is stream-ordered after the fallback's stores.
template <int TopK>
__global__ void __launch_bounds__(kThreadsPerBlock)
sorted_emit_rows_kernel(int32_t* __restrict__ output, uint32_t num_rows) {
if (blockIdx.x >= num_rows) return;
sorted_emit_row<TopK>(output + static_cast<size_t>(blockIdx.x) * TopK);
}
} // namespace persistent
// ============================================================================
csrc/libtorch_stable/topk.cu
+66 -7
View File
@@ -61,6 +61,7 @@ std::atomic<int> g_cached_tie_det{-1};
std::atomic<int> g_cached_exact_medium{-1};
std::atomic<int> g_cached_det_large{-1};
std::atomic<int> g_cached_exact_decode{-1};
std::atomic<int> g_cached_sorted_emit{-1};
bool publish_cached(std::atomic<int>& mirror, bool value) {
mirror.store(value ? 1 : 0, std::memory_order_release);
@@ -142,6 +143,30 @@ bool exact_decode_enabled() {
return enabled;
}
// ---------------------------------------------------------------------------
// MERLIN_TOPK_SORTED_EMIT: every persistent_topk
// row leaves the op in the canonical order the Python MERLIN_QSA_SORT imposes
// (ascending index, padding last as -1), so qsa_indexer.py can drop that sort.
// See sorted_emit_row() in persistent_topk.cuh for the contract; the selected
// SET is never touched. Resolved once per process like every switch here, so
// it is a launch constant a cudagraph bakes in. Default OFF: the maintainer
// promotes after reference-host qualification.
//
// This std::getenv is the single source of truth. Python does not parse the
// variable to decide anything: qsa_indexer.py asks THIS cache, through
// persistent_topk_toggle_state() field 6, after its first launch -- so the
// Python sort is skipped only when the kernel that ran says it sorted.
// ---------------------------------------------------------------------------
bool sorted_emit_enabled() {
static const bool enabled = [] {
const char* raw = std::getenv("MERLIN_TOPK_SORTED_EMIT");
return publish_cached(
g_cached_sorted_emit,
(raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false);
}();
return enabled;
}
// ---------------------------------------------------------------------------
// MERLIN_TOPK_LARGE_FAULT_INJECT. TEST ONLY.
//
@@ -220,16 +245,21 @@ void report_tie_policy_once() {
stderr,
"[merlin] persistent_topk tie policy: MERLIN_TIE_DET=%d "
"MERLIN_TIE_RECENT=%d MERLIN_TOPK_EXACT_MEDIUM=%d "
"MERLIN_TOPK_EXACT_DECODE=%d MERLIN_TOPK_DET_LARGE=%d -> "
"MERLIN_TOPK_EXACT_DECODE=%d MERLIN_TOPK_DET_LARGE=%d "
"MERLIN_TOPK_SORTED_EMIT=%d -> "
"decode-buffered=%s decode-restart=%s medium-buffered=%s "
"medium-restart=%s large=%s\n",
"medium-restart=%s large=%s emission=%s\n",
det ? 1 : 0, tie_recent_enabled() ? 1 : 0,
exact_medium_enabled() ? 1 : 0, exact_decode_enabled() ? 1 : 0,
det_large_enabled() ? 1 : 0, det ? dir : "atomic-arrival",
det_large_enabled() ? 1 : 0, sorted_emit_enabled() ? 1 : 0,
det ? dir : "atomic-arrival",
exact_decode_enabled() ? dir : "n/a (detect-only; counter [10])",
det ? dir : "atomic-arrival",
exact_medium_enabled() ? dir : "n/a (exact restart disabled)",
det_large_enabled() ? dir : "atomic-arrival (non-deterministic)");
det_large_enabled() ? dir : "atomic-arrival (non-deterministic)",
sorted_emit_enabled()
? "ascending-in-kernel (Python QSA sort skipped)"
: "arrival-order on buffered paths (Python QSA sort decides)");
return true;
}();
(void)reported;
@@ -394,6 +424,19 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
max_smem_per_block = device_prop->sharedMemPerBlockOptin;
}
// MERLIN_TOPK_SORTED_EMIT on the two fallbacks that bypass
// persistent_topk_kernel: re-emit every row ascending on the same stream.
// A plain launch with a fixed configuration, so it is capture-safe.
const auto sorted_emit_fallback_rows = [&]() {
if (!sorted_emit_enabled()) return;
report_tie_policy_once();
P::sorted_emit_rows_kernel<TopK>
<<<static_cast<unsigned>(num_rows), P::kThreadsPerBlock,
P::sorted_emit_smem_bytes<TopK>(), stream>>>(
output.mutable_data_ptr<int32_t>(),
static_cast<uint32_t>(num_rows));
};
// Allow static fallback storage in addition to the 128 KiB dynamic buffer.
if (num_rows > 64 &&
max_seq_len >= vllm::sampled_topk::kMinSampledLength<TopK> &&
@@ -409,6 +452,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
logits.const_data_ptr<float>(), lengths.const_data_ptr<int32_t>(),
output.mutable_data_ptr<int32_t>(), stride,
static_cast<int>(std::min(max_seq_len, logits.size(1))));
sorted_emit_fallback_rows();
} else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
@@ -417,6 +461,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
STD_TORCH_CHECK(status == cudaSuccess,
"FilteredTopK failed: ", cudaGetErrorString(status));
sorted_emit_fallback_rows();
} else {
STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
STD_TORCH_CHECK(
@@ -513,6 +558,10 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1;
top_k_per_row_decode(logits, next_n, lengths, output, num_rows,
logits.stride(0), logits.stride(1), TopK);
sorted_emit_fallback_rows();
const cudaError_t fb_err = cudaGetLastError();
STD_TORCH_CHECK(fb_err == cudaSuccess, "sorted emission failed: ",
cudaGetErrorString(fb_err));
return;
}
cudaError_t status =
@@ -524,6 +573,10 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
static_cast<uint32_t>(stride), stream);
STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ",
cudaGetErrorString(status));
sorted_emit_fallback_rows();
const cudaError_t fb_err = cudaGetLastError();
STD_TORCH_CHECK(fb_err == cudaSuccess, "sorted emission failed: ",
cudaGetErrorString(fb_err));
return;
}
@@ -622,6 +675,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
params.det_ties = det_ties_enabled();
params.tie_recent = tie_recent_enabled();
params.exact_decode = exact_decode_enabled();
params.sorted_emit = sorted_emit_enabled();
report_tie_policy_once();
#if MERLIN_TOPK_DIAG_ANY
report_diag_build_once();
@@ -936,8 +990,12 @@ int64_t persistent_topk_exact_counters_slots() {
// field 3 MERLIN_TOPK_EXACT_MEDIUM (exact_medium_enabled)
// field 4 MERLIN_TOPK_EXACT_DECODE (exact_decode_enabled)
// field 5 MERLIN_TOPK_DET_LARGE (det_large_enabled)
// bits [48, 55] = field count (6), bits [56, 62] = layout version (1). A
// field 6 MERLIN_TOPK_SORTED_EMIT (sorted_emit_enabled)
// bits [48, 55] = field count (7), bits [56, 62] = layout version (2). A
// reader must check both and refuse any other layout rather than guess.
// Layout 2 = layout 1 plus field 6; fields 0-5 are unchanged, so
// a reader that knows both can decode either. Field 6 is what
// qsa_indexer.py reads to decide whether its Python sort may be skipped.
int64_t persistent_topk_toggle_state() {
#ifndef USE_ROCM
const auto field = [](int cached) -> int64_t {
@@ -959,8 +1017,9 @@ int64_t persistent_topk_toggle_state() {
state |= field(g_cached_exact_medium.load(std::memory_order_acquire)) << 6;
state |= field(g_cached_exact_decode.load(std::memory_order_acquire)) << 8;
state |= field(g_cached_det_large.load(std::memory_order_acquire)) << 10;
state |= int64_t{6} << 48;
state |= int64_t{1} << 56;
state |= field(g_cached_sorted_emit.load(std::memory_order_acquire)) << 12;
state |= int64_t{7} << 48;
state |= int64_t{2} << 56;
return state;
#else
return 0; // layout version 0: no persistent_topk on ROCm
vllm/envs.py
+26
View File
@@ -602,6 +602,26 @@ if TYPE_CHECKING:
# unless MERLIN_TOPK_LARGE_FAULT_INJECT names a CTA.
MERLIN_TOPK_LARGE_FAULT_KIND: str | None = None
#
# MERLIN_TOPK_SORTED_EMIT -- persistent_topk
# emits every row in the canonical order MERLIN_QSA_SORT imposes
# (ascending index, padding last as -1), and qsa_indexer.py then skips its
# Python sort. Same classification as the kernel switches above: resolved
# ONCE per process by a std::getenv in csrc/libtorch_stable/topk.cu
# (sorted_emit_enabled) -- the single source of truth -- and the resolved
# value reaches the GPU as a kernel launch argument AND removes three torch
# ops from the captured sequence, so a cudagraph bakes it in. CAPTURE: yes.
# AOT: not established for the kernel half; hashed conservatively (a fork
# env knob must be a compile factor: an os.environ knob that is not
# registered is not in the cache key). Raw passthrough: Python never parses
# it to decide anything -- qsa_indexer.py reads the C++ cache through
# persistent_topk_toggle_state() field 6. Not in ignored_factors.
MERLIN_TOPK_SORTED_EMIT: str | None = None
# Debug companion, a REAL Python read site (qsa_indexer._SORTED_EMIT_CHECK,
# read once at import): when the kernel sorted, assert on device that the
# row equals its canonical form instead of trusting it. Adds ops to the
# captured sequence, so it is hashed too. Default off.
MERLIN_TOPK_SORTED_EMIT_CHECK: bool = False
#
# NOT REGISTERED, deliberately: MERLIN_TOPK_EXACT_TRAP. It is a
# preprocessor define consumed by `#ifndef` in persistent_topk.cuh at
# BUILD time (-D), not a process environment variable -- nothing in the
@@ -2692,6 +2712,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
"MERLIN_TOPK_LARGE_FAULT_KIND": lambda: os.getenv(
"MERLIN_TOPK_LARGE_FAULT_KIND", None
),
# In-kernel canonical emission and its debug check; see the
# declaration block. Both hashed: both change the captured sequence.
"MERLIN_TOPK_SORTED_EMIT": lambda: os.getenv("MERLIN_TOPK_SORTED_EMIT", None),
"MERLIN_TOPK_SORTED_EMIT_CHECK": lambda: bool(
int(os.getenv("MERLIN_TOPK_SORTED_EMIT_CHECK", "0"))
),
}
vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py
+77 -5
View File
@@ -50,17 +50,84 @@ _DECODE_BLOCK_N = 64
# ---------------------------------------------------------------------------
_QSA_SORT = envs.MERLIN_QSA_SORT
# ---------------------------------------------------------------------------
# MERLIN_TOPK_SORTED_EMIT: persistent_topk emits
# every row in the canonical order below by itself (sorted_emit_row() in
# csrc/libtorch_stable/persistent_topk.cuh), so the Python sort is skipped.
#
# The decision is NOT a Python parse of the variable. The C++ getter
# (topk.cu sorted_emit_enabled) is the single source of truth and publishes
# what it cached into persistent_topk_toggle_state() field 6; that is read
# once, host-side, right after this process's first persistent_topk launch
# (the static is filled by then) and latched. So the sort is dropped only when
# the kernel that actually ran says it sorted; an older build (layout 1, no
# field 6) or an absent op keeps the Python sort. cooperative_topk (SM90+) has
# no sorted emission and always keeps the Python sort.
#
# MERLIN_TOPK_SORTED_EMIT_CHECK=1 (debug, default off): instead of trusting
# the kernel, assert on device that the row already equals its canonical form.
# torch._assert_async is capture-safe (no host sync); a violation is a device
# assert, i.e. a dead context -- which is the point of a debug check.
# ---------------------------------------------------------------------------
_SORTED_EMIT_CHECK = envs.MERLIN_TOPK_SORTED_EMIT_CHECK
_kernel_sorted_emit_state: bool | None = None # None: not read yet
def _qsa_canonicalise(block_indices: torch.Tensor) -> None:
"""Sort selected block indices ascending, padding last. In place, same shape."""
if not _QSA_SORT:
return
def _canonical(block_indices: torch.Tensor) -> torch.Tensor:
"""Ascending, padding last as -1 -- THE definition of the canonical order,
shared by the Python sort and the debug check (and the reference the
kernel's sorted emission is tested against)."""
big = torch.iinfo(block_indices.dtype).max
tmp = torch.where(
block_indices < 0, torch.full_like(block_indices, big), block_indices
)
tmp = torch.sort(tmp, dim=-1).values
block_indices.copy_(torch.where(tmp == big, torch.full_like(tmp, -1), tmp))
return torch.where(tmp == big, torch.full_like(tmp, -1), tmp)
def _qsa_canonicalise(block_indices: torch.Tensor) -> None:
"""Sort selected block indices ascending, padding last. In place, same shape."""
if not _QSA_SORT:
return
block_indices.copy_(_canonical(block_indices))
def kernel_sorted_emit() -> bool:
"""True iff this process's persistent_topk emits canonical order.
Host-only (a tensor-free query op; no CUDA call), so it is safe inside a
capture. Latches once the C++ static is known; an uninitialised static
(no launch yet) answers False without latching.
"""
global _kernel_sorted_emit_state
if _kernel_sorted_emit_state is not None:
return _kernel_sorted_emit_state
op = getattr(torch.ops._C, "persistent_topk_toggle_state", None)
if op is None:
_kernel_sorted_emit_state = False
return False
from vllm.v1.worker.e2_toggles import decode_topk_state
decoded = decode_topk_state(int(op()))
if isinstance(decoded, str) or "sorted_emit" not in decoded:
# Unknown layout, or layout 1 (no field 6): this build cannot sort.
_kernel_sorted_emit_state = False
return False
value = decoded["sorted_emit"]
if value == "uninit":
return False # not latched: ask again after a launch
_kernel_sorted_emit_state = value == "1"
return _kernel_sorted_emit_state
def _check_sorted_emission(block_indices: torch.Tensor) -> None:
"""Debug only: device-side assert that the kernel's row is canonical."""
if not _SORTED_EMIT_CHECK:
return
torch._assert_async(
torch.eq(block_indices, _canonical(block_indices)).all(),
"MERLIN_TOPK_SORTED_EMIT: persistent_topk row is not in canonical order",
)
@triton.jit
@@ -608,6 +675,11 @@ def _topk(
block_topk,
logits.shape[1],
)
# MERLIN_TOPK_SORTED_EMIT: the kernel already emitted the canonical order,
# so the Python sort would be a no-op -- skip it (optionally check it).
if not use_cooperative_topk and kernel_sorted_emit():
_check_sorted_emission(block_indices)
return
_qsa_canonicalise(block_indices)
vllm/v1/worker/e2_toggles.py
+56 -12
View File
@@ -12,13 +12,15 @@ producer at all. Every GPU worker therefore logs exactly one line::
e2-toggles[PP<p>_TP<t>]: pid=<pid> MERLIN_FULL_K=<raw> MERLIN_QSA_SORT=<raw>
MERLIN_TIE_DET=<raw> MERLIN_TIE_RECENT=<raw> MERLIN_TOPK_EXACT_MEDIUM=<raw>
MERLIN_TOPK_EXACT_DECODE=<raw> MERLIN_TOPK_DET_LARGE=<raw>
MERLIN_TOPK_SORTED_EMIT=<raw>
effective: full_k=<0|1> qsa_sort=<0|1|invalid> tie_det=<0|1> tie_recent=<0|1>
tie_recent_effective=<0|1> exact_medium=<0|1> exact_decode=<0|1>
det_large=<0|1> source=<libc-getenv|os-environ>
det_large=<0|1> sorted_emit=<0|1> source=<libc-getenv|os-environ>
cpp-cache: <full_k=... tie_det=... tie_recent=... tie_recent_effective=...
exact_medium=... exact_decode=... det_large=... source=cpp-cache
| n-a reason=<why>>
python: qsa_sort=<0|1|n-a> source=qsa_indexer._QSA_SORT
exact_medium=... exact_decode=... det_large=... sorted_emit=...
source=cpp-cache | n-a reason=<why>>
python: qsa_sort=<0|1|n-a> sorted_emit_skip=<0|1|n-a>
source=qsa_indexer._QSA_SORT
device: switch=<0|1|n-a> switch_large=<0|1|n-a> source=counter-slots-8-9
(one line in the log). ``<raw>`` is ``unset``, a plain token, or a double-quoted
@@ -42,7 +44,11 @@ WHAT EACH SECTION IS
was first cut against (earlier v2.2 builds, not published) have none.
* ``python:`` -- ``qsa_indexer._QSA_SORT``, the module constant the QSA sort
actually branches on (looked up in ``sys.modules``, never imported here;
``n-a`` if the module was never loaded in this process).
``n-a`` if the module was never loaded in this process), and
``sorted_emit_skip`` -- ``qsa_indexer._kernel_sorted_emit_state``, i.e.
whether the Python sort is being SKIPPED because the kernel reported
MERLIN_TOPK_SORTED_EMIT through ``persistent_topk_toggle_state`` field 6
(``n-a`` until the first QSA top-k has run and latched it).
* ``device:`` -- counter slots [8]/[9], which the counter op fills from the
C++ ``exact_medium_enabled()`` / ``det_large_enabled()`` statics
(topk.cu persistent_topk_exact_counters).
@@ -93,6 +99,7 @@ logger = init_logger(__name__)
# MERLIN_TOPK_EXACT_MEDIUM csrc/libtorch_stable/topk.cu:68-74 (unset -> ON)
# MERLIN_TOPK_EXACT_DECODE csrc/libtorch_stable/topk.cu:103-109
# MERLIN_TOPK_DET_LARGE csrc/libtorch_stable/topk.cu:86-92 (unset -> ON)
# MERLIN_TOPK_SORTED_EMIT csrc/libtorch_stable/topk.cu sorted_emit_enabled
CXX_SWITCHES: tuple[tuple[str, str, bool], ...] = (
("MERLIN_FULL_K", "full_k", False),
("MERLIN_TIE_DET", "tie_det", False),
@@ -100,6 +107,7 @@ CXX_SWITCHES: tuple[tuple[str, str, bool], ...] = (
("MERLIN_TOPK_EXACT_MEDIUM", "exact_medium", True),
("MERLIN_TOPK_EXACT_DECODE", "exact_decode", False),
("MERLIN_TOPK_DET_LARGE", "det_large", True),
("MERLIN_TOPK_SORTED_EMIT", "sorted_emit", False),
)
QSA_SORT_NAME = "MERLIN_QSA_SORT"
RAW_ORDER = (
@@ -110,6 +118,7 @@ RAW_ORDER = (
"MERLIN_TOPK_EXACT_MEDIUM",
"MERLIN_TOPK_EXACT_DECODE",
"MERLIN_TOPK_DET_LARGE",
"MERLIN_TOPK_SORTED_EMIT",
)
EFFECTIVE_ORDER = (
"full_k",
@@ -120,6 +129,7 @@ EFFECTIVE_ORDER = (
"exact_medium",
"exact_decode",
"det_large",
"sorted_emit",
)
CPP_ORDER = tuple(k for k in EFFECTIVE_ORDER if k != "qsa_sort")
@@ -265,21 +275,29 @@ _TOPK_STATE_FIELDS = (
"exact_medium",
"exact_decode",
"det_large",
"sorted_emit",
)
# Known layouts: version -> field count. Layout 2 is layout 1
# plus field 6 (MERLIN_TOPK_SORTED_EMIT); fields 0-5 are unchanged.
_TOPK_STATE_LAYOUTS = {1: 6, 2: 7}
def decode_topk_state(state: int) -> dict[str, str] | str:
"""Unpack ``persistent_topk_toggle_state()``: 2 bits per field (0 off,
1 on, 2 not initialised), field count at bits 48-55 (6), layout version at
bits 56-62 (1). Any other layout is refused with a reason string."""
1 on, 2 not initialised), field count at bits 48-55, layout version at
bits 56-62 -- (1, 6) or (2, 7). Any other layout is refused with a reason
string. A layout-1 build has no sorted emission, reported as ``n-a``."""
version, count = (state >> 56) & 0x7F, (state >> 48) & 0xFF
if version != 1 or count != len(_TOPK_STATE_FIELDS):
if _TOPK_STATE_LAYOUTS.get(version) != count:
return f"unknown-layout-v{version}-{count}-fields"
names = {0: "0", 1: "1", 2: "uninit"}
return {
out = {
key: names.get((state >> (2 * i)) & 3, "invalid")
for i, key in enumerate(_TOPK_STATE_FIELDS)
for i, key in enumerate(_TOPK_STATE_FIELDS[:count])
}
for key in _TOPK_STATE_FIELDS[count:]:
out[key] = "n-a"
return out
def decode_full_k_state(state: int) -> str:
@@ -324,6 +342,17 @@ def qsa_sort_live() -> str:
return "1" if bool(module._QSA_SORT) else "0"
def sorted_emit_skip_live() -> str:
"""``qsa_indexer._kernel_sorted_emit_state``: 1 = the Python sort is being
skipped because the kernel reported sorted emission, 0 = not skipped,
``n-a`` = not latched yet (or the module was never loaded)."""
module = sys.modules.get(_QSA_MODULE)
state = getattr(module, "_kernel_sorted_emit_state", None) if module else None
if state is None:
return "n-a"
return "1" if state else "0"
def device_switches() -> dict[str, str]:
"""Counter slots [8]/[9] (switch / switch_large). One synchronizing read
(topk_counters.read refuses inside a capture), so it is only called from
@@ -350,6 +379,7 @@ def toggles_line(
cpp: dict[str, str] | str,
qsa_live: str,
device: dict[str, str],
sorted_skip: str = "n-a",
) -> str:
parts = [f"e2-toggles[{tag}]: pid={pid}"]
parts += [f"{name}={raw[name]}" for name in RAW_ORDER]
@@ -362,7 +392,12 @@ def toggles_line(
else:
parts += [f"{key}={cpp[key]}" for key in CPP_ORDER]
parts.append("source=cpp-cache")
parts += ["python:", f"qsa_sort={qsa_live}", "source=qsa_indexer._QSA_SORT"]
parts += [
"python:",
f"qsa_sort={qsa_live}",
f"sorted_emit_skip={sorted_skip}",
"source=qsa_indexer._QSA_SORT",
]
parts.append("device:")
parts += [f"{key}={device[key]}" for key in ("switch", "switch_large")]
parts.append("source=counter-slots-8-9")
@@ -402,6 +437,7 @@ def maybe_emit(
cpp_fn: Callable[[], dict[str, str] | str] = cpp_cache,
device_fn: Callable[[], dict[str, str]] = device_switches,
qsa_fn: Callable[[], str] = qsa_sort_live,
sorted_skip_fn: Callable[[], str] = sorted_emit_skip_live,
) -> str:
"""Emit the line at most once per process. One bool check when already
emitted or not yet armed, so it can sit at the head of every execute step.
@@ -431,7 +467,15 @@ def maybe_emit(
cxx_getenv, env_source = getenv_fn()
raw, effective = environ_view(cxx_getenv, os.environ.get)
line = toggles_line(
tag, os.getpid(), raw, effective, env_source, cpp, qsa_fn(), device_fn()
tag,
os.getpid(),
raw,
effective,
env_source,
cpp,
qsa_fn(),
device_fn(),
sorted_skip_fn(),
)
_emitted = True
logger.info("%s", line)