From eeea036c1216c693d050216bf08de58f2592aaf2 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 23 Sep 2026 18:59:25 -0400 Subject: [PATCH] [TopK] In-kernel canonical emission (MERLIN_TOPK_SORTED_EMIT) 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 Signed-off-by: MOVZX --- csrc/libtorch_stable/persistent_topk.cuh | 227 +++++++++++++++++- csrc/libtorch_stable/topk.cu | 73 +++++- vllm/envs.py | 26 ++ .../qwen4_exp/nvidia/ops/qsa_indexer.py | 82 ++++++- vllm/v1/worker/e2_toggles.py | 68 +++++- 5 files changed, 444 insertions(+), 32 deletions(-) diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index cb6443c9a1..e339e0ab99 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -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(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 +struct SortedEmitSmem { + typename cub::BlockScan::TempStorage scan; + uint32_t bitmap[kSortedEmitBitmapWords]; + int32_t vals[TopK]; + uint32_t fallback; +}; + +template +__host__ __device__ constexpr size_t sorted_emit_smem_bytes() { + return (sizeof(SortedEmitSmem) + 15) & ~size_t(15); +} + +template +__device__ __noinline__ void sorted_emit_row(int32_t* __restrict__ row) { + static_assert(sorted_emit_smem_bytes() <= kSmemMedium, + "sorted emission must fit the smallest persistent_topk window"); + static_assert(alignof(SortedEmitSmem) <= 16, + "SortedEmitSmem over-aligned for the dynamic window"); + using BlockScan = cub::BlockScan; + extern __shared__ __align__(16) char sorted_emit_smem[]; + auto* s = reinterpret_cast*>(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(TopK); i += kThreadsPerBlock) { + const int32_t v = row[i]; + s->vals[i] = v; + if (v >= 0) { + const uint32_t u = static_cast(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(__popc(word)), pos, + total); + while (word != 0u) { + const uint32_t b = static_cast(__ffs(word) - 1); + row[pos++] = static_cast((tx << 5) | b); + word &= word - 1u; + } + for (uint32_t i = total + tx; i < static_cast(TopK); + i += kThreadsPerBlock) { + row[i] = -1; + } + } else { + constexpr int32_t kBig = 0x7FFFFFFF; + for (uint32_t i = tx; i < static_cast(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(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 -__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(logits + logits_offset, output_indices, static_cast(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(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + if (cta_in_group == 0) { + sorted_emit_row(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( 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(row_output); } } else { - histogram_256_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( + 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(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 +__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(output + static_cast(blockIdx.x) * TopK); +} + } // namespace persistent // ============================================================================ diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 26d5dfbcb3..ce0744bf11 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -61,6 +61,7 @@ std::atomic g_cached_tie_det{-1}; std::atomic g_cached_exact_medium{-1}; std::atomic g_cached_det_large{-1}; std::atomic g_cached_exact_decode{-1}; +std::atomic g_cached_sorted_emit{-1}; bool publish_cached(std::atomic& 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 + <<(num_rows), P::kThreadsPerBlock, + P::sorted_emit_smem_bytes(), stream>>>( + output.mutable_data_ptr(), + static_cast(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 && @@ -409,6 +452,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, logits.const_data_ptr(), lengths.const_data_ptr(), output.mutable_data_ptr(), stride, static_cast(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( @@ -417,6 +461,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, static_cast(TopK), static_cast(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(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 diff --git a/vllm/envs.py b/vllm/envs.py index ca6e460e09..912e71d2b4 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -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")) + ), } diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index cfee689897..a0dbe9e044 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -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) diff --git a/vllm/v1/worker/e2_toggles.py b/vllm/v1/worker/e2_toggles.py index 0ae7ad2afc..056a905413 100644 --- a/vllm/v1/worker/e2_toggles.py +++ b/vllm/v1/worker/e2_toggles.py @@ -12,13 +12,15 @@ producer at all. Every GPU worker therefore logs exactly one line:: e2-toggles[PP

_TP]: pid= MERLIN_FULL_K= MERLIN_QSA_SORT= MERLIN_TIE_DET= MERLIN_TIE_RECENT= MERLIN_TOPK_EXACT_MEDIUM= MERLIN_TOPK_EXACT_DECODE= MERLIN_TOPK_DET_LARGE= + MERLIN_TOPK_SORTED_EMIT= 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= + det_large=<0|1> sorted_emit=<0|1> source= cpp-cache: > - 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=> + 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). ```` 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)